Finite-Width NNGP Covariance Stabilizer
Implementation & benchmark of arXiv:2607.06290 — Quantitative Gaussian-Process limits of Tensor Programs
Source paper: Quantitative Gaussian-Process limits of Tensor Programs arXiv:2607.06290 ⓘ · analyzed Aug 30, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
Add a training-time regularizer that keeps the empirical joint covariance of hidden activations on multiple inputs close to the recursively predicted NNGP covariance. The regularizer targets the finite-width fluctuations quantified by the Wasserstein result, and is particularly appropriate for recurrent networks and attention blocks with shared weights, where hidden states at different positions or time steps are statistically coupled.
Formulas
Mathematical statement
For two inputs, the paper gives \(\mathcal{W}_{p}((y^{1}/\sqrt{d_{\rm out}},y^{2}/\sqrt{d_{\rm out}}),\overline{\mathcal{G}})\leq\mathsf{c}/\sqrt{n}\), where \(y^1,y^2\) are finite-network outputs, \(d_{\rm out}\) is output width, \(n\) is hidden width, and \(\overline{\mathcal{G}}=(\mathcal{G}^1,\mathcal{G}^2)\) is the jointly Gaussian pair generated by the same shared weights. The NNGP covariance is obtained from \(K^{(1)}=(c^{(0)}/d)x^{\top}x\) and \(K^{(\ell+1)}=(c^{(\ell)}/n_\ell)\mathbb{E}[\phi^{(\ell)}(\mathcal{G}^{(\ell)})^{\top}\phi^{(\ell)}(\mathcal{G}^{(\ell)})]\). For a minibatch \(X=(x_1,\ldots,x_B)\), let \(H_r\in\mathbb{R}^{B\times n_r}\) be the normalized activation matrix at module r, and let \(C_r=H_rH_r^{\top}/n_r\) be its empirical input-input covariance. Let \(K_r(X,X)\) be the NNGP covariance computed by the recursion. Use \(L_{\rm cov}=\sum_r\alpha_r\|C_r-K_r(X,X)\|_F^2\), with \(\alpha_r\) scaled by \(n_r\) or by the estimated variance of \(C_r\).
Implementation notes
(1) Integration point: insert the regularizer after selected preactivation or postactivation tensors in a Transformer attention block, an RNN hidden-state stack, or a small MLP. Use a calibration minibatch of B = 16 to 64 examples and compute the B by B covariance over the feature dimension. For sequence models, index rows by pooled per-example representations initially, then test token-level covariance. (2) Pseudocode: compute hidden states H = forward(model, X); initialize L_cov = 0; for each selected layer r, normalize H_r across features, compute C_r = H_r @ H_r.T / n_r, compute K_r with the NNGP recursion, and add alpha[r] times squared Frobenius norm of C_r minus stop_gradient(K_r) to L_cov; optimize task_loss + lambda_cov times L_cov. (3) Compute from the paper's mathematics: the target covariance and the O(n^{-1/2}) finite-width scale. Estimate empirically: alpha[r], lambda_cov, minibatch covariance noise, and whether pre- or postactivation targets work best. For ReLU, use the analytic ReLU covariance; for GELU, use Gauss-Hermite quadrature or 32 to 128 Monte Carlo samples for each bivariate Gaussian expectation. In an RNN or shared-weight attention module, preserve repeated-weight covariance rather than assuming independent layers. (4) First cheap experiment: use a 4-layer width-256 Transformer encoder on CIFAR-10 patches or a small language-model corpus, comparing baseline training with lambda_cov in {10^-4, 10^-3, 10^-2}. Also train width-128 with and without the regularizer and compare against width-256 baseline. Track task loss, gradient norm variance, covariance deviation, and validation accuracy or perplexity. The desired signal is lower covariance deviation and smoother gradients, followed by recovery of the wider model's accuracy at half width.
Verification
Stage 1 · Toy mechanism gate: Passed ✓
Stage 2 · Mechanism transferred to benchmark: Not reproduced ✗
Stage 2 · Practical benchmark result: No measurable effect
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 training-time NNGP covariance stabilizer for a small ReLU MLP, using the analytic ReLU Gaussian covariance recursion and a detached Frobenius covariance-matching penalty. The numerical check matched Monte Carlo ReLU covariance (0.6775 vs 0.6785) and empirical covariance error scaled as width^-0.542, consistent with the claimed O(n^-1/2) behavior. Across three fixed seeds, regularization reduced held-out calibration covariance deviation from mean 0.0699 to 0.0550 (21.3% lower) and validation MSE from 0.3581 to 0.3481 (2.8% lower), so the effect is promising but modest.
- Agent confidence
- 8/10
- Baseline
- Mean over 3 seeds: validation MSE 0.35808; final covariance deviation 0.06993.
- Idea
- Mean over 3 seeds: validation MSE 0.34808; final covariance deviation 0.054997. Covariance deviation improved 21.3% and validation MSE improved 2.8%. Width-scaling log-log slope was -0.542; analytic versus Monte Carlo ReLU covariance was 0.67755 versus 0.67852.
Limitations: Only a 3-layer width-64 MLP and a synthetic regression task were tested; no Transformer, RNN, attention, token-level covariance, GELU, width-128-versus-width-256 comparison, lambda sweep, FLOP-matched comparison, or gradient-variance analysis was performed. The covariance target is analytic and detached, and the experiment does not test the full shared-weight recurrent/attention setting.
How to run: python3 experiment.py
Stage 2 — Benchmark vs. tuned baseline (arXiv:2607.06290)
No significant difference vs. the tuned baseline (+0.1%, p=0.2419, wins 3 / 8 paired seeds).
- Benchmark
- Sequence modeling
- Model
- transformer_tiny
- Paired seeds
- 8
- Baseline mean
- 0.3383
- Idea mean
- 0.3386
- Effect (Δ)
- 0.0003 (+0.1%; negative = idea better)
- Wins
- 3 / 8 paired seeds
- p-value
- 0.2419 (permutation test, 20 000 shuffles)
- Smallest detectable effect
- ±0.1%
- Mechanism
- Not confirmed ✗
- Practical effect
- No effect
- Baseline tuning
- swept over 3 configs
The advertised bench/README.md path was unavailable, although the injected bench package and APIs were importable. Only the structurally matched sequence track was tested; no vision, tabular, or dynamics transfer was run. The intervention regularized pooled input-projection states rather than token-level Transformer encoder states, and no width comparison, gradient-variance analysis, FLOP profiling, or broader lambda sweep was performed.
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 latest Failed ✗
Agent confidence: 9/10
Built a local training-time NNGP covariance stabilizer for the benchmark's structurally matched sequence forecast task using the tiny Transformer and an added pooled hidden-state covariance penalty. Across eight paired seeds, the best idea setting had MSE 0.338577 versus the tuned baseline's 0.338322, with paired delta +0.000255 and permutation p=0.2419; this is no measurable effect and is slightly worse. The trained-model mechanism signature was not confirmed: covariance deviation was 0.10427585 for baseline versus 0.10427427 for the idea, ratio 0.999985.
- Baseline
- 8-seed tuned baseline mean test MSE 0.3383217715, std 0.0461405308; best lr=0.001. Sweep means: lr=0.001 -> 0.3383218, lr=0.003 -> 0.3591854, lr=0.009 -> 1.0258603.
- Idea
- Best of 3 settings mean test MSE 0.3385768682, std 0.0462719048 at lr=0.001, lambda_cov=0.0001. Paired delta idea-baseline +0.0002550967, 3/8 idea wins, permutation p=0.2419. Mechanism covariance deviation: baseline 0.1042758487, idea 0.1042742701, ratio 0.9999849, confirmed=false.
Benchmark result
No significant difference vs. the tuned baseline (+0.1%, p=0.2419, wins 3 / 8 paired seeds).
- Benchmark
- Sequence modeling
- Model
- transformer_tiny
- Paired seeds
- 8
- Baseline mean
- 0.3383
- Idea mean
- 0.3386
- Effect (Δ)
- 0.0003 (+0.1%; negative = idea better)
- Wins
- 3 / 8 paired seeds
- p-value
- 0.2419 (permutation test, 20 000 shuffles)
- Smallest detectable effect
- ±0.1%
- Mechanism
- Not confirmed ✗
- Practical effect
- No effect
- Baseline tuning
- swept over 3 configs
The advertised bench/README.md path was unavailable, although the injected bench package and APIs were importable. Only the structurally matched sequence track was tested; no vision, tabular, or dynamics transfer was run. The intervention regularized pooled input-projection states rather than token-level Transformer encoder states, and no width comparison, gradient-variance analysis, FLOP profiling, or broader lambda sweep was performed.
How to run: python3 stage2_bench.py
Artifacts
- bench_report.json 3.6 KB View Raw
- experiment.py 5.3 KB View Raw
- report.md 1.6 KB View
- report_bench_2026-09-04T165923.md 4.5 KB View
- results.json 2.7 KB View Raw
- stage2_bench.py 4.7 KB View Raw