File size: 6,438 Bytes
# FFOLayer — A Fully First-Order Layer for Differentiable Optimization

`FFOLayer` is a PyTorch-friendly library for **differentiable optimization layers** that computes hypergradients using **only first-order information**. It is designed as a practical, drop-in alternative to implicit-differentiation-based layers when memory or backward-time is the bottleneck.

---
## Installation
FFOLayer is available on pip:

```bash
pip install ffolayer
```

You may also need to install cvxtorch:

```bash
git clone https://github.com/cvxpy/cvxtorch.git
cd cvxtorch
pip install -e .
```

---

## Usage

FFOLayer follows the same workflow as differentiable layers like [CvxpyLayer](https://github.com/cvxpy/cvxpylayers/):

> define a [CVXPY](https://github.com/cvxpy/cvxpy) problem → wrap it as a layer → call it in PyTorch → backprop

### API

#### `FFOLayer(problem, parameters, variables, eps=..., **kwargs)`

**Arguments**
- `problem`: a `cvxpy.Problem` (must satisfy DPP when using CVXPY parameters).
- `parameters`: CVXPY `Parameter` (the inputs to the layer).
- `variables`: CVXPY `Variable` (the outputs returned by the layer).
- `alpha`: perturbation scale (\delta in Eq.4) used in the finite-difference hypergradient approximation. Bigger alpha leads to a smaller delta: less bias, more numerical noise.
- `eps`: solver tolerence for forward pass
- `backward_eps`: solver tolerence for backward pass
- `max_workers`: maximum number of worker threads/processes used to parallelize solver calls 

#### Calling the layer
```python
(outputs,) = layer(*torch_parameters, solver_args={...})
```

**Arguments**
- `*torch_parameters`: PyTorch tensors matching `parameters` in shape (optionally batched).
- `solver_args` (optional): forwarded to `problem.solve(...)` inside CVXPY.

**Return value**
- A tuple of PyTorch tensors corresponding to `variables`.


### Solver-agnostic differentiation!!!

FFOLayer is **solver-agnostic**: it treats the solver as a **black box** and computes hypergradients by re-solving **perturbed** problems, instead of differentiating through solver internals.  
This means you can use **any CVXPY solver** (e.g., GUROBI, MOSEK, ECOS, SCS) without requiring custom backward implementations.

```python
(solution,) = layer(*torch_parameters, solver_args={"solver": cp.GUROBI, "eps": 1e-5})
```



#### Example

```python
import torch

from ffolayer import FFOLayer

batch, n, m = 8, 2, 3
x = cp.Variable(n)
A = cp.Parameter((m, n))
b = cp.Parameter(m)
constraints = [x >= 0]
objective = cp.Minimize(0.5 * cp.pnorm(A @ x - b, p=1))
problem = cp.Problem(objective, constraints)
assert problem.is_dpp()

# layer = CvxpyLayer(problem, parameters=[A, b], variables=[x])
layer_ffo = FFOLayer(problem, parameters=[A, b], variables=[x], eps=1e-8)
A_tch = torch.randn(batch, m, n, requires_grad=True)
b_tch = torch.randn(batch, m, requires_grad=True)

# solve the problem
(solution,) = layer_ffo(A_tch, b_tch)

# compute the gradient of the sum of the solution with respect to A, b
solution.sum().backward()
```
---
## Code Structure

- `src/`: core layer implementations (FFOCP / FFOQP variants)
- `synthetic_task/`: synthetic decision-focused learning (QP) benchmark
- `sudoku/`: Sudoku as an optimization layer benchmark
- `baselines/`: reference baselines used in experiments
- `tests/`: basic checks / utilities
- `plot_results_*.ipynb`: notebooks for plotting paper figures

### Variants in this repo

We provide two main variants (same core idea, different specialization):

- **FFOCP:** applies to general convex programs.
- **FFOQP:** specializes to QP layers, exploiting quadratic structure for efficiency.


---
## Reproducing the paper experiments

### 1) Synthetic Qradratic Program (QP)

**Key files**
  - `synthetic_task/main_synthetic.py`: entrypoint for all methods
  - `models.py`: all models' definitions and settings

To run the code, if your cluster supports SLURM, please use:
```bash
sh scripts/loop_synthetic_per_seed.sh
```
If not, please use:
```bash
python synthetic_task/main_synthetic.py --method ffocp_eq --ydim 800 --epoch 1 --batch_size 8
python synthetic_task/main_synthetic.py --method ffoqp_eq --ydim 800 --epoch 1 --batch_size 8
python synthetic_task/main_synthetic.py --method lpgd --ydim 800 --epoch 1 --batch_size 8
python synthetic_task/main_synthetic.py --method cvxpylayer --ydim 800 --epoch 1 --batch_size 8
python synthetic_task/main_synthetic.py --method qpth --ydim 800 --epoch 1 --batch_size 8
python synthetic_task/main_synthetic.py --method bpqp --ydim 800 --epoch 1 --batch_size 8
python synthetic_task/main_synthetic.py --method dqp --ydim 800 --epoch 1 --batch_size 8
```

To plot the results, please run `plot_results_synthetic.ipynb`.


---

### 2) Sudoku task
**Key files**
  - `sudoku/main_sudoku.py`: entrypoint for all methods
  - `models_sudoku.py`: all models' definitions and settings

To run the code, if your cluster supports SLURM, please use:
```bash
sh scripts/loop_sudoku_per_seed.sh
```
If not, please use:
```bash
python sudoku/main_sudoku.py --method ffocp_eq --n 3 --epoch 1 --batch_size 8
python sudoku/main_sudoku.py --method ffoqp_eq --n 3 --epoch 1 --batch_size 8
python sudoku/main_sudoku.py --method lpgd --n 3 --epoch 1 --batch_size 8
python sudoku/main_sudoku.py --method cvxpylayer --n 3 --epoch 1 --batch_size 8
python sudoku/main_sudoku.py --method qpth --n 3 --epoch 1 --batch_size 8
python sudoku/main_sudoku.py --method bpqp --n 3 --epoch 1 --batch_size 8
python sudoku/main_sudoku.py --method dqp --n 3 --epoch 1 --batch_size 8
```

To plot the results, please use
```bash
python sudoku/plot_results.py
```

---

### 3) `Synthetic Second-order Cone Progrem (SOCP))`
**Key files**
- `synthetic_task/main_synthetic_general.py`: entrypoint for all methods
- `models.py`: all models' definitions and settings

To run the code, if your cluster supports SLURM, please use:
```bash
sh scripts/loop_synthetic_general_per_seed.sh
```
If not, please use:
```bash
python synthetic_task/main_synthetic_general.py --method ffocp_eq --ydim 800 --epoch 1 --batch_size 8
python synthetic_task/main_synthetic_general.py --method lpgd --ydim 800 --epoch 1 --batch_size 8
python synthetic_task/main_synthetic_general.py --method cvxpylayer --ydim 800 --epoch 1 --batch_size 8
python synthetic_task/main_synthetic_general.py --method bpqp --ydim 800 --epoch 1 --batch_size 8
```

To plot the results, please run `plot_results_synthetic.ipynb`.

---