💬 NLP

Revisiting the Shape Convention of Transformer Language Models

This paper challenges the conventional narrow-wide-narrow MLP design in Transformers by proposing a deeper hourglass-shaped FFN, demonstrating through empirical validation that this architecture not only matches or exceeds the performance of standard models but also enables more efficient parameter allocation, such as expanding hidden dimensions, to achieve superior results within fixed budgets.

Original authors: Feng-Ting Liao, Meng-Hsi Chen, Guan-Ting Yi, Da-shan Shiu

Published 2026-02-09

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Feng-Ting Liao, Meng-Hsi Chen, Guan-Ting Yi, Da-shan Shiu

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 a Transformer language model (the kind of AI that writes stories, answers questions, and chats with you) as a giant, multi-story factory. Inside this factory, every piece of information (a word or a token) passes through two main workstations on every floor:

The Attention Station: This is where the factory workers look at the whole sentence to understand the context. "Oh, this word refers to that word over there."
The FFN Station (Feed-Forward Network): This is where the workers do the heavy lifting of processing and refining that information.

For years, the "standard blueprint" for the FFN station has been a Narrow-Wide-Narrow shape. Think of it like a funnel that squeezes information down, then suddenly explodes it into a massive, wide room to do complex calculations, and then squeezes it back down. The industry assumed this "wide room" was necessary because it held most of the factory's workers (parameters).

The Big Idea: The Hourglass
The authors of this paper asked a simple question: Does this wide room actually need to be that wide? Or is it just a habit we got stuck in?

They proposed a new blueprint called the Hourglass FFN. Instead of one giant wide room, imagine a series of smaller, stacked "hourglass" shapes.

The Shape: It starts wide, squeezes down into a narrow bottleneck, and then expands back out.
The Stack: Instead of doing this once, they stack several of these hourglasses on top of each other, connected by "residual pathways" (like a conveyor belt that lets information skip steps if needed).

The Analogy: The Traffic Jam vs. The Express Lane
Think of the traditional "Wide" FFN as a massive, crowded highway. It has tons of lanes (parameters) to handle traffic, but it's expensive to build and maintain.

The new "Hourglass" FFN is like a smart traffic system with a few narrow tunnels.

The Trick: By forcing the traffic through a narrow tunnel and then letting it expand again, the system is forced to be more efficient. It filters out the noise and focuses on the most important signals.
The Bonus: Because the "tunnel" is narrower, you save a huge amount of construction money (parameters).

What Did They Do With the Saved Money?
This is the most exciting part. In the old design, the FFN station ate up about 75% of the budget. With the new Hourglass design, they saved a lot of that budget.

Instead of just making the factory cheaper, they reinvested the savings. They took the saved workers and moved them to the Attention Station.

The Result: The factory now has a much better "Attention" team that can understand context and relationships between words much better, while the "Processing" team is leaner but deeper (more layers).

The Findings (The Race Results)
The authors built several factories of different sizes (from small 113-million-parameter models up to 1-billion-parameter models) to test this.

Small to Medium Factories (Up to ~900M parameters): The Hourglass design won. It produced better results (lower confusion, better reasoning) than the traditional Wide design. It proved that you don't need a giant wide room to do good work; a series of efficient, stacked hourglasses works better.
The 1-Billion Parameter Factory: At this massive scale, the Hourglass design performed just as well as the traditional design. It didn't lose, but it didn't win by a huge margin either. This suggests that while the Hourglass is great, there is still a minimum amount of "processing room" needed for very large models.
The Sweet Spot: They found that the best balance wasn't just about making things wider or deeper. It was about finding a specific "U-shape" balance where the model isn't too skinny (too deep) or too fat (too wide).

The Takeaway
For a long time, AI engineers thought the "Narrow-Wide-Narrow" shape was the only way to build a good language model. This paper shows that shape is actually a bit of a habit, not a law of physics.

By switching to a Deep Hourglass shape, we can:

Make models that are just as smart but use fewer resources.
Shift the focus to making the "Attention" part of the brain stronger, which helps the model understand context better.
Prove that sometimes, going "narrower but deeper" is smarter than going "wider and shallower."

In short: The paper suggests we stop building giant, wide processing rooms and start building efficient, stacked hourglasses, allowing us to put more brainpower into understanding the conversation itself.

Technical Summary: Revisiting the Shape Convention of Transformer Language Models

Problem Statement

Dense Transformer language models have historically adhered to a rigid architectural convention: the Feed-Forward Network (FFN) utilizes a "narrow–wide–narrow" Multilayer Perceptron (MLP) shape. In this design, the model dimension ( $d_{model}$ ) is expanded to a wider intermediate dimension ( $d_h$ ) with an expansion ratio typically between 2 and 4, before projecting back. This structure dominates parameter allocation, often accounting for the majority of a model's parameters and floating-point operations.

The authors question whether this specific shape is essential or merely a convenient instantiation. Recent theoretical work suggests that residual "wide–narrow–wide" (hourglass) MLPs can serve as optimal function approximators, potentially offering superior expressivity with fewer parameters. Furthermore, empirical evidence in generative vision models and Mixture-of-Experts (MoE) Transformers has shown that hourglass structures can outperform standard designs. The core problem addressed is whether the conventional narrow–wide–narrow FFN is a necessary constraint or if a deeper, narrower hourglass FFN can offer a more efficient allocation of capacity between attention and FFN modules.

