Variable-Width Transformers
About
Scaling model size, specifically depth and width, has driven significant progress in transformer-based language models. However, most architectures maintain a constant width across all layers, allocating a fixed parameter and computation budget evenly despite different layers potentially playing distinct computational roles. In this work, we empirically investigate nonuniform capacity allocation across network depth by proposing a $\times$-shaped > <former architecture. This design maintains wider early and late layers while narrowing the middle layers, utilizing a parameter-free residual resizing mechanism. Across decoder-only language models ranging from 200M to 2B parameters (dense) and 3B parameters (MoE), our > <former consistently outperforms parameter-matched uniform baselines on language modeling loss. By reducing the average layer width, this architecture also requires fewer overall FLOPs (22% reduction under fitted loss-matched scaling curves) and smaller KV cache memory and I/O cost (15% reduction). In analysis, we show that this bottleneck structure results in qualitatively different representations in residual streams. Overall, our results demonstrate that nonuniform width allocation can result in more resource-optimal scaling of language models.
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| Commonsense Reasoning | WinoGrande | Accuracy60.2 | 1581 | |
| Commonsense Reasoning | HellaSwag | HellaSwag Accuracy58.7 | 897 | |
| Question Answering | ARC Challenge | Accuracy (ARC)34.4 | 631 | |
| Commonsense Reasoning | PIQA | Accuracy75.2 | 400 | |
| Common Sense Reasoning | COPA | Accuracy80 | 288 | |
| Reading Comprehension | BoolQ | Accuracy (BoolQ)60.9 | 258 | |
| Question Answering | ARC Easy | Accuracy63.3 | 246 | |
| Language Modeling | WikiText-2 (val) | Perplexity (BVS)15.98 | 179 | |
| Language Modeling | LAMBADA | Accuracy56.2 | 50 | |
| Reading Comprehension | RACE | Accuracy (RACE)34.3 | 26 |