# Finite-Width NNGP Covariance Stabilizer

- ID: 471
- Canonical URL: https://synthcore.org/idea/471/finite-width-nngp-covariance-stabilizer
- API JSON: https://synthcore.org/api/idea/471.json
- API Markdown: https://synthcore.org/api/idea/471.md
- Verification status: failed_benchmark
- Source: [arXiv:2607.06290](https://arxiv.org/abs/2607.06290)
- Category: regularization
- Solves: stability, generalization, accuracy
- ML areas: transformer, attention, rnn, regularization, loss
- Math tags: probability, statistics, linear-algebra, measure-theory
- Ratings: usefulness 6/10; difficulty 5/10; novelty 6/10

## 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.

## 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\).

## Key formulas

- $$\mathcal{W}_{p}\left(\left(\frac{y^{1}}{\sqrt{d_{\rm out}}},\frac{y^{2}}{\sqrt{d_{\rm out}}}\right),\overline{\mathcal{G}}\right)\leq\frac{\mathsf{c}}{\sqrt{n}}.$$
- $$K^{(1)}(x_i,x_j)=\frac{c^{(0)}}{d}x_i^{\top}x_j,\qquad K^{(\ell+1)}(x_i,x_j)=\frac{c^{(\ell)}}{n_\ell}\mathbb{E}_{(u,v)\sim\mathcal{N}(0,\left[\begin{smallmatrix}K^{(\ell)}_{ii}&K^{(\ell)}_{ij}\\K^{(\ell)}_{ji}&K^{(\ell)}_{jj}\end{smallmatrix}\right])}[\phi^{(\ell)}(u)\phi^{(\ell)}(v)].$$
- $$C_r=\frac{1}{n_r}H_rH_r^{\top},\qquad L_{\rm cov}=\sum_{r\in\mathcal{R}}\alpha_r\left\|C_r-K_r(X,X)\right\|_F^2.$$

## 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

- Status: failed_benchmark
- Mechanism evidence: yes
- Mechanism confirmed: no
- Practical verdict: no_effect
- Verdict: 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.

### Mechanism check

- Verdict: 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.
- Confidence: 8/10
- 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.

### Practical benchmark

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

## Artifacts

- [bench_report.json](https://synthcore.org/code/1315/bench_report.json)
- [experiment.py](https://synthcore.org/code/1315/experiment.py)
- [report.md](https://synthcore.org/code/1315/report.md)
- [report_bench_2026-09-04T165923.md](https://synthcore.org/code/1315/report_bench_2026-09-04T165923.md)
- [results.json](https://synthcore.org/code/1315/results.json)
- [stage2_bench.py](https://synthcore.org/code/1315/stage2_bench.py)
- [Download all files as ZIP](https://synthcore.org/download/1315)

## Disclaimer

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