Variable-Width Transformers

Zhaofeng Wu; Oliver Sieberling; Shawn Tan; Rameswar Panda; Yury Polyanskiy; Yoon Kim

Source-linked AI summary

Variable-Width Transformers

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

arXiv:2606.18246v1Jun 16, 2026cs.CL

TL;DR

Transformers usually assign the same width to every layer, though layers may require different capacities. This paper introduces > <formers, which widen early and late layers while narrowing the middle, and finds they outperform uniform baselines with lower FLOPs and KV-cache costs.

Problem
It remains unclear whether transformer capacity should be distributed uniformly across depth or allocated nonuniformly under fixed depth and parameter budgets.
Method
The paper empirically trains parameter-matched decoder-only models with ×-shaped variable widths and a fixed global residual stream across layers.
Results
Parameter-matched > <formers outperform constant-width baselines across 200M–3B-parameter evaluations while reducing FLOPs and KV-cache memory.
Takeaways & Limitations
Nonuniform width allocation is a resource-efficient strategy for scaling decoder-only language models within the evaluated parameter range.
Takeaways & Limitations
Efficient training requires specialized kernels, while heterogeneous widths complicate standard tensor and pipeline parallelism.

Abstract

from arXiv · show

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.

1 Introduction

The introduction challenges the standard assumption that transformer width should remain constant across depth. It proposes empirically studying nonuniform width allocation, identifying ×-shaped models and fixed-residual implementation as promising directions with efficiency benefits.

Motivation: Uniform-width transformers allocate approximately the same computation and parameter budget to every layer, although different layers may play different computational roles.The introduction frames even allocation under fixed depth and parameter budgets as an open question.
Nonuniform allocation: The study compares growing, narrowing, 3-shaped, and ×-shaped decoder-only transformers under fixed parameter and depth constraints.The ×-shaped design is wide in early and late layers and narrower in the middle.
Residual-stream mechanism: A fixed global residual dimension lets variable-width blocks read and write layer-specific residual slices while unused coordinates bypass each block.This construction avoids projection bottlenecks and changes to the skip path caused by naïvely varying residual dimensions.
Efficiency: Nonuniform width allocation reduces training and inference FLOPs, KV cache memory, and activation I/O cost relative to constant-width transformers.The efficiency rationale is that parameter count scales quadratically with width, whereas attention FLOPs and KV cache size scale linearly.

2 Variable-Width Transformers

Variable-width Transformers allow layer-specific hidden dimensions, motivated by the distinct functions performed across transformer depth. The proposed parameter-free resizing and bottleneck-shaped width schedules preserve parameter matching while reducing average width and resource costs relative to constant-width baselines.

Motivation: Different transformer layers may require different capacities, motivating a layer-specific dimension d_ℓ instead of a constant model dimension.The paper grounds this design choice in prior evidence that layers perform distinct functions.
Resizing mechanism: The architecture resizes hidden states without parameters by truncating dimensions when shrinking and restoring previously truncated coordinates when expanding.Coordinates are recovered from the most recent preceding layer that actively processed them; otherwise, they are padded with 0.
Resizing mechanism: Variable-width layers can be viewed as reading from and writing to subsets of a wider residual stream, whose width equals the widest layer.This provides a uniform-width conceptualization while allowing individual layers to use different active dimensions.
Width schedules: A ×-shaped model narrows through early layers and widens through late layers using geometric width changes around a bottleneck layer.It uses α− < 1 for early layers and α+ > 1 for late layers; the schedule is controlled by ℓ∗, d_ℓ∗, α+, and α−.
Resource implications: Parameter-matched variable-width models have strictly fewer overall FLOPs, lower average layer width, and consequently lower KV cache size and activation I/O cost than constant-width baselines.Layer parameters scale approximately as P_ℓ ≈ Kd_ℓ^2, so matching parameter count constrains the sum of squared layer dimensions while nonconstant widths reduce their average.
Resource implications: Under parameter matching, dense linear-projection FLOPs remain identical to the baseline, while attention dot-product FLOPs scale linearly with each layer’s hidden dimension.Attention compute scales as FLOPs_ℓ ∝ N^2d_ℓ, where N is sequence length.

3 > <former

This section establishes that the ×-shaped variable-width transformer is the best-performing shape and that a cross-size bottleneck schedule consistently improves loss and downstream performance over constant-width baselines while reducing resource costs.

