Zhaofeng Wu, Oliver Sieberling, Shawn Tan, Rameswar Panda, Yury Polyanskiy, Yoon Kim
4 min
Abstract
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.
Sam: What does that mean in plain terms?
Alex: Imagine a study group where, over time, everyone just starts copying the same answer. At first, each person had a different perspective—one noticed the grammar, another caught the logic, another spotted the tone. But gradually they all converge on one identical response, and the group loses all its diversity of thought. In an AI model, something similar can happen. The internal layers are supposed to each learn something slightly different about the data. But sometimes they drift toward producing nearly identical outputs, which means the model is wasting capacity—it has many workers all doing the same job.
Sam: And the bottleneck in the middle prevents that? Because the layers are forced to be more selective about what they process?
Alex: That's what the paper suggests. When a layer has limited capacity, it can't afford to be redundant. It has to contribute something distinct. So the architectural constraint—the narrowing itself—acts as a kind of discipline that keeps the layers doing genuinely different work.
Sam: That's a counterintuitive result. You'd think constraining the model would hurt it, but the constraint is actually part of what makes it work well.
Alex: It is counterintuitive, and the researchers are careful not to overclaim. This is a proposed architecture with supporting evidence, not a settled conclusion. But the underlying logic is coherent: the shape of a model isn't just an engineering detail. It affects what the model learns and how efficiently it learns it.
Sam: So the headline isn't just "smaller is cheaper." It's more that the distribution of capacity across a model's layers matters as much as the total amount of capacity.
Alex: That's the core argument. We tend to think about AI models in terms of their total size—how many parameters they have. This paper suggests we should also be asking how that capacity is arranged. An X-shaped model, wider where the work is hardest and narrower where it isn't, may be a more principled way to build these systems than simply stacking identical layers from start to finish. Thanks for listening to ResearchPod.