# Reverse-Sweep Backward for Block-Implicit Layers

- ID: 2000
- Canonical URL: https://synthcore.org/idea/2000/reverse-sweep-backward-for-block-implicit-layers
- API JSON: https://synthcore.org/api/idea/2000.json
- API Markdown: https://synthcore.org/api/idea/2000.md
- Verification status: mech_ok_no_baseline
- Source: [arXiv:2608.08559](https://arxiv.org/abs/2608.08559)
- Category: architecture
- Solves: memory, speedup, stability
- ML areas: graph-nn, training, memory
- Math tags: optimization, linear-algebra, dynamical-systems, numerical-analysis
- Ratings: usefulness 8/10; difficulty 5/10; novelty 6/10

## Idea description

Replace unrolled autodiff through an ordered block-implicit neural layer with a custom reverse sweep that solves one small transposed local system per forward block update. The backward computes the exact gradient of the executed finite-depth solver while avoiding a global Jacobian and retaining only compact block information.

## Mathematical statement

The paper starts from the implicit stationary problem \(\nabla G(\bm{x})=\bm{0}\), with \(G(\bm{x})=\tfrac{1}{2h^{2}}(\bm{x}-\bm{y})^{\top}M(\bm{x}-\bm{y})+E(\bm{x})\), where \(\bm{x}\) is the state, \(\bm{y}\) is a prediction, \(M\) is a mass or metric matrix, \(h\) is a step size, and \(E\) is an energy. The transferable construction is a finite ordered composition of local implicit maps. For block \(i\), define a residual \(r_i(x_i;z_i,\theta)=0\), where \(x_i\) is the block being solved, \(z_i\) contains already-available neighboring states, and \(\theta\) contains network parameters or inputs. Let \(A_i=\partial r_i/\partial x_i\), \(B_i=\partial r_i/\partial z_i\), and \(C_i=\partial r_i/\partial\theta\), evaluated at the executed forward solution. Differentiating the local solve gives \(A_i\,dx_i+B_i\,dz_i+C_i\,d\theta=0\). If \(\bar{x}_i=\partial L/\partial x_i\) is the reverse cotangent entering the block, solve \(A_i^{\top}\lambda_i=\bar{x}_i\), then propagate \(\bar{z}_i\mathrel{+}= -B_i^{\top}\lambda_i\) and \(\bar{\theta}\mathrel{+}= -C_i^{\top}\lambda_i\). Applying these updates in the exact reverse order of the forward block overwrites yields the transpose of the executed solver without a global adjoint matrix.

## Key formulas

- $$\nabla G(\bm{x})=\bm{0},\qquad G(\bm{x})=\tfrac{1}{2h^{2}}(\bm{x}-\bm{y})^{\!\top}M(\bm{x}-\bm{y})+E(\bm{x}).$$
- $$r_i(x_i;z_i,\theta)=0,\qquad A_i\,dx_i+B_i\,dz_i+C_i\,d\theta=0,$$
- $$A_i^{\top}\lambda_i=\bar{x}_i,\qquad \bar{z}_i\mathrel{+}= -B_i^{\top}\lambda_i,\qquad \bar{\theta}\mathrel{+}= -C_i^{\top}\lambda_i.$$
- $$\texttt{for }k=K,\ldots,1:\quad \texttt{for }i=m,\ldots,1:\quad A_{k,i}^{\top}\lambda_{k,i}=\bar{x}_{k,i};\ \bar{z}_{k,i}\mathrel{+}= -B_{k,i}^{\top}\lambda_{k,i};\ \bar{\theta}\mathrel{+}= -C_{k,i}^{\top}\lambda_{k,i}.$$

## Implementation notes

Integrate this as a custom autograd Function around a block-implicit graph layer. The layer should maintain node or channel states \(x\) and perform \(K\) colored Gauss-Seidel-like sweeps; in each sweep, block \(i\) solves a small residual equation such as \(r_i(x_i;z_i,\theta)=x_i-\phi_i(W_i x_i+U_i z_i+b_i)\), or a Newton/proximal variant. During forward, execute the actual solver, save only the solved block values, block ordering, coloring, and any factorization or low-dimensional Jacobian representation needed for each local solve. Do not save every elementary operation. During backward, initialize cotangents from the loss and iterate sweeps and blocks in exact reverse order. For each block, form or reuse \(A_{k,i}=\partial r_{k,i}/\partial x_{k,i}\), solve \(A_{k,i}^{\top}\lambda_{k,i}=\bar{x}_{k,i}\), and use vector-Jacobian products to add \(-B_{k,i}^{\top}\lambda_{k,i}\) to neighboring-state cotangents and \(-C_{k,i}^{\top}\lambda_{k,i}\) to parameter and input gradients. Detach the Newton iterations inside each local solve: differentiate the solved residual, not the iterations used to find it. For a first implementation use block width 32 or 64 and dense batched solves; estimate Jacobians with autodiff JVP/VJP rather than materializing them for larger blocks. Validate gradients against full unrolled autodiff at depths \(K=1,2,4,8\) in float64, measuring relative error and peak memory. The first cheap experiment is a small CIFAR-10 graph or channel-block equilibrium classifier, comparing unrolled reverse-mode, equation-level fixed-point differentiation, and this reverse sweep at equal forward sweeps. Success means relative gradient error below \(10^{-5}\) versus unrolling, substantially lower peak memory, and faster loss descent or larger feasible \(K\) at fixed GPU memory.

## Verification

- Status: mech_ok_no_baseline
- Mechanism evidence: yes
- Mechanism confirmed: yes
- Practical verdict: no_effect
- Verdict: Built a scalar block-implicit layer with detached local Newton solves and an exact reverse-sweep custom autograd backward. The predicted local-Jacobian singularity at w=1 and geometric neighbor sensitivity gamma^K were confirmed numerically, while gradients matched unrolled autodiff at K=1,2,4,8 to 1.36e-16 relative error. The toy timing was 5.10 seconds for reverse sweep versus 19.42 seconds for unrolled autodiff over 100 K=4 runs, although this is not a reliable large-model memory benchmark.

### Mechanism check

- Verdict: Built a scalar block-implicit layer with detached local Newton solves and an exact reverse-sweep custom autograd backward. The predicted local-Jacobian singularity at w=1 and geometric neighbor sensitivity gamma^K were confirmed numerically, while gradients matched unrolled autodiff at K=1,2,4,8 to 1.36e-16 relative error. The toy timing was 5.10 seconds for reverse sweep versus 19.42 seconds for unrolled autodiff over 100 K=4 runs, although this is not a reliable large-model memory benchmark.
- Confidence: 8/10
- Limitations: Only a scalar CPU toy implementation was tested. No CIFAR or graph task, batched dense block solves, CUDA fallback path, peak-memory measurement, FLOP-normalized comparison, or training/loss-descent experiment was performed.

### Practical benchmark

- Paired seeds: 8
- Baseline mean: 0
- Idea mean: 0
- p-value: 1
- Paired wins: 2/8
- Benchmark verdict: no measurable effect

## Artifacts

- [bench_report.json](https://synthcore.org/code/757/bench_report.json)
- [bench_reverse_sweep.py](https://synthcore.org/code/757/bench_reverse_sweep.py)
- [prediction_results.json](https://synthcore.org/code/757/prediction_results.json)
- [report.md](https://synthcore.org/code/757/report.md)
- [report_bench_2026-09-01T024533.md](https://synthcore.org/code/757/report_bench_2026-09-01T024533.md)
- [results.json](https://synthcore.org/code/757/results.json)
- [reverse_sweep_mvp.py](https://synthcore.org/code/757/reverse_sweep_mvp.py)
- [validate_predictions.py](https://synthcore.org/code/757/validate_predictions.py)
- [Download all files as ZIP](https://synthcore.org/download/757)

## Disclaimer

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