3 > <former: The training recipe selects a bottleneck layer index ℓ∗ and dimension dℓ∗ for pre-training variable-width models against constant-width baselines.The resulting > <formers achieve better loss and downstream task performance with smaller pre-training FLOPs and KV cache size.
3.1 Training Setup: Experiments cover dense models from 200M to 2B parameters and a 3B-total/1B-active-parameter MoE model, with parameter matching based on total parameters for the MoE comparison.The variable-width MoE model has 3% fewer active parameters than its baseline.
3.2 The × Shape Works Best: The ×-shaped model performs best among the tested variable-width shapes at the 500M-parameter scale.The comparison includes constant-width, 3-shaped, ×-shaped, ∨-shaped, and ∧-shaped models across multiple hyperparameter choices.
3.3 Finding a Specific Width Schedule: The bottleneck hyperparameters are parameterized as ratios ℓ∗ = rℓL and dℓ∗ = rdd, producing a relatively consistent performance pattern across model sizes.The ratios are swept at 200M, 500M, and 1B scales to seek a recipe that avoids per-size searches.
3.4 > <former Outperforms Constant-Width Transformer: At every tested model size, > <former outperforms the constant-width transformer while requiring fewer FLOPs and a smaller average layer size.Average layer size is used as a proxy for KV cache size during inference.
3.4 > <former Outperforms Constant-Width Transformer: 77.8% FLOPs and 85.1% average layer width are sufficient for > <former to match the 2B constant-width transformer’s loss of 2.751 under fitted scaling curves.Both loss-versus-FLOPs and loss-versus-average-layer-size curves fit tightly; > <former has a smaller intercept and slightly steeper scaling exponent.
3.4 > <former Outperforms Constant-Width Transformer: > <formers consistently outperform constant-width transformers on perplexity-based downstream tasks, while the 2B > <former wins on most natural language understanding tasks.Evaluation uses standard zero-shot LM benchmarks, including NLU and perplexity-based tasks.

4 Analysis

The analysis finds that > <former uses its nonuniform width to avoid the representational collapse seen in constant-width transformers, with denser intermediate activations, higher residual-stream entropy, and distinct prediction dynamics.

Representation strategy: > <former mitigates the middle-layer representational collapse observed in constant-width transformers and uses its capacity more effectively.This finding is based on inspecting both MLP intermediate activations and the residual stream after each layer.
MLP activations: > <former consistently activates MLP dimensions more densely and has fewer signs of under-utilized capacity than the constant-width model.Marginal utilization is measured as how often each dimension activates across tokens, thresholded at 0.1; low utilization and dead dimensions indicate under-utilized capacity.
Residual stream: The constant-width baseline develops a severe compression valley, with normalized entropy dropping near zero in middle layers, whereas > <former maintains higher entropy through bottleneck and final layers.Higher matrix entropy indicates more even use of the representation space; the result may reflect a high-entropy manifold encouraged by the constrained parameter space.
Logit lens: > <former assigns higher target-token probability and changes its decoded token distribution more gradually across early-to-middle layers, before concentrating probability on the target token in final layers.The gradual change is reflected by lower layer-to-layer KL divergence, while final-layer changes become rapid again.
Dimension expansion ablation: Copying features through the residual stream performs best among the tested methods for expanding dimensions at the 500M scale.The alternatives were padding with zeros and training a projection layer to predict extra dimensions; each method used the best result across multiple hyperparameter configurations.

5 Limitations

The approach’s main limitations are implementation complexity and training instability rather than an intrinsic algorithmic drawback. Efficient deployment would require specialized kernels and may incur overhead from the fixed-residual construction.

Implementation complexity: Efficient training requires kernels for many layer shapes, each with distinct latency, memory, and compute profiles.The fixed-residual construction may also add kernel launches through slicing, copying, and zero-padding around a wider global residual stream.
Training stability: Training a projection for the entire new layer representation is empirically unstable and diverges.The attempted projection predicted the full representation rather than only the extra dimensions.
Algorithmic versus implementation limits: The reported efficiency gap reflects infrastructure optimized for uniform width, not an intrinsic property of variable-width transformers.The authors expect purpose-built kernels to close much of the gap between theoretical and realized efficiency.
Adoption outlook: The authors do not advocate immediate adoption but encourage future architecture research to explore variable width as a design degree of freedom.This frames the architecture as a research opportunity rather than a near-term deployment recommendation.

6 Related Work

Prior work varies transformer width across depth, compresses representations across sequence length, and uses bottlenecks in non-transformer architectures. > <former is also conceptually related to Hyper-Connections, but differs by using deterministic single-stream resizing rather than learned residual-stream mixing.

