PaperlayerView exact revision on arXiv

Source-linked AI summary

Tapered Language Models

Reza Bayat, Ali Behrouz, Aaron Courville

arXiv:2606.23670v1cs.LGcs.AIcs.CL

TL;DR

Uniform layer widths may waste capacity because layers contribute unequally; this paper introduces tapered MLP widths under fixed budgets and finds consistent gains across architectures and scales without added parameters or compute.

  • Problem

    Language models uniformly allocate parameters across identical layers despite evidence that layers contribute non-uniformly, motivating depth-aware capacity allocation.

  • Method

    Tapered Language Models smoothly decrease MLP intermediate width across depth under a fixed total parameter budget.

  • Results

    Across three scales and four architectures, cosine-tapered MLP widths consistently improve perplexity and downstream benchmark performance over uniform baselines at matched parameters and FLOPs.

  • Takeaways & Limitations

    Depth-aware capacity allocation is an architecture-agnostic design axis that improves language-model performance without increasing parameter count or training compute.

  • Takeaways & Limitations

    The selected cosine schedule was tuned only on the 440M Transformer, so it is not established as optimal for every model.

Abstract

from arXiv · show

Modern language models, including transformer, recurrent, and memory-based variants, share a common chassis: a stack of identical layers in which parameters are allocated uniformly across depth. This is a default inherited from the original transformer and largely unchanged since, yet a growing body of evidence suggests that layers contribute non-uniformly to the final output, with later layers refining the residual stream rather than transforming it. We ask whether parameter capacity should reflect this asymmetry. Our controlled experiment shows that, under a fixed budget, allocating more capacity to earlier layers and less to later layers improves perplexity over a uniform-width baseline, while the reverse allocation hurts. Building on this result, we introduce Tapered Language Models (TLMs), an architectural principle in which a parameter-bearing component is monotonically tapered across depth under a fixed total budget. MLPs are the natural site for this instantiation: they dominate parameter count across all modern LM families and expose width as a single, clean axis of variation. Across three model scales and four architectures (Transformer, Gated Attention, Hope-attention, and Titans), tapering MLP width via a smooth cosine schedule consistently improves perplexity and downstream benchmark performance over uniform baselines, at no additional parameter or compute cost. These findings establish depth-aware capacity allocation as a simple, architecture-agnostic axis of language model design, a free lever hidden in plain sight.

1. Introduction

Tapered Language Models allocate more MLP capacity to earlier layers and less to later layers under a fixed budget, improving language-model performance over uniform-width designs. The principle transfers across model scales and architecture families without additional parameters or compute.

  • Tapered Language Models: Cosine tapering reaches validation perplexity 14.44 versus 16.28 for the uniform baseline, with all configurations matched in total parameters and FLOPs.Step-wise and linear schedules fall between the uniform baseline and cosine taper.
  • Front-loading capacity: Under a fixed parameter budget, wider-early allocation reaches perplexity 15.96, while wider-late raises it to 17.29 and middle-focused allocation reaches 16.61.The wider-early configuration improves over the uniform baseline by 0.32 points; wider-late increases perplexity by more than a point.
  • Tapered Language Models: TLMs monotonically taper a parameter-bearing component across depth, replacing sharp block transitions with smooth decay while preserving total parameter count and FLOPs.The MLP instantiation uses linear, cosine, and sigmoid decay schedules.
  • Tapered Language Models: MLP width is the natural TLM target because MLPs dominate parameter storage and expose the independently adjustable intermediate dimension d_ff.The MLP structure provides a single clean width axis without changing the surrounding architecture.
  • Generalization: Across 440M, 760M, and 1.3B scales and Transformer, Gated Attention, Hope-attention, and Titans architectures, tapering improves perplexity and downstream benchmark performance.The reported gains occur without additional parameter or compute cost.
  • Mechanistic motivation: MLP outputs become progressively more aligned with the residual stream at greater depths, providing a mechanistic basis for front-loading capacity.This analysis supports allocating greater capacity to earlier layers and less to later layers.

2. Tapered Language Models

