Wisdom in an Hourglass: When Transformers Learn to Slim Their Waist

*A deep-dive on Variable-Width Transformers. Paper: Wu et al., "Variable-Width Transformers", arXiv:2606.18246, by researchers from MIT and IBM (Zhaofeng Wu, Oliver Sieberling, Shawn Tan, Rameswar Panda, Yury Polyanskiy, Yoon Kim).*

Why Must All Layers Be the Same Width?

Since the original 2017 Transformer, every layer has used the same hidden dimension — a choice made for simplicity, not because uniform width was ever proven optimal. Scaling-law research (Kaplan et al., 2020) studied "how big," but skipped "how to allocate." The authors' core hypothesis: different layers play different roles, so a fixed parameter budget need not be spread uniformly across depth. Early layers handle noisy, high-dimensional raw input; late layers need expressive capacity for output; middle layers may just pass along compressed information.

Design: Parameter-Free Residual Resizing

The hard problem is the residual stream: if layer 5 is 1000-dimensional and layer 6 is only 300, how does information cross? The '> <former' solution treats the residual stream as a fixed-width conveyor belt equal to the widest layer:

• Narrowing: extra dimensions are simply left untouched and carried forward unchanged.
• Widening: new dimensions are copied from the most recent layer that processed them ("carry-forward").

Ablations (Table 4) show carry-forward is best: loss 3.099 vs 3.124 for zero padding and 3.150 for learned projections (which even underperform the baseline). Equivalently, the model can be seen as a uniform-width network where each layer reads/writes only a subset of the residual stream — easy to implement with standard infrastructure.

The Math: Why an Hourglass Saves Compute

For a baseline of L layers of width d, parameters scale as Ld². Keeping the average of squared widths equal to d² and applying Jensen's inequality to the strictly convex square function:

\[\left(\frac{1}{L}\sum_{\ell=1}^L d_\ell\right)^2 < \frac{1}{L}\sum_{\ell=1}^L d_\ell^2 = d^2\]

so the mean width is strictly less than d. Since attention cost scales with width, FLOPs drop — about 22% in the paper's configuration — with no heuristics, just a strict mathematical guarantee.

Experiments

Shape comparison (500M scale, Fig. 2)

| Shape | Description | Result | |:---|:---|:---| | ∨ | narrow-to-wide | poor | | ∧ | wide-to-narrow | poor | | 3 | narrow-wide-narrow | mid | | × | wide-narrow-wide (hourglass) | best |

Counterintuitively, wider middle layers (the authors' initial hypothesis) hurt — they cause representation dilution. Narrow middles act as a focusing lens.

Performance (Table 2)

| Scale | Baseline loss | '> <former' loss | FLOPs saved | Avg. width saved | |:---|:---|:---|:---|:---| | 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 wins on both loss and compute. Fitted scaling laws show a lower intercept (3.12 vs 3.16) and a slightly steeper exponent (-0.0503 vs -0.0498), suggesting the advantage may grow with scale. At matched loss, a 2B-quality '> <former' needs only 77.8% of FLOPs and 85.1% of average width. Grid search across 200M–1B consistently finds the optimum at bottleneck position r_ℓ = 0.75 (75% depth) and width r_d = 0.3 (30% of baseline) — off-center and aggressive, yet consistent across scales. Gains are larger with MoE, suggesting complementary sparsity along width (MoE) and depth ('> <former'). The ~10% narrower average width also cuts KV cache size and memory bandwidth by 10% — significant for long-context inference.

Why Slimming Helps: Internal Analyses

• Activation utilization (Figs. 5–6): baseline MLPs contain near-dead dimensions; the bottleneck acts as a structural regularizer, packing representations more densely.
• Residual-stream entropy (Fig. 7): the baseline shows a mid-layer "compression valley" (normalized entropy near 0); '> <former' keeps entropy high through the bottleneck and beyond, avoiding representation collapse.
• Participation ratio (Fig. 9): baseline middle layers have <5% effective dimensions; '> <former' sustains ~1000 effective dimensions despite narrower layers.
• Logit Lens (Fig. 8): predictions form earlier and evolve more smoothly, with rapid convergence in final layers.

Engineering Reality

Open challenges: heterogeneous layer widths need custom GPU kernels, break assumptions in tensor/pipeline parallelism, and underuse hardware optimized for uniform shapes. The authors stress these are infrastructure limits, not algorithmic ones — and even with current frameworks, FLOPs and KV-cache savings are already real.

Takeaway

Like the entasis of the Parthenon's columns, a subtle taper beats uniformity. Questioning the "all layers equal width" default reveals a long-forgotten design knob: width allocation across depth. Set it to "hourglass," and models become 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.