aiqbench analysis
This research proposes allocating Transformer FFN widths based on geometric properties of token representations rather than using uniform or heuristic schedules like cosine tapering. By quantifying the induced geometric change in hidden states, the method demonstrates measurable improvements in validation loss across various model sizes compared to standard allocation strategies. The findings suggest that geometry-guided capacity planning offers a more efficient path for optimizing large-scale language models without increasing parameter counts.
Geometry-Guided Layerwise FFN Width Allocation in Transformers
arXiv:2608.02064v1 Announce Type: cross 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.
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.
Machine Learning (cs.LG); Artificial Intelligence (cs.AI); Computation and Language (cs.CL)
arXiv:2608.02064 [cs.LG]
(or arXiv:2608.02064v1 [cs.LG] for this version)
https://doi.org/10.48550/arXiv.2608.02064 Focus to learn more arXiv-issued DOI via DataCite (pending registration)
arXivLabs is a framework that allows collaborators to develop and share new arXiv features directly on our website.
Both individuals and organizations that work with arXivLabs have embraced and accepted our values of openness, community, excellence, and user data privacy. arXiv is committed to these values and only works with partners that adhere to them.
Have an idea for a project that will add value for arXiv's community? Learn more about arXivLabs.
Full text under the publisher's open licence · 263 words
Related coverage
Shaping Latent Mixture-of-Experts: From Expert Geometry to Layerwise Capacity Allocation
arXiv · Oct 15, 2026
PIRAMID: Progress and Plans
LessWrong · Jul 27, 2026
MixFrag: Fragility-Guided Mixed-Precision Post-Training Quantization for Vision Transformers
arXiv · Jul 31, 2026
Predictive Geometry of Hidden Trajectories in Transformers
arXiv · Oct 1, 2026