Tapering replaces uniform per-layer capacity with a monotonically decreasing allocation across depth while preserving the average dimension and total budget. This work applies the principle to MLP width, using smooth schedules that transfer capacity from later to earlier layers across multiple language-model architectures.

  • Definition: Tapering varies a per-layer architectural dimension monotonically across depth while preserving the average dimension and associated total parameter budget.The framework applies to any depth-wise dimension controlling parameter count, although this work empirically studies MLP width.
  • MLP instantiation: MLP width is the intervention because MLPs dominate parameter storage and expose a single adjustable width shared across the studied architectures.This makes tapering applicable to transformers, gated attention models, and memory-based architectures without changing surrounding architectural components.
  • MLP instantiation: Tapered MLPs replace constant d_ff with d_ff(l) that decreases from d_start to d_end, where d_start > d_end.The implementation uses smooth continuous decay rather than the motivating experiment’s discrete block-wise assignment.
  • Decay schedules: The linear, cosine, and sigmoid schedules distribute the width transition differently, with cosine providing soft endpoint plateaus and a gradual midpoint transition.Linear decay spreads change uniformly, whereas sigmoid decay concentrates it near the midpoint and produces a near-binary allocation.
  • Budget preservation: Preserving the average MLP width preserves total parameters and training and inference FLOPs, although it redistributes where computation is spent across layers.Per-layer MLP parameter count and forward FLOPs are linear in d_ff(l).

3. Experiments

Experiments select cosine tapering with a 1.5/0.5 start-to-end width ratio, then evaluate it across four architectures and two model scales at matched parameter and compute budgets. Tapering improves commonsense accuracy consistently, improves LAMBADA perplexity in every comparison, and improves WikiText perplexity in most comparisons.

  • Experimental setup: The experiments vary only per-layer MLP intermediate dimension while holding total parameters, training FLOPs, and inference FLOPs fixed.Models use the Llama 3 tokenizer with a 32K vocabulary and 4K-token training sequences.
  • Schedule and width selection: Across all five width ratios, cosine achieves the lowest perplexity, linear the second-lowest, and sigmoid the worst.Cosine’s worst configuration, 1.75/0.25 at 15.49, beats linear’s best, 1.625/0.375 at 15.64.
  • Schedule and width selection: Cosine tapering reaches its minimum perplexity at width ratio 1.5/0.5, with 14.44, forming a U-shaped pattern across the five ratios.The selected cosine schedule and 1.5/0.5 ratio are carried unchanged into subsequent experiments.
  • Main results: Tapering improves average commonsense accuracy across all four architectures and both scales, without exception.The evaluated architectures are Transformer, Gated Attention, Hope-attention, and Titans at 760M and 1.3B parameters.
  • Main results: LAMBADA perplexity improves in all eight comparisons, while WikiText perplexity improves in seven.The gains persist at both 760M and 1.3B across every architecture, with perplexity moving in the same direction across configurations.

4. Layer-wise Novelty

Layer-wise measurements show that MLP outputs become increasingly aligned with the residual stream at greater depth, indicating diminishing novelty and motivating narrower later MLPs. The same depth-wise rise appears at the block level, suggesting that tapering may extend beyond MLPs to other parameter-bearing axes.

  • Mechanism: MLP outputs become progressively more aligned with the existing residual as depth increases, reinforcing content already present rather than computing new features.The paper interprets this increasing alignment as diminishing novelty and uses it to motivate tapering later MLPs.
  • Mechanism: A rising alignment trend with depth signals diminishing novelty, while ρ≈0 denotes orthogonal content and ρ>0 denotes reinforcement of an existing residual direction.The block-update quantity measures the layer’s full additive contribution, whereas the MLP-only quantity isolates the tapered component.
  • Evidence: Both measured quantities rise from r=0.27 to r=0.71 across depth, showing that redundancy increases in the MLP itself and across the layer as a whole.The pattern is reported across pretrained GPT-2 checkpoints evaluated on WikiText-2, with first and last layers omitted as boundary cases.
  • Implications: Tapering reduces hidden dimension where residual alignment is largest and reallocates saved parameters to earlier layers, while the block-level rise suggests attention heads, key-value dimensions, and recurrent state size could also be tapered.The paper leaves these non-MLP tapering axes for future work.

