Geometry-Guided Layerwise FFN Width Allocation in Transformers
2026-08-03 • Machine Learning
Machine LearningArtificial IntelligenceComputation and Language
AI summaryⓘ
The authors studied how the size of certain parts (called feed-forward networks) in Transformer models can be better assigned rather than kept the same throughout the model. They measured changes in the shape of data as it passes through each layer using geometric and topological methods. Their experiments on several language models showed that adjusting layer sizes based on these shape changes can improve model performance compared to fixed or hand-tuned sizes. They found that some geometry-based methods work better than others for guiding this adjustment.
Feed-forward networksTransformer modelsLayer widthGromov-Wasserstein distortionPersistent homologyLayer sensitivityResidual normValidation lossModel optimizationGeometry in machine learning
Authors
Timur Mudarisov, Mikhail Burtsev, Radu State
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.