PaperlayerView exact revision on arXiv

Source-linked AI summary

Revisiting the Shape Convention of Transformer Language Models

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

arXiv:2602.06471v1cs.CLcs.AIcs.LG

TL;DR

Dense Transformers conventionally devote substantial capacity to narrow-wide-narrow FFNs, leaving open whether this shape is necessary. The paper replaces that FFN with deeper residual hourglass sub-MLPs and evaluates architectural and parameter-allocation choices across scales. Hourglass FFNs are reported as competitive or superior alternatives, while larger-scale extrapolation and interactions with advanced attention remain limited.

  • Problem

    The paper asks whether the standard narrow-wide-narrow FFN shape is essential or merely one residual-compatible design, given evidence supporting residual hourglass MLPs.

  • Method

    The paper replaces conventional FFNs with stacked wide-narrow-wide hourglass sub-MLPs connected by residuals and evaluates width, depth, layer count, and parameter allocation from 113M to 1B parameters.

  • Results

    Hourglass FFNs are reported as a competitive and efficient alternative to standard baselines across model scales, with deeper configurations improving validation perplexity from 40.153 at K = 1 to 32.832 at K = 8.

  • Takeaways & Limitations

    The findings support reconsidering the narrow-wide-narrow convention and reallocating parameters between FFN structure and attention within matched budgets.

  • Takeaways & Limitations

    Comprehensive parameter searches were limited to 113M, larger-scale findings had limited ablation, and the interaction with advanced attention mechanisms was not explored.

Abstract

from arXiv · show

Dense Transformer language models have largely adhered to one consistent architectural shape: each layer consists of an attention module followed by a feed-forward network (FFN) with a narrow-wide-narrow MLP, allocating most parameters to the MLP at expansion ratios between 2 and 4. Motivated by recent results that residual wide-narrow-wide (hourglass) MLPs offer superior function approximation capabilities, we revisit the long-standing MLP shape convention in Transformer, challenging the necessity of the narrow-wide-narrow design. To study this, we develop a Transformer variant that replaces the conventional FFN with a deeper hourglass-shaped FFN, comprising a stack of hourglass sub-MLPs connected by residual pathways. We posit that a deeper but lighter hourglass FFN can serve as a competitive alternative to the conventional FFN, and that parameters saved by using a lighter hourglass FFN can be more effectively utilized, such as by enlarging model hidden dimensions under fixed budgets. We confirm these through empirical validations across model scales: hourglass FFNs outperform conventional FFNs up to 400M and achieve comparable performance at larger scales to 1B parameters; hourglass FFN variants with reduced FFN and increased attention parameters show consistent improvements over conventional configurations at matched budgets. Together, these findings shed new light on recent work and prompt a rethinking of the narrow-wide-narrow MLP convention and the balance between attention and FFN towards efficient and expressive modern language models.

1. Introduction

Dense Transformer language models have standardized on narrow-wide-narrow FFNs, but this paper revisits whether that shape is necessary by testing deeper residual hourglass FFNs. The study evaluates these variants across model scales and explores reallocating parameters across FFN width, depth, attention, and model dimensions.

  • Modern dense Transformers typically use narrow-wide-narrow FFNs with expansion ratios between 2 and 4, and the FFN accounts for most parameters.
  • The paper asks whether narrow-wide-narrow is essential or merely one residual-compatible transformation, motivated by theoretical and empirical evidence favoring residual hourglass MLPs.
  • The proposed Hourglass FFN stacks wide-narrow-wide sub-MLPs with residual connections, enabling a deeper, narrower alternative to the conventional FFN.
  • The experiments compare hourglass and conventional FFN Transformers from 113M to 1B parameters while searching over model width, layer depth, and hourglass structure.
  • The paper reports that parameter reallocation from FFN width to attention width or internal depth improves performance at 113M, and that Hourglass FFNs remain competitive through 1B parameters.

2. Background and Related Works

