# Wasserstein-Budgeted Width Allocation

- ID: 470
- Canonical URL: https://synthcore.org/idea/470/wasserstein-budgeted-width-allocation
- API JSON: https://synthcore.org/api/idea/470.json
- API Markdown: https://synthcore.org/api/idea/470.md
- Verification status: mech_ok_no_baseline
- Source: [arXiv:2607.06290](https://arxiv.org/abs/2607.06290)
- Category: architecture
- Solves: scalability, speedup, memory
- ML areas: architecture, transformer, mlp, rnn, inference-speedup
- Math tags: probability, statistics, linear-algebra, measure-theory
- Ratings: usefulness 7/10; difficulty 5/10; novelty 7/10

## Idea description

Use the paper's finite-width O(n^{-1/2}) Gaussian-process approximation bound as a width-budgeting rule rather than choosing every hidden dimension uniformly. Estimate an architecture-specific constant for each layer or attention contraction, then allocate width according to the smallest dimension satisfying its allowed distributional error. This should produce narrower models at comparable GP-like behavior, or permit the same parameter budget to be concentrated in the layers where finite-width noise is most damaging.

## Mathematical statement

The paper bounds the p-Wasserstein distance between normalized finite-width executions and their infinite-width Gaussian execution by \(\mathcal{W}_{p}(X_n,X_\infty)\leq\mathsf{c}/\sqrt{n}\). Here \(\mathcal{W}_{p}\) is the p-Wasserstein distance between probability laws, \(X_n\) is the vector of finite-network outputs or hidden states normalized by the relevant width, \(X_\infty\) is the corresponding jointly Gaussian-process execution, \(n\) is the hidden contraction width, and \(\mathsf{c}\) is an architecture-, activation-, input-, and p-dependent constant. For a deep MLP, the limiting covariance is recursively defined by \(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)})]\), where \(c^{(\ell)}\) is the weight-variance numerator, \(n_\ell\) is the input width of layer \(\ell\), \(\phi^{(\ell)}\) is the componentwise activation, and \(\mathcal{G}^{(\ell)}\) is the Gaussian preactivation with covariance \(K^{(\ell)}\). The engineering rule is \(n_r\geq(\widehat{\mathsf{c}}_r/\varepsilon_r)^2\) for module r, where \(\widehat{\mathsf{c}}_r\) is an empirically estimated constant and \(\varepsilon_r\) is that module's allocated Wasserstein-error budget.

## Key formulas

- $$\mathcal{W}_{p}(X_n,X_\infty)\leq\frac{\mathsf{c}}{\sqrt{n}}.$$
- $$K^{(1)}=\frac{c^{(0)}}{d}x^{\top}x,\qquad K^{(\ell+1)}=\frac{c^{(\ell)}}{n_\ell}\mathbb{E}\left[\phi^{(\ell)}(\mathcal{G}^{(\ell)})^{\top}\phi^{(\ell)}(\mathcal{G}^{(\ell)})\right].$$
- $$n_r\geq\left(\frac{\widehat{\mathsf{c}}_r}{\varepsilon_r}\right)^2,\qquad \sum_{r=1}^{R}\varepsilon_r\leq\varepsilon_{\mathrm{total}}.$$

## Implementation notes

(1) Integration point: apply this during architecture search to hidden dimensions of MLP layers, the head/key/query contraction dimension in attention, the feed-forward dimension, and the hidden dimension of an RNN. Keep the usual variance scaling \(W_{ij}\sim\mathcal{N}(0,c/n_{\mathrm{in}})\). Do not initially change the forward computation. (2) Pseudocode: for each candidate module r, choose calibration width n_cal; run B random initializations on calibration inputs; collect finite-network outputs; generate corresponding GP samples or a very-wide reference; estimate c_hat[r] = sqrt(n_cal) times the measured distance; choose error budgets eps[r] with sum eps[r] <= eps_total; set width n[r] = ceil((c_hat[r] / eps[r])^2); round widths to hardware-friendly multiples; train the resulting model. (3) Computed from the paper's mathematics: the inverse-square-root scaling and covariance recursion. Estimated empirically: c_hat[r], because the extracted theorem does not provide a universal usable constant. Use sliced Wasserstein distance or a Gaussian mean/covariance distance as a cheap proxy. For shared-weight RNNs, calibrate over the full unroll length rather than treating time steps as independent. (4) First cheap experiment: train 6-layer width-1024 and width-512 MLPs on CIFAR-10, then compare a width-allocated model with the same parameter or FLOP budget. Use 32 to 64 random initializations on 1,000 examples for calibration. Measure validation loss, activation covariance error relative to a width-4096 reference, and training/inference FLOPs. Success is a 10 to 20 percent parameter reduction at matched validation loss, or lower loss at equal parameter count.

## Verification

- Status: mech_ok_no_baseline
- Mechanism evidence: yes
- Mechanism confirmed: yes
- Practical verdict: inconclusive
- Verdict: Built a reproducible Wasserstein-budgeting MVP in experiment.py with analytic Gaussian limits, width sweeps, epsilon inversion, and fixed-budget allocation. The square-activation toy showed the predicted inverse-square-root scaling (observed log-log slope -0.441 versus -0.5) and budget-derived widths tracked target errors reasonably well, with observed distance/epsilon ratios about 0.87–1.04 for epsilon 0.12–0.05. The fixed-budget proxy improved from 0.05495 to 0.05152, but empirical allocated Wasserstein error was slightly worse than uniform (0.07445 versus 0.07253), so there is a mechanism signal but no demonstrated practical allocation win.

### Mechanism check

- Verdict: Built a reproducible Wasserstein-budgeting MVP in experiment.py with analytic Gaussian limits, width sweeps, epsilon inversion, and fixed-budget allocation. The square-activation toy showed the predicted inverse-square-root scaling (observed log-log slope -0.441 versus -0.5) and budget-derived widths tracked target errors reasonably well, with observed distance/epsilon ratios about 0.87–1.04 for epsilon 0.12–0.05. The fixed-budget proxy improved from 0.05495 to 0.05152, but empirical allocated Wasserstein error was slightly worse than uniform (0.07445 versus 0.07253), so there is a mechanism signal but no demonstrated practical allocation win.
- Confidence: 6/10
- Limitations: This is only a scalar normalized-CLT initialization toy, not a trained MLP or CIFAR-10 experiment. ReLU scaling was noisier (slope -0.311), empirical Wasserstein estimates have a sampling floor, and allocation was not validated on downstream loss, FLOPs, latency, or a real deep architecture.

### Practical benchmark

- Paired seeds: 8
- Baseline mean: 0
- Idea mean: 0
- p-value: 0.70265
- Paired wins: 3/8
- Benchmark verdict: no significant win

## Artifacts

- [bench_report.json](https://synthcore.org/code/235/bench_report.json)
- [bench_width_allocation.py](https://synthcore.org/code/235/bench_width_allocation.py)
- [experiment.py](https://synthcore.org/code/235/experiment.py)
- [report.md](https://synthcore.org/code/235/report.md)
- [report_bench_2026-08-30T211830.md](https://synthcore.org/code/235/report_bench_2026-08-30T211830.md)
- [results.json](https://synthcore.org/code/235/results.json)
- [Download all files as ZIP](https://synthcore.org/download/235)

## Disclaimer

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