Variable-Width Transformers: When Transformer Layers Learn to 'Cinch the Waist'

Forum topic · 小凯 · 2026-06-17

Summary

A deep-dive post on zhichai.net explains the Variable-Width Transformers paper (Wu et al., arXiv:2606.18246) by MIT and IBM researchers including RoPE co-author Zhaofeng Wu and Yoon Kim. Dubbed "> <former" for its hourglass shape, the architecture challenges the seven-year-old default that all Transformer layers must share the same hidden width. Its key innovation is parameter-free residual resizing: a fixed-width residual stream where narrow layers simply skip dimensions and expanding layers copy carried-forward dimensions, which ablations show beats zero padding or learned projections. By Jensen's inequality, redistributing a fixed parameter budget across an hourglass profile strictly lowers average layer width, cutting FLOPs by roughly 2.5-4.6% and KV cache size by ~10% across 200M to 3B (MoE) models while consistently achieving lower loss. Experiments found the optimal shape narrows to 30% width at 75% depth — contradicting the authors' initial intuition that middle layers should be wider, since over-wide middles cause representational collapse. Analysis shows higher activation density, greater residual-stream entropy at the bottleneck, and more effective dimensions. Remaining challenges are engineering-level (kernels, parallelism), not algorithmic.

Variable-Width Transformers: When Transformer Layers Learn to "Cinch the Waist"

*A deep-dive into the Variable-Width Transformers paper*

> Paper: Wu et al., "Variable-Width Transformers", arXiv:2606.18246

The Parthenon Analogy

Ancient Greek architects discovered that perfectly uniform columns look "fat" due to optical illusion, so they applied *entasis* — a slight mid-column taper. Two millennia later, MIT and IBM researchers did something analogous to Transformers: they made the middle layers narrower, and the resulting model was not only more elegant but smarter.

The architecture is nicknamed "> <former" — the name itself draws an hourglass: wide opening on the left, a narrow bottleneck in the middle, wide opening on the right. The author list is notable: Zhaofeng Wu (co-author of RoPE), Oliver Sieberling, Shawn Tan, Rameswar Panda (IBM Watson), Yury Polyanskiy (MIT information theory), and Yoon Kim (MIT NLP).

> Background: The Transformer underlies all modern LLMs (GPT, Claude, Gemini). It can be viewed as a multi-layer "thinking pipeline" where each layer refines the previous layer's output. Traditionally every layer has the same width (hidden dimension) — like a 24-station factory where every station is sized identically regardless of its role.

Why Must All Layers Be the Same Width?

The uniform-width convention traces back to Vaswani et al.'s 2017 "Attention Is All You Need" — chosen for simplicity, not because uniformity was ever proven optimal. Scaling-law research (Kaplan et al., 2020) studied *how big* models should be but avoided *how to allocate* capacity across depth. Earlier efforts (DeLighT, OpenELM) only scaled FFN intermediate dimensions, not the full layer width.

The authors' core hypothesis: different layers play different computational roles, and a fixed parameter budget need not be spread uniformly across depth. Early layers process noisy, high-dimensional raw input; late layers need expressive capacity to generate diverse outputs; middle layers may just be relaying and compressing information.

Design: Parameter-Free Residual Resizing

The hardest technical problem is the residual connection: if layer 5 has width 1000 and layer 6 has width 300, how does information traverse the sudden change?

"> <former" solves this with a fixed-width residual stream equal to the widest layer:

Shrinking: a narrow layer simply truncates the extra dimensions, leaving them untouched — they pass through unchanged.
Expanding: extra dimensions are copied from the most recent layer that processed them (carry-forward).

Ablations (Table 4) confirm carry-forward is best: loss 3.099 vs 3.124 for zero padding and 3.150 for a learned projection (which even underperforms the baseline and can diverge). Equivalently, the model can be seen as uniform-width where each layer reads/writes only a subset of the residual stream — making implementation straightforward with existing infrastructure.

The Math: Why an Hourglass Saves Compute

For a baseline of L layers of width d, parameters scale as Ld². Replacing it with an hourglass profile under the constraint:

\frac{1}{L}\sum_{\ell=1}^L d_\ell^2 = d^2

Jensen's inequality (squaring is strictly convex) gives:

\frac{1}{L}\sum_{\ell=1}^L d_\ell < d

The average width is strictly less than the baseline — so attention FLOPs decrease. This is a rigorous guarantee, not a heuristic.

Experiments

Shape comparison (500M scale)

| Shape | Description | Result | |:---|:---|:---| | ∨ | narrow→wide | poor | | ∧ | wide→narrow | poor | | 3 | narrow-wide-narrow | mediocre | | × | wide-narrow-wide (hourglass) | best ✓ |

Counter-intuitively, making middle layers *wider* (shape 3) hurts — over-wide middles cause "representational collapse," like diluting espresso into a swimming pool. The narrow bottleneck acts as a focusing lens. The authors candidly note their initial intuition favored shape 3, but experiments forced the × shape.

Performance (Table 2)

| Scale | Baseline loss | > <former loss | FLOPs saved | Avg. width reduction | |:---|:---|:---|:---|:---| | 200M | 3.452 | 3.430 | 3.2% | 10.0% | | 500M | 3.138 | 3.099 | 3.7% | 11.0% | | 1B | 2.926 | 2.890 | 2.6% | 10.5% | | 2B | 2.751 | 2.726 | 2.5% | 10.9% | | 3B MoE | 2.726 | 2.710 | 4.6% | 10.9% |

Every row is lower loss + less compute. Matched-loss scaling analysis: an hourglass model matching a 2B baseline needs only 77.8% of FLOPs and 85.1% average width. Fitted scaling curves show a lower intercept (3.12 vs 3.16) and slightly steeper exponent (-0.0503 vs -0.0498), suggesting the advantage may grow with scale.

Optimal bottleneck configuration

Grid searches at 200M, 500M, and 1B all converge on r_ℓ = 0.75 (bottleneck at 75% depth), r_d = 0.3 (30% width) — an aggressive shrink, consistent across scales, with the bottleneck *not* at the midpoint, leaving a longer expansion phase for output generation.

Synergy with MoE

The relative gains are larger on MoE models (4.6% FLOPs at 3B), suggesting width-sparsity (MoE) and depth-wise non-uniform allocation are complementary.

KV cache implications

A 10% reduction in average layer width directly cuts KV cache size and memory bandwidth by ~10% — significant for long-context inference, which happens on every request.

Why Narrowing Helps: Internal Analysis

Activation utilization: baseline MLPs contain "dead dimensions"; > <former packs representations more densely — the bottleneck acts as a structural regularizer.
Residual-stream entropy: baselines show a mid-layer "compression valley" (normalized entropy near 0); > <former keeps higher entropy at and after the bottleneck, avoiding collapse.
Participation ratio: baseline middle layers use <5% effective dimensions; > <former maintains ~1000 effective dimensions despite less total width.
Logit Lens: predictions form earlier and evolve more smoothly, with rapid final-layer convergence.

Engineering Caveats

Remaining challenges are engineering, not algorithmic: heterogeneous layer widths need custom GPU kernels, break assumptions in tensor/pipeline parallelism, and mismatch hardware/software stacks (Tensor Cores, FlashAttention) designed for uniform widths. The authors expect tooling to catch up if the advantage holds.

Epilogue

Like the Parthenon's entasis — which modern engineering later showed also improves structural stability — questioning a seven-year-old default ("why must layers be equal width?") revealed a forgotten design knob: width allocation across depth. Set it to "hourglass," and the model gets both cheaper and smarter.

References

1. Wu, Z., et al. (2026). *Variable-Width Transformers*. arXiv:2606.18246. 2. Kaplan, J., et al. (2020). *Scaling Laws for Neural Language Models*. arXiv:2001.08361. 3. Vaswani, A., et al. (2017). *Attention Is All You Need*. NeurIPS 2017. 4. Levine, Y., et al. (2020). *Depth-Width Trade-offs for Neural Networks*. arXiv:2001.10648. 5. Mehta, S., et al. (2021). *DeLighT: Deep and Light-weight Transformer*. ICLR 2021.