The conventional Transformer FFN applies a positionwise residual nonlinear transformation that expands and then contracts each token representation. Related theory and empirical work motivates studying residual hourglass structures as potentially more parameter-efficient alternatives.

  • 2.1. Narrow-wide-narrow MLP in Transformer FFN: A conventional FFN applies the same positionwise nonlinear transformation independently to each token representation and adds a residual connection.The transformation uses normalization, a nonlinear activation, and expansion and projection matrices.
  • 2.1. Narrow-wide-narrow MLP in Transformer FFN: The conventional MLP expands dmodel to an intermediate width dh greater than dmodel, then projects back to dmodel.
  • Because FFNs usually dominate parameters and floating-point operations, their expansion ratio, model dimension, and layer count shape capacity allocation across width, depth, and attention.
  • Residual hourglass MLPs retain universal approximation capability with sufficiently narrow hidden layers when enough depth is available.
  • Bottleneck and separable structures in architectures such as U-Net, MobileNet, and Wide ResNets illustrate efforts to improve parameter-compute efficiency through structured depth and width.
  • Earlier Transformer scaling work favored narrow-wide-narrow MLPs, whereas later generative-vision results report that residual hourglass MLPs outperform conventional MLPs.

3. Transformer with Hourglass FFN

The Hourglass FFN replaces the conventional wide-expansion FFN with stacked narrow-bottleneck sub-MLPs and residual connections. Its independently adjustable width, depth, and model dimension enable parameter redistribution between FFN and attention under fixed budgets.

  • Hourglass FFN design: The Hourglass FFN replaces the conventional wide-expansion MLP with a compression-expansion structure whose bottleneck dimension satisfies dh < dmodel.Each sub-block uses a down-projection, nonlinear activation, and up-projection.
  • Layer architecture: Each Transformer layer combines global sequence aggregation through attention with local per-token refinement through an Hourglass FFN.Both modules use residual connections and layer normalization.
  • Hourglass FFN design: The Hourglass FFN refines representations through K stacked hourglass-shaped MLP sub-blocks connected by residual updates.The sub-block index runs from i = 0 to K − 1, and the final layer output follows the stacked refinement process.
  • Design space: The architecture independently controls FFN depth K and bottleneck width dh, while its parameter budget depends jointly on K, L, dh, and dmodel.The stated parameter count is L(Attn(dmodel) + K(3dhdmodel)).
  • Parameter allocation: Narrowing the FFN permits saved parameters to be reinvested in model dimension or attention capacity, shifting resources toward attention interactions.The paper describes this redistribution as decoupling model depth from FFN-driven parameter growth.
  • Design space: Stacking narrower sub-blocks trades extreme FFN width for deeper compositions of nonlinear transformations without expanding activation dimensionality.The paper presents this deep-and-narrow topology as enabling more complex feature refinement than shallow, wide baselines.

4. Experiments

Experiments compare Hourglass FFNs with conventional Transformers under matched parameter budgets, varying FFN depth, width, layer count, parameter allocation, and model scale. Hourglass variants generally improve smaller-scale validation performance, enable attention-heavy reallocations, and remain comparable at 1B parameters, while deeper or excessively compressed FFNs reveal efficiency trade-offs.

  • FFN Width and Depth: At dh/dmodel ≈0.4, K = 4 generally achieves lower cross-entropy loss than K = 2 and outperforms the baseline across most tested ratios.The results support deeper, narrower FFN stacks when paired with an appropriate intermediate dimension ratio.
  • Parameter Redistribution: At approximately 113M parameters, reallocating capacity toward attention yields similar performance across settings, with (K = 4, L = 6) reaching validation perplexity 35.101 and (K = 2, L = 12) reaching 35.346.The compared configurations have validation perplexity differences from 0.05 to 1.34 and validation-loss differences of at most 0.046.
  • Width-Depth Landscape: Across width-depth sweeps, validation loss has a U-shaped relationship with dmodel/L, with a balanced range of roughly 100 to 250 and K = 4 generally performing best.The K = 4 curve lies below K = 2 and K = 1 across the sweep and nearly always below the conventional baseline.
  • Scalability: Across scales, Hourglass FFNs outperform conventional baselines at 113M, 403M, and 906M, while remaining comparable at 1B parameters.At 906M, validation perplexity improves from 22.473 to 22.282; at 1B, the values are 20.082 versus 20.002.
  • Ablation Studies: Increasing K improves validation perplexity from 40.153 at K = 1 to 32.832 at K = 8, but the model size simultaneously rises from 67M to 175M.The strongest gains occur from shallow to moderate depth, making K = 4 or K = 6 practical compromises between accuracy and parameter budget.

