Space-Time Onsager Optimizer
Implementation & benchmark of arXiv:2608.04920 — Resolving coupled transport in space and time from molecular fluctuations in confined fluids
Source paper: Resolving coupled transport in space and time from molecular fluctuations in confined fluids arXiv:2608.04920 ⓘ · analyzed Aug 31, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
Replace an instantaneous diagonal optimizer with a causal convolution of recent gradients, where cross-layer or cross-module gradient correlations define a finite-memory Onsager response matrix. Estimate the response at several parameter-block pairs and lags, integrate it to obtain a finite-time transport matrix, and use its regularized inverse or symmetric part to precondition the update. This targets optimization regimes in which gradients propagate between blocks with measurable delay, such as deep transformers, recurrent networks, and diffusion models.
Formulas
Mathematical statement
Let $g_a(k)\in\mathbb{R}^{d_a}$ be the gradient of the loss at optimization step $k$ for parameter block $a$, and let $\tau$ denote a nonnegative lag. The paper's spatially resolved response kernel $\mathbf{R}_{ab}(\mathbf r,\mathbf r',t)$ is adapted to a block-and-lag kernel $R_{ab}(\tau)$ estimated from gradient fluctuations. Its local-global and global-global reductions become $K_{ab}(\tau)=\sum_b R_{ab}(\tau)$ and $C_{ab}(\tau)=\sum_a K_{ab}(\tau)$, while the finite-time integrated response is $\mathcal M_{ab}(T)=\sum_{\tau=0}^{T}K_{ab}(\tau)$. The long-time Onsager matrix is $M_{ab}=\lim_{T\to\infty}\mathcal M_{ab}(T)$ when correlations are summable. Use a symmetrized, ridge-regularized matrix $P=(M+M^{\mathsf T})/2+\lambda I$ as a preconditioner, or use its stable inverse square root. For a locally quadratic loss with Hessian $H$, the update $\theta_{k+1}=\theta_k-\eta P^{-1}g_k$ is linearly stable when $0<\eta<2/\lambda_{\max}(P^{-1/2}HP^{-1/2})$. Persistent and rapidly relaxing response modes correspond to different learned gradient-memory times.
Implementation notes
1. Integration point: partition the network parameters into blocks, such as attention projections, MLP projections, normalization parameters, or recurrent state matrices. Maintain a FIFO queue of the last $T$ block gradients after flattening each block and optionally projecting it to a small sketch dimension with a fixed random matrix. Estimate lagged cross-block covariances online using exponential averaging, avoiding full parameter-sized matrices. Begin with one scalar covariance per block pair or low-rank factors; use blockwise diagonal or Kronecker approximations for scalability. 2. Pseudocode: compute gradients $g_{a,k}$; update means $\bar g_a$ and covariance estimates $R_{ab}(\tau)\leftarrow(1-\beta)R_{ab}(\tau)+\beta(g_{a,k}-\bar g_a)(g_{b,k-\tau}-\bar g_b)^{\mathsf T}$; form $M=\sum_{\tau=0}^{T}w_\tau R(\tau)$; set $P=(M+M^{\mathsf T})/2+\lambda I$; solve $P d_k=g_k$ with conjugate gradients or a blockwise inverse; update $\theta\leftarrow\theta-\eta d_k$. A safer variant adds the causal memory term in the fourth formula and clips its norm. 3. Computed from the paper's construction: the kernel hierarchy, finite-time integration, and separation of transient from persistent modes. Estimated empirically: gradient covariances, correlation decay, memory horizon $T$, and the largest eigenvalue of $P^{-1/2}HP^{-1/2}$ using Hessian-vector products. Use backtracking if the quadratic model increases or an estimated spectral radius exceeds one. 4. First cheap experiment: train a 12-layer GPT-style transformer on a small character-level WikiText subset, comparing AdamW, block-diagonal Onsager preconditioning, and the causal-memory variant at matched compute. Use 8 to 16 blocks, $T\in\{4,8,16,32\}$, and low-rank gradient sketches. Measure gradient cross-correlations, loss autocorrelation, wall-clock convergence, and Hessian-vector spectral estimates. The quantitative prediction is that $\widehat M(T)$ plateaus once $T$ exceeds the dominant gradient-correlation time; larger $T$ should provide no further conditioning gain and may destabilize training. On a locally quadratic probe, the measured maximum stable learning rate should agree within 20% with $\eta_c=2/\lambda_{\max}(P^{-1/2}HP^{-1/2})$. Removing cross-block lag terms should eliminate most improvement when off-diagonal correlations are weak, but should significantly reduce convergence time when they are strong.
Verification
Stage 1 · Toy mechanism gate: Passed ✓
Stage 2 · Mechanism transferred to benchmark: Not reproduced ✗
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 a reproducible two-block Space-Time Onsager toy probe in onsager_probe.py. The exact finite-memory kernel matched the predicted tail q^(T+1) to 2e-10 relative error, and the measured quadratic stability transition occurred between 0.99 and 1.01 times the predicted eta_c=0.96555. Cross-block preconditioning showed no gain at rho=0 and increasingly strong conditioning/convergence gains as rho rose, with idealized equal-step speedups from 6x at rho=0.2 to 180x at rho=0.95. This establishes the mechanism in a controlled quadratic setting, but not a practical optimizer win.
- Agent confidence
- 8/10
- Baseline
- Diagonal/no-cross-block preconditioner: condition numbers 1.0, 1.5, 3.0, 9.0, 39.0 and estimated optimal-step counts 1, 12, 27, 83, 360 for rho=0, .2, .5, .8, .95.
- Idea
- Full response preconditioner: condition numbers approximately 1.0 across all rho and estimated optimal-step counts 1, 2, 2, 2, 2; speedup is 1x at rho=0, then 6x, 13.5x, 41.5x, and 180x. Exact memory-tail prediction max relative error was 1.9e-10; stability prediction classified 0.90x and 0.99x eta_c stable, 1.01x and 1.10x unstable.
Limitations: No transformer or neural-network training was run. The cross-block comparison uses an idealized response matrix equal to the quadratic Hessian rather than an online learned covariance preconditioner. Finite-sample empirical covariance tails were inaccurate when q^(T+1) was extremely small, despite the exact-kernel prediction being confirmed.
How to run: python3 onsager_probe.py
Stage 2 — Benchmark vs. tuned baseline (arXiv:2608.04920)
Significantly WORSE than the tuned baseline (+948.5%, p=0.0081, wins 0 / 8 paired seeds).
- Benchmark
- Tabular regression (Friedman)
- Model
- mlp_tiny
- Paired seeds
- 8
- Baseline mean
- 15.1542
- Idea mean
- 158.8961
- Effect (Δ)
- 143.7419 (+948.5%; negative = idea better)
- Wins
- 0 / 8 paired seeds
- p-value
- 0.0081 (permutation test, 20 000 shuffles)
- Smallest detectable effect
- n/a
- Mechanism
- Not confirmed ✗
- Practical effect
- Hurts
- Baseline tuning
- swept over 0 configs
Only the tabular optimizer track was tested. Vision, sequence, and dynamics transfer were not evaluated; the implementation used scalar block summaries rather than full low-rank cross-block matrices.
How to run: python3 bench_onsager.py
Verdict computed by deterministic test code from paired-seed statistics — not by the language model.
Stage 2 — Benchmark latest Failed ✗
Agent confidence: 9/10
Implemented a causal block-gradient Space-Time Onsager optimizer and evaluated it against matched Adam on the structurally appropriate tabular Friedman#1 benchmark. Across 8 paired seeds, the idea was significantly worse: mean test MSE 158.896 versus 15.154, paired delta +143.742, permutation p=0.0081. The horizon-4 mechanism signature was not confirmed, with observed integrated-response fraction 0.676.
- Baseline
- Adam, selected lr=0.006 and weight_decay=0.0001; mean test MSE 15.154220581054688, std 2.16180483804436.
- Comparison
- delta_mean=+143.74190545082092; idea_wins=0/8; permutation p_value=0.0081; verdict=idea worse (significant).
- Idea
- Space-Time Onsager, selected lr=0.006, weight_decay=0.0001, horizon=4; mean test MSE 158.8961260318756, std 73.99530089820898.
Benchmark result
Significantly WORSE than the tuned baseline (+948.5%, p=0.0081, wins 0 / 8 paired seeds).
- Benchmark
- Tabular regression (Friedman)
- Model
- mlp_tiny
- Paired seeds
- 8
- Baseline mean
- 15.1542
- Idea mean
- 158.8961
- Effect (Δ)
- 143.7419 (+948.5%; negative = idea better)
- Wins
- 0 / 8 paired seeds
- p-value
- 0.0081 (permutation test, 20 000 shuffles)
- Smallest detectable effect
- n/a
- Mechanism
- Not confirmed ✗
- Practical effect
- Hurts
- Baseline tuning
- swept over 0 configs
Only the tabular optimizer track was tested. Vision, sequence, and dynamics transfer were not evaluated; the implementation used scalar block summaries rather than full low-rank cross-block matrices.
How to run: python3 bench_onsager.py
Artifacts
- bench_onsager.py 7.5 KB View Raw
- bench_report.json 2.7 KB View Raw
- onsager_probe.py 4.1 KB View Raw
- report.md 1.7 KB View
- report_bench_2026-08-31T210231.md 2.9 KB View
- results.json 6.3 KB View Raw