Geometry-Guided Layerwise FFN Width Allocation in Transformers
Timur Mudarisov, Mikhail Burtsev, Radu State
cs.LG, cs.AI, cs.CL
Submitted: 2026-08-03
Project page: http://skylion007.github.io
License: http://creativecommons.org/licenses/by/4.0/
The gist: Feed-forward networks (FFNs) account for a large fraction of Transformer parameters, yet their hidden width is usually constant across depth.
Terminology
Abstract
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.
Sources
- Layer Normalization
- The Llama 3 Herd of Models
- Gemma 2: Improving Open Language Models at a Practical Size
- Gemma 3 Technical Report
- Mistral 7B
- Attention Is All You Need
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks