Variable-Width Transformers

Abstract: Scaling model size, specifically depth and width, has driven significant progress in transformer-based LLMs. 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 LLMs 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 LLMs.

• The paper introduces a nonuniform, depth-dependent width allocation that expands early and late layers while contracting middle layers to boost efficiency.
• Empirical results show a 3% improvement in perplexity and a 2.5–4.6% FLOPs reduction across models ranging from 200M to 3B parameters.
• The fixed global residual stream and parameter-free resizing mechanism enhance training stability and lower memory overhead.

Variable-Width Transformers: Revisiting Width Allocation in Transformer Architectures

Introduction

The conventional transformer architecture maintains a fixed hidden size (width) across all layers, implicitly assuming that uniform allocation of capacity is optimal under a fixed parameter or computational budget. This work, "Variable-Width Transformers" (2606.18246), systematically interrogates this design choice and empirically demonstrates that nonuniform, depth-dependent allocation of width yields both improved model quality and computational efficiency.

The central contribution is the introduction of the > <former (x-shaped transformer), which features expanded width at the boundaries—early and late layers—and a constricted bottleneck in the middle layers. The > <former employs a parameter-free residual resizing mechanism for efficient inter-layer communication. This design is empirically validated across a set of dense and mixture-of-experts (MoE) decoder-only LLMs ranging from 200M to 3B parameters.

Architecture and Methodology

Rather than enforcing a constant hidden dimension $d$ throughout the $L$ layers, the > <former sets layer-specific widths $d_\ell$, defining a geometric schedule that expands and contracts the network’s capacity. Critical to this architecture is the fixed global residual stream: each transformer block reads from and writes to a layer-specific slice of a residual stream with width equal to the maximum used by any layer. This enables seamless bypass of inactive coordinates, mitigating bottleneck-induced representational collapse and avoiding explicit projections that introduce additional parameters or instability.

The > <former’s layer widths are governed by a set of hyperparameters:

• Bottleneck location $l^*$ (as a fraction of $L$)
• Bottleneck width $d^*$
• Geometric rates for contraction (early layers) and expansion (late layers).

Layer widths are computed to maintain parameter parity with a matched constant-width baseline, ensuring fair comparison on a fixed resource budget.

Empirical Results

Training and Evaluation

Models are pretrained on DCLM, spanning 200M to 2B parameters (dense) and up to 3B (MoE), with each > <former directly compared to parameter-matched constant-width counterparts. Evaluation encompasses both language modeling loss and standard downstream LM benchmarks via the lm-evaluation-harness suite, focusing on both accuracy and perplexity.

Key empirical findings:

• The x-shaped width profile (wide-shallow–narrow-mid–wide-deep) consistently outperforms other nonuniform shapes (such as V-shaped or A-shaped) as well as the constant-width baseline.
• Across the tested spectrum, > <formers yield ≈3% relative improvement in perplexity and 2.5–4.6% reduction in pre-training FLOPs, with average layer width (and thus KV cache size) reduced by ≈11%.
• Downstream evaluation reveals > <former superiority across perplexity-based tasks and most natural language understanding benchmarks at equal or lower resource consumption.
• Scaling law analysis indicates that > <formers have both lower intercepts and steeper exponents on loss vs. FLOPs and loss vs. layer-width regression plots, suggesting greater efficiency at increasing scale.

Representation Analysis

The work provides a detailed mechanistic analysis:

• MLP Utilization: > <former yields higher and more evenly distributed MLP activation density across layers, especially mitigating underutilization and dead dimensions in the bottleneck.
• Matrix Entropy: The variable-width bottleneck structure raises the effective residual stream entropy in mid-to-late layers, counteracting representational “compression valleys” that are typical in constant-width transformers.
• Logit Lens: Intermediate layer activations decoded through the logit lens demonstrate that > <formers sustain higher target-token log probabilities and maintain less abrupt changes in the decoded distribution, suggesting more stable and informative intermediary representations.

Theoretical and Practical Implications

Theoretical Implications

This study systematically challenges the standard assumption of uniform capacity allocation in transformers, demonstrating that the fixed-width paradigm is suboptimal from both an efficiency and an expressivity perspective. The > <former’s bottleneck acts as an architectural regularizer, enforcing more efficient use of representational capacity and fostering robust network dynamics. The findings align with and extend recent literature on parameter redistribution and architectural bottlenecks, showing that full-block width variation (not merely MLP or attention specialization) can meaningfully impact LLM performance and scaling trends.

Practical Implications

Practically, the adoption of > <former-style width schedules enables training of more capable models under the same or even reduced computational budgets, while incurring lower KV cache memory and I/O costs. This is particularly pertinent given ever-increasing context lengths and scaling ambitions in current LLM research. While the method introduces some engineering complexity, primarily due to per-layer width heterogeneity and associated kernel optimizations, these are implementation concerns amenable to purpose-built kernel fusion and architectural innovation.

The insights from > <former design—both in shaping layerwise capacity profiles and in residual-stream management—provide actionable guidance for hardware-efficient LLM design, particularly in settings constrained by FLOPs, memory bandwidth, or inference latency.

Limitations and Future Directions

The primary limitations are engineering: variable-width architectures challenge current deep learning infrastructure tuned for uniform-width regimes, primarily impacting kernel efficiency and parallelism strategies. These are not inherent to the core algorithm but require tooling investment. Furthermore, while the work explores width schedules extensively, questions remain regarding optimality for very large models ($>$20B parameters), integration with advanced context-extension methods, and transfer to encoder-decoder or multimodal transformers.

Future directions include automated discovery of optimal layerwise width allocation via architecture search, extending the method to settings with sequence length bottlenecks or hierarchical design, and further investigation of the relationships between representational bottlenecks, generalization, and mechanistic interpretability.

Conclusion

"Variable-Width Transformers" decisively demonstrates that nonuniform allocation of hidden dimensions across depth—specifically, the x-shaped > <former profile—confers strong empirical and theoretical advantages for decoder-only LMs. The findings motivate a reevaluation of fixed-width architectural dogma in transformer research, establishing variable-width design as a compelling axis for efficient scaling and architectural regularization. This work is likely to inspire further research into bespoke transformer geometries optimized for emerging computational and task-specific requirements.