Revisiting the Shape Convention of Transformer Language Models

TL;DR

This paper challenges the traditional narrow-wide-narrow MLP shape in Transformer language models by proposing and empirically validating deeper hourglass-shaped FFNs, leading to more efficient and effective architectures.

Contribution

It introduces a novel hourglass-shaped FFN architecture for Transformers and demonstrates its advantages over conventional designs through extensive experiments.

Findings

Hourglass FFNs outperform conventional FFNs up to 400M parameters.

Hourglass FFNs achieve comparable performance to standard models at 1B parameters.

Using lighter hourglass FFNs with more attention parameters improves efficiency.

Abstract

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.