🤖 machine learning

Geometry-Guided Layerwise FFN Width Allocation in Transformers

This paper proposes a geometry-guided method to dynamically allocate Feed-Forward Network (FFN) width across Transformer layers based on geometric metrics like Gromov-Wasserstein distortion and persistent homology, demonstrating that such data-driven schedules significantly reduce validation loss compared to uniform or hand-designed width allocations.

Original authors: Timur Mudarisov, Mikhail Burtsev, Radu State

Published 2026-08-04

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Timur Mudarisov, Mikhail Burtsev, Radu State

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are building a giant, multi-story library where every floor is a room dedicated to organizing a specific type of information. In the world of artificial intelligence, these "libraries" are called Transformers, and they are the brains behind many modern chatbots and language tools. Inside each room, there is a team of workers (called a Feed-Forward Network, or FFN) whose job is to take the information coming in, think about it, and pass it along. For a long time, engineers built these libraries with a strict rule: every single room gets the exact same number of workers, no matter what floor it is on. It's like giving the same amount of budget to the basement storage room as you do to the fancy top-floor reading lounge.

But what if that rule is wasteful? What if the basement needs fewer workers because the boxes are simple, while the top floor needs a huge team to handle complex, abstract ideas? This is the question researchers have started asking. They know that as information travels through the library, it changes shape and gets more complex. The big question is: can we measure exactly how much "work" each room is doing and then move workers around to where they are needed most, without hiring any extra people? This paper dives into that idea, using a mix of geometry (the study of shapes and distances) and math to figure out the perfect staffing plan for each room.

The authors of this paper, Timur Mudarisov, Mikhail Burtsev, and Radu State, decided to stop guessing and start measuring. They treated the information flowing through the AI like a "cloud of points" (think of it as a swarm of fireflies moving through a room). As the fireflies move from one room to the next, they shift, stretch, and twist. The team developed a way to measure exactly how much "geometric work" happens in each room. They used three different ways to measure this: looking at how far the fireflies moved (shift), how the distances between them changed (Gromov-Wasserstein), and how the loops and holes in their formation changed (topology).

Here is the twist: they found that how you measure matters. If you just look at raw distances, the later rooms seem to do the most work, but that's often just because the fireflies are getting bigger (a scale issue). However, when they normalized the data—looking at the shape of the movement rather than just the size—they discovered a different pattern. The "work" is actually highest at the very beginning of the library and tapers off as you go deeper. This suggests that the early rooms need the most brainpower, while the later rooms can get by with fewer workers.

To test this, they created a new "staffing schedule" based on these measurements. Instead of giving every room the same number of workers, they gave the early rooms more and the later rooms less, all while keeping the total number of workers exactly the same as the old, uniform model. They tested this on several different AI models, ranging from small ones (128 million parameters) to a larger one (440 million parameters).

The results were promising. In their experiments, the models that used this "geometry-guided" staffing plan performed better than the standard models where every room was equal. In fact, when they compared their new method to a popular, hand-designed plan called a "cosine taper" (which simply assumes work decreases smoothly like a slide), their data-driven approach did even better. At the 440 million parameter scale, their geometry-based plan reduced the model's errors significantly more than the cosine plan did.

However, the authors are careful not to claim they have solved the problem forever. They note that their method relies on measuring a "reference" model and then applying those rules to a new one. While the results suggest that moving capacity to the early layers is a winning strategy, they admit that more testing is needed to be absolutely certain about which specific measurement tool (topology vs. geometry) is the best. They also point out that their "anti-topological" control group (a plan that deliberately did the opposite of what the data suggested) performed worse than the standard model, which confirms that the direction of their new plan is indeed correct.

In short, this paper suggests that AI models don't need to be built with equal rooms. By using geometry to measure where the real work happens, we can rearrange the workers to make the whole library smarter, without spending a single extra dollar on hiring. It's a clever, data-driven way to squeeze more performance out of the same amount of computing power.

Technical Summary: Geometry-Guided Layerwise FFN Width Allocation in Transformers

Problem Statement
In standard Transformer architectures, the hidden width of Feed-Forward Networks (FFNs) is typically constant across all layers, despite evidence that layer functions and representation dynamics change with depth. While recent work (Tapered Language Models) demonstrated that reallocating capacity from deeper to shallower layers improves performance, existing methods rely on hand-designed, monotonic schedules (e.g., cosine tapers). These prescribed shapes do not adapt to the specific computational demands of a given model or data distribution. The paper asks whether a data-dependent, layer-specific width allocation can be derived directly from the geometric behavior of activations during a forward pass, rather than relying on a fixed architectural prior.

Methodology
The authors propose a framework that treats FFN layers as transporters of token representation clouds. The core methodology involves three stages:

Geometric Work Measurement:
The paper defines "geometric work" as the change induced by an FFN layer on a cloud of token representations. To ensure the metric reflects the FFN's actual contribution rather than residual scaling artifacts, the authors evaluate three distinct work statistics under three geometric frameworks:
- Raw Geometry: Uses Euclidean distance on unnormalized states. This is sensitive to scale and serves as a control.
- Spherical Geometry: Uses unit-normalized states and angular distance.
- Hyperbolic Geometry: Maps unit-normalized states into a Poincaré ball, acting as a nonlinear reweighting of angular distances.
For each geometry, three work statistics are computed:
- Shift: The mean displacement of tokens (correspondence-preserving).
- Metric Distortion: Measured via Gromov-Wasserstein (GW) distance, quantifying how pairwise relations are preserved or altered.
- Topological Change: Measured via the 2-Wasserstein distance between degree-one persistent homology diagrams (Vietoris–Rips filtration), capturing the creation or destruction of loops.
Approximation Surrogate and Optimization:
The authors introduce a surrogate model to map geometric work to optimal layer widths. They posit that the approximation error of an FFN layer scales with its width according to an intrinsic dimension ( $d_{int}$ ) and a loss sensitivity ( $\lambda$ ).
- Assumptions: (1) Loss effects are locally additive; (2) Error scales as $w^{-a}$ where $a$ depends on intrinsic dimension; (3) Measured geometric work is a proxy for the difficulty of approximation.
- Optimizer: A constrained optimization problem is solved to minimize the sum of weighted approximation errors subject to a fixed total parameter budget. This yields a closed-form solution for the optimal width $w_\ell$ for each layer, incorporating lower bounds and depth-dependent exponents.
Schedule Construction:
A reference model (uniform width) is run to collect activation clouds. The geometric work profile is estimated, smoothed, and mapped to a width schedule using the derived optimizer. This schedule is then applied to train a new model from scratch.

Key Contributions

Definition of Geometric Work: The paper formalizes three complementary forward-pass work statistics (shift, GW distortion, topological change) under raw, spherical, and hyperbolic metrics, distinguishing between raw magnitude and scale-normalized structural changes.
Exact Budget-Optimal Solver: It derives a unique continuous optimizer for layerwise width allocation under a fixed parameter budget, handling lower bounds and varying intrinsic dimensions across layers.
Comprehensive Diagnostics: The study evaluates scale sensitivity, local loss additivity, intrinsic dimension, and the correlation between geometric work and perturbation-based layer sensitivity across seven pretrained models (Llama 3.1/3.2, Gemma 2/3, Mistral).
Empirical Validation: The paper compares geometry-derived schedules against uniform width, a hand-designed cosine taper, and an "anti-topological" control across three model scales (128M, 256M, and 440M).

Results

Scale Sensitivity: Raw geometric work (Euclidean) largely tracks residual norm growth, whereas normalized metrics (spherical/hyperbolic) reveal a "front-loaded" profile where the earliest layers exhibit the highest geometric work.
Sensitivity vs. Allocation: While Gromov-Wasserstein work shows a stronger correlation with perturbation-based layer sensitivity, topological/hyperbolic allocation yields the lowest mean validation loss in training runs. This suggests that sensitivity and approximation difficulty are distinct factors.
Training Performance:
- In 128M and 256M runs, normalized geometry schedules generally outperform uniform width and the anti-topological control, though results are exploratory due to limited seeds.
- In the 440M five-seed experiment, the topological/hyperbolic schedule achieved a mean validation loss reduction of 0.019 relative to uniform width. This is approximately 6.3 times the reduction achieved by the TLM-style cosine taper (0.003).
- The anti-topological/raw control performed worse than uniform width, confirming that the direction of the allocation matters.

Significance and Claims
The paper claims that geometric work provides a principled, data-driven method to replace global FFN width constants with layer-dependent schedules. It establishes that:

FFN capacity should not be assumed uniform; layer updates change roles across depth.
Fixed monotonic tapers are not the only effective alternative; activation-derived profiles can capture non-monotonic structures and achieve larger performance gains than hand-designed schedules under identical conditions.
Normalized geometric changes (specifically topological/hyperbolic) reveal a transferable early-heavy profile that improves model efficiency.

The authors remain modest in their conclusions, noting that the results are descriptive mean comparisons until paired-difference uncertainty is fully reported. They do not claim to have proven a universal superiority of one specific geometry (e.g., topology vs. GW) but rather that normalized geometric change is a useful signal for capacity allocation. The work positions activation-derived allocation as a practical extension of tapering, motivating future work to directly calibrate geometric work to layerwise approximation demand.

Technical Summary: Geometry-Guided Layerwise FFN Width Allocation in Transformers

More like this