Variable-Width Transformers

Most AI models use the same amount of "brain power" for every step of their thinking process. Whether it is the first layer or the fiftieth, the hidden dimension (the vector size representing a token) remains constant. This allocates a fixed parameter and computation budget across the entire depth of the network. This assumes that every layer plays an equally important role in processing information. This is rarely the case in complex reasoning tasks.

A new paper from MIT explores whether we can break this uniformity to achieve more resource-optimal scaling. They propose an architecture called the "> <former." This model varies the width of the transformer layers throughout the network. By making the middle layers narrower and the early and late layers wider, they create an X-shaped capacity profile. The goal is to move away from the "one size fits all" approach. Instead, they aim to distribute parameters where they actually drive performance.

The inefficiency of uniform width

Current transformer architectures operate on the assumption that a constant hidden dimension $d$ is the most stable way to scale. While this makes implementation straightforward, it ignores the reality of distinct computational roles. In practice, deep, constant-width language models often suffer from "compression valleys." These are phenomena where middle layers experience a collapse in representational capacity.

The authors observe that in standard models, middle layers often compress computations into a highly degenerate, low-rank subspace. This means the model wastes the width it was allocated. This is visible when looking at the normalized matrix entropy (a metric measuring how evenly a model uses its representational space) of the residual stream. Entropy drops toward zero in the center of the model . Essentially, you pay the FLOP (floating point operation) and memory tax for a wide layer. However, the model only uses a tiny fraction of that available dimensional space.

The X-shaped architecture and residual resizing

To address this, the authors propose the "> <former." This architecture implements a non-uniform width schedule. They specifically focus on an X-shaped profile. The model starts wide, narrows to a bottleneck in the middle, and widens again toward the final layers .

This design is parameterized by the bottleneck layer index $\ell^$ and the bottleneck dimension $d_{\ell^}$.

The core engineering challenge is changing dimensions between layers without breaking the residual stream (the highway of information that skips across layers). Standard approaches might use learned projection layers. Instead, the authors implement a parameter-free resizing mechanism. They maintain a fixed global residual dimension. This dimension equals the width of the widest layer. Each block reads from and writes to a specific slice of that stream.

When the model transitions from a wide layer to a narrower one, it simply truncates the extra dimensions. When it moves from a narrow layer back to a wider one, it restores the previously truncated dimensions. It does this by "carrying forward" features from the most recent layer that actually processed them [Equation 1]. This deterministic copying ensures that expanded dimensions are populated with relevant information from earlier in the stack.

Better loss with lower hardware costs

The results of this nonuniform allocation are significant for compute budgets. The authors report that across models from 200M to 2B parameters, the "> <former" achieved a loss of 2.726. This is better than the 2.751 seen in the constant-width baseline [Table 2].

These gains come directly from improved efficiency. Because parameter counts scale quadratically with width, matching the parameter count of a uniform model forces the "> <former" to have a lower average layer width. Under fitted loss-matched scaling curves, this architecture provides a 22% reduction in training FLOPs. It also provides a 15% reduction in KV cache memory (the memory used to store past keys and values during inference) and I/O cost .

The efficiency extends to the internal mechanics of the model. Analysis of the MLP (multi-layer perceptron) activation density shows that the "> <former" utilizes its dimensions more effectively. It exhibits denser activations and better load-balancing across the activation space compared to the baseline [, Figure 6].

Implementation hurdles and scaling tensions

Despite the strong empirical results, there are clear engineering trade-offs. The primary limitation is the increased complexity of the training and inference stack. Current deep learning frameworks are heavily optimized for the uniform-width regime. Implementing heterogeneous per-layer widths requires custom kernels (low-level code for GPU operations). Developers must handle slicing, copying, and padding without incurring massive overhead from extra kernel launches.

Furthermore, this architecture creates friction with standard distributed training techniques. Most high-performance scaling recipes rely on tensor or pipeline parallelism. Both techniques assume predictable, uniform shapes across the model's dimensions. Managing a system where tensor shapes change at every layer boundary will likely require bespoke orchestration. This is necessary to avoid significant throughput degradation.

The verdict

If you are currently mid-run on a massive pre-training cluster, do not attempt to pivot to variable-width architectures immediately. The specialized kernel requirements and the tension with existing parallelism strategies make it a high-risk move for established pipelines.

However, for teams designing the next generation of scaling laws or optimizing inference-heavy workloads, this is a definitive "yes." The ability to achieve lower loss while simultaneously slashing KV cache footprint and FLOPs is a major advantage. The "> <former" proves that the bottleneck doesn't just save money. It also acts as a structural regularizer that prevents representational collapse. If you have the bandwidth to write custom kernels, this is a highly viable path for more resource-optimal scaling.

Code is reportedly available at https://github.com/ZhaofengWu/variable-width-transformers.

Figures from the paper