Nonuniform allocation of width in transformers: Transformer variants allocate width nonuniformly across depth using block-wise, layerwise, framed, reverse, and crown scaling profiles.DeLighT narrows earlier blocks and widens later ones, while OpenELM varies attention and feed-forward dimensions across layers.
Bottleneck across sequence length: Sequence-length bottleneck models compress hidden states or activations before recovering token-level or hierarchical representations.Funnel-Transformer shortens hidden-state sequences, Hourglass Transformers downsample and upsample activations, and Perceiver models distill inputs into compact latent bottlenecks.
Bottleneck designs outside Transformers: Outside transformers, U-Net and stacked hourglass networks repeatedly reduce and recover spatial resolution, while ResNets use channel bottlenecks to reduce deep-network cost.The spatial-resolution architectures often use skip connections to preserve highresolution information.
Hyper-Connections: > <former differs from Hyper-Connections because it uses deterministic slicing and carry-forward within one global residual stream instead of learned mixing across multiple residual streams.In narrower layers, inactive coordinates bypass the block and return when width expands, avoiding learned residual-mixing matrices associated with large-scale HC instability.

7 Conclusion

The paper introduces > <former, a variable-width transformer that challenges uniform capacity allocation across depth. Across dense and MoE models from 200M to 3B parameters, parameter-matched > <formers outperform uniform baselines while reducing FLOPs and KV cache memory.

7 Conclusion: > <former challenges uniform capacity allocation across transformer depth with a variable-width architecture.The architecture is presented as the paper’s central design contribution.
7 Conclusion: 200M to 3B parameters: parameter-matched > <formers outperform uniform baselines across dense and MoE evaluations.The comparison spans both dense and mixture-of-experts models.
7 Conclusion: FLOPs and KV cache memory are both mathematically and empirically reduced by the bottleneck design.The passage reports reductions in both resource categories without specifying their magnitudes.
7 Conclusion: Analyses suggest that the bottleneck design may act as a structural regularizer.The passage attributes this interpretation to the paper’s analyses.

A Parameter-Matched Width Calculation

The width schedule is parameter-matched to a constant-width baseline by determining a geometric bottleneck shape, correcting for fixed-width embeddings, and solving for valid endpoint and bottleneck dimensions. Continuous widths are then rounded to attention-head-compatible multiples.

Width schedule: The residual stream uses layer-dependent widths with symmetric endpoints, d1 = dL = ¯d, while input and output embeddings retain baseline width d.Resizing between adjacent layers is parameter-free.
Width schedule: A geometric progression on each side of bottleneck layer ℓ∗ is fully determined by α−∈(0, 1], with α+ constrained by symmetric endpoints.The dimensionless factors cℓ(α−) define each width as dℓ = ¯d cℓ(α−).
Parameter matching: Parameter matching accounts for dominant SwiGLU block parameters, baseline embeddings, and endpoint corrections caused by padding initial embeddings and truncating final unembeddings.For dense blocks, per-layer parameters scale with Kdℓ^2, while the baseline count is Pbase = 2vd + LKd^2; the endpoint correction is Wend(¯d) = 1{¯d > d}(3 + E)¯d(¯d − d).
Parameter matching: The valid endpoint width ¯dα− is obtained by solving the parameter equation under each indicator state and selecting the self-consistent positive quadratic root.A 1D numerical solver over α−∈(0, 1] then finds the bottleneck dimension dℓ∗, after which widths are rounded to multiples of attention head dimension Q.

B Dataset Statistics

This section reports the statistics of the evaluation datasets and provides task evaluation configurations and metrics.

The evaluation dataset statistics are reported in Table 3.
Table 5 lists task evaluation configurations and metrics.

C Additional Results

The paper evaluates activation-energy concentration using Participation Ratio (PR), distinguishing raw effective dimensionality from width-normalized utilization. Compared with the constant-width baseline, > <former uses middle-layer dimensions more evenly, maintaining roughly 1,000 effective dimensions while the baseline falls below 5% normalized utilization.

Participation Ratio: Participation Ratio measures the effective number of activation dimensions by computing how evenly MLP activation energy is distributed across coordinates.The analysis reports both absolute PR and width-normalized PR to account for activation magnitude and outlier dimensions.
Participation Ratio: Normalized PR ranges from near zero when outlier dimensions dominate to 1 when energy is uniformly distributed across the layer width.Absolute PR is bounded by the layer width, while normalized PR estimates the fraction of effectively utilized dimensions.
Results: < 5%: The constant-width baseline’s width-normalized energy utilization collapses to near zero by around layer 10.This collapse indicates that most dimensions contribute little computationally under the baseline’s activation-energy distribution.
Results: roughly 1,000 effective dimensions: > <former maintains this absolute PR in middle layers, yielding more even energy utilization and a richer representation manifold.Figure 9 likewise reports higher PR for > <former in the middle layers than for the 2B constant-width transformer.