△ Mechanism confirmed, baseline not beaten 2026

Memory-Light Differentiable Learned Optimizer

Implementation & benchmark of arXiv:2608.08559 — Differentiate the Solver, Not the Equation: Reverse-Sweep Adjoints for Block Implicit Simulation

Usefulness7/10

Difficulty6/10

Novelty5/10

Optimizer Memory Speedup Accuracy Optimizer Fine tuning Training Memory Optimization Linear algebra Dynamical systems Numerical analysis

Source paper: Differentiate the Solver, Not the Equation: Reverse-Sweep Adjoints for Block Implicit Simulation arXiv:2608.08559 ⓘ · analyzed Sep 1, 2026

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

Idea description

Construct a learned optimizer whose update is an ordered sequence of local implicit parameter-block solves, then differentiate the finite optimization trajectory with reverse local adjoints. This enables training optimizer hyperparameters or meta-gradients through many inner steps without storing all intermediate tensor operations or replacing the executed trajectory by an idealized fixed-point gradient.

Formulas

$$r_{k,i}(w_{k,i};s_{k,i},\eta)=0,$$

$$A_{k,i}\,dw_{k,i}+B_{k,i}\,ds_{k,i}+C_{k,i}\,d\eta=0,\qquad A_{k,i}=\frac{\partial r_{k,i}}{\partial w_{k,i}},\quad B_{k,i}=\frac{\partial r_{k,i}}{\partial s_{k,i}},\quad C_{k,i}=\frac{\partial r_{k,i}}{\partial\eta},$$

$$A_{k,i}^{\top}\lambda_{k,i}=\bar w_{k,i},\qquad \bar s_{k,i}\mathrel{+}= -B_{k,i}^{\top}\lambda_{k,i},\qquad \bar\eta\mathrel{+}= -C_{k,i}^{\top}\lambda_{k,i}.$$

$$\frac{\partial L_{\mathrm{outer}}}{\partial\eta}=\sum_{k=K}^{1}\sum_{i=m}^{1}-C_{k,i}^{\top}\lambda_{k,i},\qquad A_{k,i}^{\top}\lambda_{k,i}=\bar w_{k,i}.$$

Mathematical statement

Let the inner optimizer state be partitioned into blocks $w=(w_1,\ldots,w_m)$, with momentum or auxiliary state included in each block. At inner step $k$, block $i$ is defined by an implicit update residual $r_{k,i}(w_{k,i};s_{k,i},\eta)=0$, where $s_{k,i}$ contains the current gradient, neighboring parameter blocks, and auxiliary states, and $\eta$ are learned optimizer parameters such as step sizes, preconditioners, or damping values. The local differential identity is $A_{k,i}dw_{k,i}+B_{k,i}ds_{k,i}+C_{k,i}d\eta=0$, with $A_{k,i}=\partial r_{k,i}/\partial w_{k,i}$, $B_{k,i}=\partial r_{k,i}/\partial s_{k,i}$, and $C_{k,i}=\partial r_{k,i}/\partial\eta$. Reverse differentiation of the executed inner trajectory solves $A_{k,i}^{\top}\lambda_{k,i}=\bar w_{k,i}$ in reverse time and reverse block order, then propagates $-B_{k,i}^{\top}\lambda_{k,i}$ and $-C_{k,i}^{\top}\lambda_{k,i}$. The construction is a transpose of the finite solver map, not the adjoint of the converged optimization equation, so it remains correct when the inner optimizer uses only a small fixed number of sweeps.

Implementation notes

Use this inside a meta-learning or differentiable optimization benchmark. Partition model parameters into layer blocks, for example one block per Transformer layer or one block per convolutional stage, and define an implicit damped update for block $i$: $r_{k,i}(w_{k,i};s_{k,i},\eta)=w_{k,i}-w_{k-1,i}+\alpha_i g_{k,i}(w_{k,i})+D_i^{-1}(w_{k,i}-w_{k-1,i})$, where $g_{k,i}$ is the block gradient, $\alpha_i>0$ is a learned step size, and $D_i$ is a positive diagonal damping vector. Solve each block with one or two Newton or fixed-point iterations, while treating the resulting block value as the output of the local implicit solve. During forward, run $K$ inner steps and save block outputs, gradients, and diagonal or factored local Jacobians. During backward from the outer validation loss, traverse $k=K$ to 1 and blocks in reverse order. Form the local VJP for $A_{k,i}$, solve the transposed system; with diagonal damping this is elementwise division, while a small dense layer block can use batched LU. Propagate cotangents into earlier parameters and accumulate meta-gradients for $\eta=(\alpha,D)$ using the displayed $-C^{\top}\lambda$ term. Compare against unrolled autodiff through Adam or SGD on few-shot MNIST or CIFAR-10 with a 20-step MAML inner loop. Measure peak memory, wall-clock backward time, meta-gradient cosine similarity to unrolling, and outer validation accuracy. The target is 5--20x lower activation memory at equal inner steps, gradient cosine similarity above 0.99, and no degradation in final validation accuracy.

