Variable-Width Transformers

2026-06-16

AI Summary

This work addresses the inefficiency of uniform layer width in conventional Transformers, which fails to allocate computational resources according to the distinct roles of different layers. The authors propose a deformable-width Transformer architecture that employs a parameter-free residual rescaling mechanism to enable non-uniform width allocation—maintaining wider representations in early and late layers while narrowing intermediate ones. This approach provides the first systematic validation of non-uniform width effectiveness in language models, demonstrating consistent improvements over uniform-width baselines across both dense and Mixture-of-Experts (MoE) decoder-only models at scales ranging from 200M to 3B parameters. The method reduces FLOPs by 22% and decreases KV cache and I/O overhead by 15%, substantially enhancing both computational and memory efficiency.

Abstract

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.

Problem

• Transformer
• model scaling
• nonuniform width
• language models
• parameter allocation

Innovation

• Variable-Width Transformers
• nonuniform width allocation
• parameter-free residual resizing
• resource-optimal scaling
• bottleneck architecture