This project demonstrates how to wrap JuMP.jl / DiffOpt.jl optimization models with in a Tesseract, enabling automatic differentiation through Julia-based optimization problems for ML workflows with optimization-in-the-loop in JAX.
The integration of ML pipelines and optimization routines is becoming increasingly popular (see e.g. this review article). However, combining Python-based ML frameworks (PyTorch, JAX) with optimization ecosystems in other languages (e.g., JuMP in Julia) while maintaining end-to-end differentiability remains challenging. This project demonstrates one way of bridging this gap using tesseract-core.
- Julia/JuMP Optimization as a Tesseract: Definition & solution of JuMP optimization models within a tesseract (container) with AD-ready interface
- Automatic Differentiation: Forward and reverse-mode AD through optimization problems via DiffOpt.jl
- JAX Integration: Composition in JAX pipelines using Tesseract-JAX
- Thermal Generation Example: Demonstration on simple example with neural network demand prediction + quadratic program for optimal dispatch decisions
-
Clone the repository
git clone https://github.com/llueg/diffopteract.git cd diffopteract -
Install Python dependencies These are packages required to run the end-to-end example in ´main.py´. It is recommended to install these into a fresh virtual environment.
pip install -r requirements.txt
-
Build the Tesseract
Note that this will take a while - Julia, etc. needs to be setup inside the container from scratch.
./buildall.sh
-
Run the thermal generation example
Note that training loop with the current implementation is relatively slow. It took around 2000s (!) on my laptop to generate the results shown in ´plots´.
python main.py
An example of a typical workflow with optimization in the loop can be visualized as follows:
The tesseract implemented in this project covers the ´Optimization model´ component: a user-defined problem is formulated and solved in JuMP and tesseract-core is used to wrap this step into a interface that maps from parameters to variables, while maintinaing differentiability. In the provided example, we excercise an entire workflow as depicted above on a simple example problem. In that case, the rest of the above workflow is implemented directly in JAX - an end-to-end differentiable workflow is achieved by using the Tesseract-JAX package.
The API of the diffopteract Tesseract is kept very simple and generic:
class InputSchema(BaseModel):
parameters: Differentiable[Array[(PARAM_DIM,), Float64]]
class OutputSchema(BaseModel):
variables: Differentiable[Array[(VAR_DIM,), Float64]]Julia code to define the interface to DiffOpt/JuMP (solving model, evaluating sensitivities) is located in tesseracts/diffopteract/src/diffopt_interface.jl. To link specific optimization problems, three functions need to be implemented in Julia: generate_model(...), get_parameter_refs(...) and get_variable_refs(...). For the example problem in this case, these are implemented in tesseracts/diffopteract/src/thermalgen.jl. The Julia functions are called using JuliaCall in tesseracts/diffopteract/tesseract_api.py. Only minor adjustments (e.g. importing correct files, setting PARAM_DIM and VAR_DIM, setting auxiliary parameters) need to be made to tesseract_api.py when using the Tesseract for different optimization problems.
We exercise the workflow on a simple example, adapted from the DiffOpt.jl documentation.
The thermal generation dispatch problem determines the optimal allocation of electricity generation across multiple thermal generators to meet demand while minimizing cost. The mathimatical formulation is given by:
$$ \begin{align} \min_{g, \phi} \quad & \sum_{i=1}^{N} c_i g_i + \sum_{i=1}^{N} q_i g_i^2 + c_\phi \phi \ \text{s.t.} \quad & 0 \leq g_i \leq G_i, \quad i = 1, \ldots, N \ & \sum_{i=1}^{N} g_i + \phi = d, \end{align} $$ where:
-
$g_i$ : Generation output from generator$i$ -
$\phi$ : Slack variable for unmet demand -
$d$ : Electricity demand (parameter) -
$c_i$ : Linear cost coefficient for generator$i$ -
$q_i$ : Quadratic cost coefficient for generator$i$ -
$c_\phi$ : Penalty cost for unmet demand (set high to discourage) -
$G_i$ : Maximum capacity of generator$i$
In this example, we only wish to train an ML model to predict the parameter corresponding to the demand (tesseract_api.py.
# Generator capacities
g_sup = [40.0, 80.0, 150.0]
# Linear cost coefficients
c_g = [1.0, 2.0, 3.0]
# Quadratic cost coefficients
q_g = [1.0, 0.8, 0.6]
# Penalty for unmet demand
c_phi = 1000.0We generated a data set of time-varying demand (t,d) and resulting optimal dispatch and slack values (g, φ). The data is found in data/thermalgen_data.csv. We wish to train a NN that predicts demand based on time, however we do not use the true demand in the loss function, but instead compute predicted dispatch and slack values based on predicted demand, and compare that to the known optimal dispatch levels at a given time. This way, the optimization problem must be integrated into the learning pipeline.
The JAX code to implement and train this pipeline is given in main.py. There, the Tesseract is transformed into a JAX function using Tesseract-JAX, with the resulting loss function combining ML model and optimization Tesseract, e.g.:
def apply_thermalgen(params: Array) -> Array:
"""Apply the gen_cost Tesseract to compute dispatch from demand."""
output = apply_tesseract(
thermalgen_tesseract,
{"parameters": params},
)
return output["variables"]
def loss_fun(model: nnx.Module, data: Array, targets: Array):
pred_demand = model(data) # NN predicts demand
preds = apply_thermalgen(pred_demand) # Solve optimization
loss = optax.squared_error(preds, targets).mean()
return loss, predsThe ML model is implemented using flax. Note that in order to ensure feasibility of the optimization problem, predicted demand must be non-negative. Hence, a relu activation function is applied to the final output of the model.
The trained neural network learns to predict demand from time variables, and these predictions are used to compute optimal dispatch decisions. The model successfully learns the relationship between time and demand:
Consequently, the predicted decisions also align closely with those in the dataset:
Gradient-based training of NN models usually requires thousands of steps from the optimizer. With an optimization problem in the loop, multiple (here, one for each element in a batch) optimization problems need to be solved in each step, which is significantly more expensive than forward evaluation of typical ML models. Consequently, even a simple training pipeline, as shown here, requires significant training time. Performance improvements, particularly for the concurrent evaluation of optimization problems, seem like a worthwile extension to this code.
Furthermore, it is well-known that for certain optimization problems (e.g. LPs) the sensitivity of the solution w.r.t. some parameters is often equal to zero. This creates issues for the differentiability of the entire pipeline. Implementing some gradient smoothing techniques to improve this behavior is also a useful avenue for future work.
This project builds upon several excellent open-source tools:
- Tesseract Core: Container-based composable computing framework
- Tesseract-JAX: JAX integration for Tesseract
- DiffOpt.jl: Automatic differentiation for optimization problems in Julia
- JuMP.jl: Mathematical optimization modeling language for Julia
- JAX: High-performance numerical computing and autodiff
- Flax: Neural network library for JAX
Special thanks to the Tesseract hackathon organizers and the JuMP/DiffOpt communities for building these powerful tools.
Licensed under Apache License 2.0.
All submissions must use the Apache License 2.0 to be eligible for the Tesseract Hackathon. See LICENSE file for details.