Verification

Mechanism confirmed, baseline not beaten

Stage 1 · Toy mechanism gate: Passed ✓

Stage 2 · Mechanism transferred to benchmark: Confirmed ✓

Stage 2 · Practical benchmark result: Significantly worse ✗

Methodology: 8 paired seeds, permutation test (p<0.05) against a learning-rate-tuned baseline. How verification works

Stage 1 — Mechanism check agent confidence 8/10

Built and numerically verified a diagonal implicit block solver with reverse local adjoints, plus an autodiff-unrolled reference. The contraction law q=(1+d)/(1+d+alpha*lambda) matched exactly across the sweep, implicit stability persisted beyond the explicit Euler boundary x=2, and finite-horizon gradients matched the predicted K*q^(2K-1) scaling and autodiff to approximately 1.8e-15 relative error. In the vector toy, the local reverse pass was 13.9x faster and used an estimated 10x smaller saved-state footprint, but this is not evidence of a neural-network-level win because the footprint is estimated and the task is diagonal quadratic.

Agent confidence: 8/10

Baseline: Autodiff unroll: gradient cosine 1.000000, relative disagreement 1.77e-15, 0.04435 s; estimated saved states 655,360 bytes.
Idea: Local implicit adjoint: gradient cosine 1.000000, loss disagreement 5.55e-16, 0.00319 s; estimated saved states 65,536 bytes, 10.0x ratio. Contraction and stability predictions matched to <=1.4e-17; normalized finite-horizon gradient prediction was 4.32 for every tested K.

Limitations: No MNIST/CIFAR or neural-network meta-learning benchmark was tested; no actual peak-memory profiler was used, and the local-adjoint implementation covers a scalar/diagonal quadratic block system without neighboring blocks, momentum, Newton iterations, CUDA fallback, or learned optimizer training.

How to run: python3 experiment.py

Stage 2 — Benchmark vs. tuned baseline (arXiv:2608.08559) Sep 2, 2026

Significantly WORSE than the tuned baseline (+17.6%, p=0.0081, wins 0 / 8 paired seeds).

Benchmark: Tabular regression (Friedman)
Model: mlp_tiny
Paired seeds: 8
Baseline mean: 7.1538
Idea mean: 8.4152
Effect (Δ): 1.2615 (+17.6%; negative = idea better)
Wins: 0 / 8 paired seeds
p-value: 0.0081 (permutation test, 20 000 shuffles)
Smallest detectable effect: ±14.4%
Mechanism: Confirmed ✓
Practical effect: Hurts
Baseline tuning: swept over 6 configs

Limitations:

The benchmark used the optimizer-matched tabular track and a small MLP, not meta-learning through many inner tasks. Peak memory was not profiled; the implementation used a diagonal proximal local update without neighboring-block coupling, momentum, Newton iterations, or learned outer hyperparameters.

How to run: python3 stage2_bench.py

Verdict computed by deterministic test code from paired-seed statistics — not by the language model.

Stage 2 — Benchmark Sep 1, 2026 latest Failed ✗

Agent confidence: 9/10

The matched tabular/Friedman#1 mlp_tiny benchmark completed with tuned Adam baseline, equal learning-rate coverage, and 8 paired seeds. The trained-model local implicit displacement prediction was confirmed, but the idea was significantly worse on test MSE, so it does not provide a transfer win.

Baseline: Mean test MSE 7.153750777244568; best lr=0.01, betas=[0.8,0.99].
Comparison: delta_mean=1.2614733576774597, p_value=0.0081, idea_wins=0/8.
Idea: Mean test MSE 8.415224134922028; best lr=0.01, damping=1.0.
Mechanism signature: predicted_step_norm=0.03727731504874156, observed_step_norm=0.037277315091206756, mean_residual=5.612762866054795e-08, confirmed=true.

Benchmark result