5. Discussions and Future Work

The Hourglass FFN challenges the standard narrow-wide-narrow constraint by using narrower, deeper FFN stacks and reallocating parameters toward attention. However, larger-scale validation and interactions with advanced attention mechanisms remain open questions.

  • Hourglass FFNs achieve competitive or superior performance while reducing FFN parameters and enabling attention to receive a larger share of the budget.At the 113M scale, attention uses 66M parameters versus 47M for the FFN.
  • Performance remains robust across depth allocations and intermediate ratios dh/dmodel ≈0.4–0.6, indicating architectural flexibility.
  • Comprehensive parameter searches were limited to 113M, while extrapolation to models up to 1B received limited ablation support.
  • Scalability to models with tens or hundreds of billions of parameters remains an open empirical question.
  • The study fixes attention to MHA, leaving interactions between Hourglass FFNs and advanced attention mechanisms unexplored.Group-Query Attention and Multi-head Latent Attention are identified as alternatives that could change the attention–FFN ratio landscape.
  • Future work must examine larger scales and potential signal-propagation problems as deeper Hourglass FFNs accumulate residual noise.The paper identifies diminishing gradient updates in later layers as a possible consequence of increased depth.

Impact Statement

The work shows that alternative Hourglass FFN structures can improve Transformer parameter and computational efficiency without sacrificing performance, supporting efforts to make large language models more efficient.

  • Hourglass structures can improve parameter and computational efficiency without sacrificing performance.
  • The findings may help reduce computational resources and energy consumption required for training and using large language models.

A.1.1. EXPERIMENTAL ENVIRONMENT AND TRAINING CORPUS

Experiments use official OLMo training scripts on NVIDIA RTX 6000 Ada and B200 GPUs with a fixed random seed. Models are trained on fixed-order prefixes of the OLMo-2 1B Stage 1 pre-training corpus.

  • Experiments use official OLMo code and NVIDIA RTX 6000 Ada and B200 GPUs with fixed random seed 6198.
  • Training uses the OLMo-2 1B Stage 1 pre-training corpus with its exact data ordering preserved.
  • The 113M, 403M, 906M, and 1074M models train on the first 2.5B, 7B, 16B, and 21B tokens, respectively.

A.1.2. VALIDATION DATASETS

Validation averages cross-entropy loss and perplexity across four datasets using the official OLMo-2 1B validation configuration.

  • Validation uses Dolma Common Crawl, Dolma The Stack, M2D2, and WikiText.
  • Cross-entropy loss and perplexity are reported as average scores across the four validation datasets.

A.1.3. DOWNSTREAM EVALUATION DATASETS

Downstream evaluation uses seven benchmark datasets, with accuracy for most tasks and perplexity for TriviaQA and NaturalQS. The experiments also report validation loss and perplexity across four model sizes and four validation datasets, alongside specified training settings.

  • Seven benchmarks are used for downstream evaluation: Arc Easy, HellaSwag, PIQA, SciQ, CommonsenseQA, TriviaQA, and NaturalQS.
  • Accuracy evaluates all downstream tasks except TriviaQA and NaturalQS, which use perplexity.
  • Training uses AdamW with cross-entropy plus an auxiliary softmax loss, while all models use SwiGLU and RoPE.
  • Additional model-size hyperparameters, including learning rates, warmup tokens, and attention heads, are provided in Table 7.
  • Validation loss and perplexity are reported for 113M, 403M, 906M, and 1074M models across four validation datasets.
Loading 2602.06471v1…