💬 NLP

Tapered Language Models

This paper introduces Tapered Language Models (TLMs), an architecture-agnostic design principle that improves perplexity and downstream performance by monotonically tapering MLP width from earlier to later layers under a fixed parameter budget, demonstrating that non-uniform capacity allocation outperforms traditional uniform-width stacks.

Original authors: Reza Bayat, Ali Behrouz, Aaron Courville

Published 2026-06-23

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Reza Bayat, Ali Behrouz, Aaron Courville

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine a modern language model (like the AI behind this chat) as a long assembly line in a factory. This factory has many stations (layers) stacked on top of each other. At every single station, the workers are given the exact same amount of tools, the same amount of space, and the same amount of time to do their job. This has been the standard rule since the very first "Transformer" models were built.

The paper "Tapered Language Models" asks a simple question: What if the workers at the beginning of the line actually need more tools than the workers at the end?

The Discovery: Not All Stations Are Equal

The researchers noticed that in these AI factories, the later stations often just polish or refine what the earlier stations have already built, rather than creating something entirely new. It's like a team of editors: the first few editors do the heavy lifting of finding the story and fixing the plot, while the last few editors just check for typos.

If you give the "typos-checking" team the same massive workspace and budget as the "story-creating" team, you are wasting resources. The paper proves that if you take that extra space from the end and give it to the beginning, the whole factory works better.

The Experiment: Rearranging the Tools

To test this, the researchers took a standard AI model and did a "budget swap." They kept the total number of tools and the total cost exactly the same, but they moved them around:

The "Wider-Early" Setup: They gave the first few layers (the start of the assembly line) bigger, more powerful tools and made the later layers smaller.
- Result: The AI got significantly smarter. It made fewer mistakes (lower "perplexity") and understood language better.
The "Wider-Late" Setup: They did the opposite, giving the big tools to the end of the line and small tools to the start.
- Result: The AI got much worse. It struggled to understand the basics because the early layers were too weak to build a solid foundation.

The Solution: The "Tapered" Design

Instead of just making a sharp jump between big and small tools, the authors propose a smooth transition called Tapered Language Models (TLMs).

Think of it like a funnel or a pyramid:

The top (the beginning of the model) is wide and powerful.
As you go down the stack, the width gradually and smoothly shrinks.
The bottom (the end of the model) is narrow and efficient.

They tested three ways to shape this funnel:

Linear: A straight ramp down.
Sigmoid: A steep drop in the middle, with flat tops and bottoms.
Cosine: A smooth, gentle curve that starts wide, slows down in the middle, and tapers off gently at the end.

The Winner: The Cosine curve (the smooth, gentle funnel) worked the best across the board.

Why Does This Work?

The paper digs into why this happens. They looked at what the AI is actually "thinking" at each layer.

Early Layers: These are doing the heavy lifting, creating new features and understanding the core meaning. They need a lot of "brainpower" (parameters).
Late Layers: These are mostly just repeating or reinforcing what was already found. They don't need as much new capacity; they just need to polish the work.

By shrinking the later layers, the model stops wasting energy on redundant work and focuses its power where it's actually needed: at the start.

The Results

The researchers tested this idea on four different types of AI architectures (not just the standard kind) and three different sizes (from small to large). In every single case:

The "Tapered" models performed better than the standard "Uniform" models.
They did this without adding any extra cost, extra tools, or extra computing power. It was a "free upgrade" just by rearranging the existing resources.

The Bottom Line

For years, AI designers have assumed that every layer in a neural network should be identical. This paper shows that's not true. By treating the AI like a funnel—giving it more capacity at the start and less at the end—we can make it smarter, faster, and more efficient, all without spending a penny more. It's a simple design tweak that was hiding in plain sight.

Technical Summary: Tapered Language Models

Problem Statement
Modern language models (LMs), including Transformers, recurrent variants, and memory-based architectures, share a common structural default: a stack of identical layers where parameters are allocated uniformly across depth. Despite growing evidence suggesting that layers contribute non-uniformly to the final output—where early layers often capture shallow syntactic patterns and later layers refine residual streams rather than transforming them—this uniform allocation remains the standard. The paper posits that this "one-size-fits-all" approach fails to reflect the asymmetry of layer importance, potentially wasting parameter capacity in later layers where it is less utilized.

