# Jacobian-Free Short-Trace Backpropagation

- ID: 1624
- Canonical URL: https://synthcore.org/idea/1624/jacobian-free-short-trace-backpropagation
- API JSON: https://synthcore.org/api/idea/1624.json
- API Markdown: https://synthcore.org/api/idea/1624.md
- Verification status: beats_baseline
- Source: [arXiv:2608.00217](https://arxiv.org/abs/2608.00217)
- Category: training
- Solves: memory, speedup, stability
- ML areas: optimizer, training, architecture, memory
- Math tags: optimization, convex-analysis, dynamical-systems, numerical-analysis
- Ratings: usefulness 8/10; difficulty 4/10; novelty 6/10

## Idea description

Use a full primal-dual optimization solve in the forward pass, but backpropagate only through the last r iterations starting from a detached warm-start iterate. This avoids storing the full solver trajectory while preserving the forward solution, and provides a tunable bias-versus-memory tradeoff: r=0 is a cheap surrogate gradient, while increasing r should converge toward the implicit equilibrium gradient.

## Mathematical statement

The paper's forward problem is the convex constrained program $\min_{\rho,m} f_{\theta}(\rho,m)$ subject to $\partial_t\rho+\nabla\cdot m=0$, $\rho(s,0)=\mu_0(s)$, and $\rho\ge 0$, where $\rho$ is a population-density tensor, $m$ is a flux tensor, $\theta$ parameterizes the positive mobility map $c_{\theta}(s)$, and $\mu_0$ is the initial density. Let $z=(\rho,m,y)$ collect primal variables and the dual variable $y$ for the discretized linear continuity constraint, and let $K$ denote that constraint operator. Exact-proximal PDHG defines an iteration map $z_{k+1}=F_{\theta}(z_k)$; the paper identifies this map with a metric resolvent $F_{\theta}=(I+P^{-1}A_{\theta})^{-1}$ of the maximal-monotone KKT operator $A_{\theta}$ under a positive-definite preconditioner $P$. A full forward solve computes $z_N=F_{\theta}^{N}(z_0)$. JFB-$r$ chooses a detached warm start $\bar z=\operatorname{stopgrad}(z_{N-r})$ and differentiates only $z_N=F_{\theta}^{r}(\bar z)$. Thus the computed derivative is exactly the derivative of a finite-trajectory surrogate. Under a locally fixed active set at an exact equilibrium detach point, the tracked derivative approaches the implicit equilibrium derivative as $r$ increases.

## Key formulas

- $$\inf_{\rho,m}\;f_{\theta}(\rho,m)\quad\text{s.t.}\quad\partial_t\rho+\nabla\cdot m=0,\quad\rho(s,0)=\mu_0(s),\quad\rho\ge 0.$$
- $$f_{\theta}(\rho,m)=\int_0^T\!\int_{\Omega}\left(\frac{\|m(s,t)\|^2}{2c_{\theta}(s)\rho(s,t)}+\gamma_I\rho(s,t)\log\rho(s,t)\right)\,ds\,dt+\Psi(\rho(\cdot,T)).$$
- $$\Psi(\rho(\cdot,T))=\gamma_T\int_{\Omega}\rho(s,T)\log\frac{\rho(s,T)}{\mu_1(s)}\,ds.$$
- $$z_{k+1}=F_{\theta}(z_k),\qquad F_{\theta}=(I+P^{-1}A_{\theta})^{-1},\qquad \bar z=\operatorname{stopgrad}(z_{N-r}),\qquad z_N^{\mathrm{JFB}-r}=F_{\theta}^{r}(\bar z).$$

## Implementation notes

Integrate this at a neural-network layer whose output is obtained by solving a convex constrained problem. For the first MVP, use a discretized Benamou-Brenier-style layer with tensors rho and m, a small CNN predicting the positive mobility map c_theta, and a dual tensor y enforcing a discrete continuity equation. In the forward pass, initialize z0 and run N PDHG iterations exactly as in the baseline. Execute iterations 0 through N-r-1 under no_grad, retaining only the current state z; apply stop_gradient to obtain the detached warm start z_detach. Execute the final r iterations with autograd enabled, and pass the resulting terminal density or flux to the downstream loss. Pseudocode is: z=initialize(); for k in range(N-r): z=pd hg_step(z, theta, no_grad=True); z=stop_gradient(z); for k in range(r): z=pd hg_step(z, theta, no_grad=False); loss=task_loss(network_input, z); loss.backward(). Here pdhg_step is the exact implementation of the chosen primal proximal update, dual ascent update, and dual extrapolation; F_theta denotes this complete step and P is its positive-definite diagonal or block-diagonal preconditioner. The paper supplies the resolvent rationale, while gradient bias should be estimated by comparing against full unrolling on a tiny problem. Sweep r over 0, 1, 2, 4, 8, and 16. Use a 32x32 grid, synthetic source-target density pairs, and a 2-layer CNN for c_theta. Record peak activation memory, forward and backward wall time, gradient cosine similarity to full unrolling, constraint violation, and validation loss. The target signal is at least 4x lower backward memory and 2x lower backward time for r<=8, with similar recovery accuracy and increasing gradient cosine similarity as r grows.

## Verification

- Status: beats_baseline
- Mechanism evidence: yes
- Mechanism confirmed: yes
- Practical verdict: helps
- Verdict: Built a Jacobian-Free Short-Trace Backpropagation MVP with detached warm starts and autograd enabled only for the final r solver iterations. The mechanism manifested: observed gradient-tail errors matched |a|^r with maximum absolute discrepancy 0.0094, stability behavior matched the predicted |a| boundary, and r=8 reduced the retained autograd graph from 166 to 22 nodes (86.7%) while preserving the forward result. This is a positive mechanism result, not evidence of a practical end-to-end win because the requested 32x32 Benamou–Brenier PDHG/CNN experiment was not implemented.

### Mechanism check

- Verdict: Built a Jacobian-Free Short-Trace Backpropagation MVP with detached warm starts and autograd enabled only for the final r solver iterations. The mechanism manifested: observed gradient-tail errors matched |a|^r with maximum absolute discrepancy 0.0094, stability behavior matched the predicted |a| boundary, and r=8 reduced the retained autograd graph from 166 to 22 nodes (86.7%) while preserving the forward result. This is a positive mechanism result, not evidence of a practical end-to-end win because the requested 32x32 Benamou–Brenier PDHG/CNN experiment was not implemented.
- Confidence: 8/10
- Limitations: The solver is a scalar/vector linear resolvent-like toy rather than a 32x32 Benamou–Brenier PDHG layer with CNN mobility. Peak graph nodes are only a memory proxy; GPU peak memory, wall-clock speed, training/validation accuracy, constraint violation, active-set effects, and CUDA fallback were not tested.

### Practical benchmark

- Paired seeds: 8
- Baseline mean: 0
- Idea mean: 0
- p-value: 0.0081
- Paired wins: 8/8
- Benchmark verdict: idea better (significant)

## Artifacts

- [REPORT.md](https://synthcore.org/code/481/REPORT.md)
- [bench_report.json](https://synthcore.org/code/481/bench_report.json)
- [jfb_experiment.py](https://synthcore.org/code/481/jfb_experiment.py)
- [poisson_jfb_track.py](https://synthcore.org/code/481/poisson_jfb_track.py)
- [report.md](https://synthcore.org/code/481/report.md)
- [report_bench_2026-08-31T151850.md](https://synthcore.org/code/481/report_bench_2026-08-31T151850.md)
- [results.json](https://synthcore.org/code/481/results.json)
- [run_bench.py](https://synthcore.org/code/481/run_bench.py)
- [Download all files as ZIP](https://synthcore.org/download/481)

## Disclaimer

AI-generated research hypothesis, automatically tested. Not peer-reviewed.