61 lines
2.5 KiB
Markdown
61 lines
2.5 KiB
Markdown
---
|
||
title: "Tapered Language Models (Raw)"
|
||
source: https://arxiv.org/abs/2606.23670
|
||
authors: Reza Bayat, Ali Behrouz, Aaron Courville
|
||
institutions: Mila, Cornell University, Université de Montréal, CIFAR
|
||
arxiv: 2606.23670v1
|
||
category: cs.LG
|
||
date: June 22, 2026
|
||
---
|
||
|
||
# Tapered Language Models
|
||
|
||
## Abstract
|
||
|
||
Modern LMs share a common chassis: a stack of identical layers with parameters allocated uniformly across depth. Evidence suggests layers contribute non-uniformly—later layers refine rather than transform the residual stream. We ask: should parameter capacity reflect this asymmetry? Under fixed budget, allocating more capacity to earlier layers and less to later layers improves perplexity; reverse allocation hurts. We introduce Tapered Language Models (TLMs), where a parameter-bearing component is monotonically tapered across depth. MLPs are the natural site: they dominate parameter count and expose width as a clean axis. Across three scales and four architectures, tapering MLP width via cosine schedule consistently improves perplexity and downstream benchmarks at no additional cost.
|
||
|
||
## 1. Introduction
|
||
|
||
Uniform layer width is an inherited default from Vaswani et al. (2017). As models scaled, this uniformity remained unexamined.
|
||
|
||
## 2. Related Work
|
||
|
||
- Mixture-of-Experts (MoE): conditional parameter allocation
|
||
- LayerDrop, stochastic depth: layer-wise computation pruning
|
||
- Depth-wise scaling: changing number of layers
|
||
|
||
## 3. Tapered Language Models
|
||
|
||
**Principle**: Under fixed total parameter budget, monotonically decrease parameter allocation from early to late layers.
|
||
|
||
**Why MLP width**:
|
||
- MLP accounts for majority of parameters in Transformer, Gated Attention, Mamba, Titans
|
||
- Width (d_ff) is a single clean axis of variation
|
||
- Token-mixing modules vary across architectures, making uniform comparison difficult
|
||
|
||
**Cosine schedule**: w_ℓ = w_max · (cos(πℓ/(2L)))^p, where p controls steepness
|
||
|
||
## 4. Experiments
|
||
|
||
**Controlled study** (440M Transformer):
|
||
- Uniform: 16.28 perplexity
|
||
- Cosine taper (1.50→0.50 × d_ff): 14.44 perplexity
|
||
- Cosine dominates linear at every taper range
|
||
- U-shape: optimal at 1.50→0.50, extreme tapers underperform
|
||
|
||
**Architecture sweep** (440M):
|
||
- Transformer ✓
|
||
- Gated Attention ✓
|
||
- Hope-attention ✓
|
||
- Titans ✓
|
||
|
||
**Scale sweep**: 440M, 1B, 3B — consistent improvement
|
||
|
||
**Downstream**: HellaSwag, ARC-E, PIQA, WinoGrande — taper improves all
|
||
|
||
## 5. Analysis
|
||
|
||
- Early layers benefit more from MLP capacity (transformation)
|
||
- Late layers benefit less (refinement of residual stream)
|
||
- Cosine > Linear > Step-wise > Uniform
|