Methodology
The authors introduce Tapered Language Models (TLMs), an architectural principle where a parameter-bearing component is monotonically tapered across depth under a fixed total budget.

Target Component: The study instantiates this principle on the MLP intermediate width ( $d_{ff}$ ). MLPs are chosen because they dominate the parameter count in modern LMs and expose width as a single, clean axis of variation independent of the token-mixing module.
Tapering Schedules: Instead of discrete block-wise allocations, the authors propose smooth, continuous decay schedules for $d_{ff}(l)$ $d_{f f} (l)$ across layers $l \in \{0, \dots, L-1\}$ $l \in {0, \dots, L - 1}$ . Three schedules are evaluated:
1. Linear: Constant rate of decay.
2. Cosine: Follows a half-cosine curve, decaying slowly at endpoints and steeply in the middle.
3. Sigmoid: Concentrates the transition in a narrow band around the midpoint.
Constraints: The tapering is strictly budget-preserving. The average per-layer width is constrained to equal the baseline uniform width ( $\frac{1}{L}\sum d_{ff}(l) = d_{ff}^{baseline}$ ), ensuring that total parameter count and FLOPs (both training and inference) remain identical to the uniform baseline.
Experimental Setup: The principle is evaluated across four distinct architectures (Standard Transformer, Gated Attention, Hope-attention, and Titans) and three model scales (440M, 760M, and 1.3B parameters). The study uses a fixed training setup (Llama 3 tokenizer, 4K sequence length) and compares tapered models against uniform baselines on in-distribution validation perplexity, out-of-distribution perplexity (WikiText, LAMBADA), and eight commonsense reasoning benchmarks.

Key Results

Direction of Allocation Matters: A motivating experiment on a 440M Transformer partitioned layers into early, middle, and late groups. Allocating more capacity to early layers ("Wider-early") improved perplexity by 0.32 points over the uniform baseline. Conversely, allocating more to late layers ("Wider-late") degraded performance by over 1.0 point, demonstrating that capacity direction is critical.
Optimal Schedule and Ratio: A sweep of width ratios and schedules on the 440M model revealed that the Cosine schedule consistently outperformed Linear and Sigmoid schedules. The optimal configuration was found at a start-to-end width ratio of 1.50 $\to$ 0.50, achieving a validation perplexity of 14.44 (a 1.84 point improvement over the uniform baseline of 16.28).
Generalization: Applying the fixed optimal configuration (Cosine, 1.5/0.5) to the 760M and 1.3B models across all four architectures consistently improved perplexity and downstream commonsense reasoning accuracy. Gains were observed regardless of whether the token-mixing module used softmax attention, gated attention, or recurrent/memory-based mechanisms.
Mechanistic Insight: Analysis of pretrained transformers revealed that MLP outputs become progressively more aligned with the existing residual stream as depth increases (measured via cosine similarity). This suggests later MLPs reinforce existing content rather than computing new features. Tapering aligns the architecture with this pattern by reducing capacity where it is least used (later layers) and reallocating it to earlier layers where it is more orthogonal and useful.

Significance and Claims
The paper establishes depth-aware capacity allocation as a simple, architecture-agnostic design lever for language models. The authors claim that tapering MLP width offers a "free" performance improvement—enhancing perplexity and downstream task accuracy without increasing parameter count or compute cost.

The significance lies in challenging the inherited uniformity of the Transformer chassis. The results suggest that the principle of tapering concerns the general allocation of parameters across depth rather than properties specific to attention mechanisms. While the authors note that the specific optimal schedule (Cosine, 1.5/0.5) was derived from a 440M model and may not be the global optimum for all scales, the consistent gains across diverse architectures and sizes indicate that the principle is robust. The paper concludes that depth-aware allocation is a "hidden lever" in plain sight that can be applied to foundation models beyond language modeling, including vision and multimodal transformers.