Methodology

The paper proposes a Transformer variant that replaces the conventional FFN with an Hourglass FFN. The methodology involves:

Architectural Redesign:
- The Hourglass FFN replaces the single wide expansion layer with a stack of $K$ hourglass-shaped sub-MLPs connected by residual pathways.
- Each sub-MLP ( $MLP_i$ ) follows a compression-expansion structure: it projects the input from $d_{model}$ down to a bottleneck dimension $d_h$ (where $d_h < d_{model}$ ), applies a non-linear activation (SwiGLU), and projects back to $d_{model}$ .
- The layer output is the result of passing the input through $K$ such residual sub-blocks.
Parameter Redistribution Strategy:
- By reducing the intermediate dimension $d_h$ (making the FFN "narrower"), the authors save parameters that are typically allocated to the FFN's width.
- These saved parameters are reallocated to either:
  - Increasing the model dimension ( $d_{model}$ ) and attention capacity.
  - Increasing the depth of the FFN stack ( $K$ ).
  - Increasing the number of Transformer layers ( $L$ ).
Experimental Setup:
- Baselines: Comparisons are made against standard LLaMA-style Transformers and the state-of-the-art OLMo-2 architecture.
- Scales: Experiments are conducted across model sizes ranging from 113M to 1B parameters.
- Constraints: Experiments are performed under matched total parameter budgets (within 0.001% difference) and fixed attention sizes to isolate the effects of the FFN shape.
- Search Space: The authors perform an architectural search varying $K$ (depth of hourglass stack), $L$ (number of Transformer layers), $d_{model}$ , and the ratio $d_h/d_{model}$ .

Key Contributions

The paper makes four primary contributions:

Challenging the Standard: It challenges the de facto "narrow–wide–narrow" FFN shape, proposing an efficient "wide–narrow–wide" Hourglass FFN that maintains expressivity through stacked, residual-connected sub-MLPs.
Parameter Reallocation: It demonstrates a strategy where redirecting parameters from FFN width to attention width or internal FFN depth ( $K > 1$ ) yields superior performance, particularly at smaller scales (113M).
Architectural Insights: It identifies a U-shaped width-depth trade-off, finding optimal ratios for $d_{model}/L$ and $d_h/d_{model}$ , and establishes that deeper Hourglass FFN structures are robust.
Scalability Validation: It validates the Hourglass FFN from 113M to 1B parameters, showing it is a competitive alternative to standard baselines, with specific configurations outperforming conventional models at smaller scales and matching them at larger scales.

Results

The empirical validation yields the following findings:

Performance at Fixed Attention: At the 113M scale with fixed attention parameters, Hourglass FFNs achieve lower validation loss and perplexity (PPL) than conventional FFNs. For instance, at $K=8$ , PPL dropped from 36.441 (conventional) to 36.179.
Optimal Configurations:
- Depth vs. Width: A ratio of $d_h/d_{model} \approx 0.4$ with $K=4$ layers yielded the lowest validation loss.
- Parameter Shift: At 113M, shifting parameters to make the attention module dominant (e.g., 66M attention vs. 47M FFN) resulted in the best performance (PPL 35.101), suggesting contextual processing is more valuable than static factual retrieval for small models.
- Width-Depth Trade-off: A distinct U-shaped curve was observed for the $d_{model}/L$ ratio, with an optimal sweet spot between 100 and 250.
Scalability:
- Small to Medium Scales (113M – 906M): Hourglass variants consistently outperformed conventional baselines in validation loss, perplexity, and downstream reasoning tasks (e.g., Arc Easy, HellaSwag, PIQA).
- Large Scale (1B): At the 1B parameter scale, the Hourglass variant achieved performance comparable to the strong conventional baseline (Val PPL 20.082 vs. 20.002) and surpassed the OLMo-2 configuration (20.466).
Ablation Studies: Increasing the hourglass depth $K$ consistently improved performance, though with diminishing returns beyond $K=6$ . Reducing $d_h/d_{model}$ below 0.4 led to significant performance degradation, indicating a lower bound for FFN compression.

Significance and Claims

The paper claims that the long-standing "narrow–wide–narrow" MLP shape in Transformers is not essential and may be redundant. By adopting an Hourglass FFN, the authors demonstrate that:

Efficiency: Models can achieve competitive or superior performance with a more efficient parameter distribution.
Resource Reallocation: The Hourglass structure enables a fundamental shift where the attention mechanism can become the dominant consumer of parameters, challenging the traditional FFN-dominated budget.
Re-evaluation of Design: The findings prompt a rethinking of the balance between attention and FFN, suggesting that for efficient and expressive modern language models, deeper, narrower FFNs combined with wider attention mechanisms may be a superior design strategy.

The authors remain modest regarding scalability to models with tens or hundreds of billions of parameters, noting that their findings were extrapolated from the 113M scale and that the "curse of depth" (signal propagation issues) requires further investigation in deeper architectures. They also note that the interaction between Hourglass FFNs and advanced attention mechanisms (e.g., Group-Query Attention) remains unexplored.