5. Related Work

Related work shows that language models commonly retain a uniform per-block design despite architectural differences, while prior studies explore non-uniform computation, parameter allocation, and layer importance. These directions motivate distinguishing static MLP-width tapering from other depth-wise allocation strategies.

  • Language model families: Across architectural families, language models use the same per-block design despite differing token-mixing modules.
  • Non-uniform allocations across depth: Prior work varies sequence length, token routing, layer depth, or block size across depth, whereas this work studies static MLP width.
  • Non-uniform allocations across depth: Earlier parameter-redistribution methods vary FFN multipliers, head counts, or layer activity, but reported results do not establish a clear allocation shape.Block-wise scaling varies FFN multipliers and head counts jointly; another fixed-budget study ablates layer-wise shapes without identifying a clear winner.
  • Layer importance: Layer-importance studies find non-uniform contributions across depth, with residual streams often reaching final predictions before the last layer and later layers removable with little performance loss.

6. Limitations

The tapering sweep was limited to the 440M Transformer, and the selected schedule was transferred unchanged to larger models and alternative architectures. Thus, the reported cosine profile is a robust default but is not established as universally optimal, since preferred tapering may vary with model and training properties.

  • Experimental scope: The schedule and width-ratio sweep was conducted only on the 440M Transformer.The selected cosine schedule was then transferred unchanged to larger models and alternative architectures.
  • Experimental scope: The transferred d_start/d_end = 1.5/0.5 configuration is a robust default, but its optimality for every model remains unestablished.The transfer design provides a deliberately strict test of transferability.
  • Configuration dependence: The preferred tapering profile may depend on model depth, hidden dimension, MLP parameter fraction, token-mixing architecture, or training budget.The best schedule or endpoint ratio could shift as models scale, motivating broader sweeps across model sizes and architectural families.

7. Discussion and Conclusion

The paper introduces Tapered Language Models as a fixed-budget principle that monotonically tapers parameter-bearing components across depth, with MLP tapering improving perplexity and downstream performance across diverse models. It frames depth-aware capacity allocation as a broader, low-cost design lever while identifying other taperable dimensions and modalities as open directions.

  • Contributions: Tapered Language Models monotonically taper a parameter-bearing component across depth under a fixed total budget, using a smooth cosine schedule over MLP width.This improves perplexity and downstream benchmark performance relative to uniform-width baselines across three model scales and four architectural families.
  • Limitations and Future Work: Other layer-wise dimensions—including attention heads, key-value dimensions, recurrent state size, memory slots, and expert count—are natural tapering candidates, but their benefits remain empirically untested.The paper presents these possibilities as an initial exploration of a broader architectural principle.
  • Broader Implications: Depth-aware capacity allocation is a simple, low-cost architectural lever that may extend beyond language modeling to vision transformers, diffusion transformers, and multimodal models.These model families inherit approximately uniform capacity across depth, which the paper suggests may be unnecessarily restrictive.

A. Long-Context Retrieval on Needle-in-a-Haystack

Tapered models preserve long-context retrieval performance on Needle-in-a-Haystack across multiple retrieval tasks and context lengths. They match or improve uniform counterparts, with gains concentrated in harder, lower-scoring cases.

  • Evaluation setup: The evaluation tests whether redistributing MLP capacity across depth affects long-context retrieval on Needle-in-a-Haystack.The study uses three single-needle variants and one multi-query variant at 4K, 8K, and 16K context lengths.
  • Evaluation setup: The single-needle tasks retrieve a passkey, numerical value, or UUID, while the multi-query task evaluates multiple retrieval queries.These are S-NIAH-1, S-NIAH-2, S-NIAH-3, and MQ-NIAH, respectively.
  • Results: Tapered models match or improve over uniform counterparts across the evaluation table.Improvements are concentrated in harder cells where absolute retrieval scores are lowest.
Loading 2606.23670v1…