Geometry-Guided Layerwise FFN Width Allocation in Transformers

2026-08-03

AI Summary

This work addresses the limitation of fixed-width feedforward networks (FFNs) in Transformers, which ignore the varying representational transformation demands across layers. The authors propose a novel resource allocation paradigm grounded in the geometric evolution of representations: by quantifying the geometric properties—such as Gromov–Wasserstein distortion and first-order persistent homology—of token representation clouds during forward propagation, they estimate the capacity required at each layer and optimize a non-uniform FFN width distribution under a fixed parameter budget. Experimental results demonstrate that this approach significantly outperforms uniform allocation and handcrafted strategies like cosine decay across multiple pretrained language models, with particularly pronounced gains observed in a 440M-parameter model.

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.

Problem

• Transformer
• Feed-forward Network
• Layerwise Allocation
• Model Width
• Geometric Analysis

Innovation

• geometry-guided allocation
• layerwise FFN width
• Gromov-Wasserstein distortion
• persistent homology
• fixed-budget optimization