# Geometry-Guided Layerwise FFN Width Allocation in Transformers

arXiv: https://arxiv.org/abs/2608.02064 · HTML: https://mlsanity.com/paper/2608.02064/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 10 of its 23 references that are in the map.
- First posted on arXiv 3 Aug 2026.

## Builds on (most cited first, 10 of 10)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [Root Mean Square Layer Normalization](https://mlsanity.com/paper/1910.07467/) (2019)
- [On Layer Normalization in the Transformer Architecture](https://mlsanity.com/paper/2002.04745/) (2020)
- [Computational Optimal Transport](https://mlsanity.com/paper/1803.00567/) (2018)
- [Error bounds for approximations with deep ReLU networks](https://mlsanity.com/paper/1610.01145/) (2016)
- [Intrinsic dimension of data representations in deep neural networks](https://mlsanity.com/paper/1905.12784/) (2019)
- [Estimating the intrinsic dimension of datasets by a minimal neighborhood information](https://mlsanity.com/paper/1803.06992/) (2018)
- [Residual Connections Encourage Iterative Inference](https://mlsanity.com/paper/1710.04773/) (2017)

## Built on by (most cited first, 1 of 1)

- [Shaping Latent Mixture-of-Experts: From Expert Geometry to Layerwise Capacity Allocation](https://mlsanity.com/paper/2610.14644/) (2026)

## Abstract (arXiv)

Feed-forward networks (FFNs) account for a large fraction of Transformer parameters, yet their hidden width is usually constant across depth. We ask whether this capacity can instead be allocated from a forward-pass measurement of layer behavior. We view each FFN as transporting a cloud of token representations and quantify the induced geometric change using correspondence-preserving shift, Gromov-Wasserstein distortion, and degree-one persistent homology under raw and scale-normalized metrics. A layerwise approximation surrogate yields an exact fixed-budget optimizer. Across seven pretrained language models, raw Euclidean work largely tracks residual-norm growth, whereas normalized work is predominantly front-loaded. Gromov-Wasserstein work is more consistently associated with perturbation-based layer sensitivity than the finite-sample topological estimate. In paired 128M and 256M training runs, several normalized-work schedules reduce mean validation loss relative to both uniform width and a hand-designed cosine taper. With the amplified paired differences at 440M, the best geometry-based allocations improve over uniform substantially larger than the cosine taper, while the anti-topological raw control is worse than uniform.