Share your thoughts, 1 month free Claude Pro on usSee more
WorkDL logo mark

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.

Zhaofeng Wu, Oliver Sieberling, Shawn Tan, Rameswar Panda, Yury Polyanskiy, Yoon Kim• 2026

Related benchmarks

TaskDatasetResultRank
Commonsense ReasoningWinoGrande
Accuracy60.2
1581
Commonsense ReasoningHellaSwag
HellaSwag Accuracy58.7
897
Question AnsweringARC Challenge
Accuracy (ARC)34.4
631
Commonsense ReasoningPIQA
Accuracy75.2
400
Common Sense ReasoningCOPA
Accuracy80
288
Reading ComprehensionBoolQ
Accuracy (BoolQ)60.9
258
Question AnsweringARC Easy
Accuracy63.3
246
Language ModelingWikiText-2 (val)
Perplexity (BVS)15.98
179
Language ModelingLAMBADA
Accuracy56.2
50
Reading ComprehensionRACE
Accuracy (RACE)34.3
26
Showing 10 of 10 rows

Other info

GitHub

Follow for update