ProCreations's picture
Publish six-claim native-scale FFOLayer reproduction
9860743 verified
|
Raw
History Blame Contribute Delete
6.44 kB
# 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`.
---