Routing Transformer

ArXiv: 2003.05997

🎯 Pitch

Self-attention’s quadratic cost usually forces models to ignore most of the past, but Routing Transformers learn to cluster queries and keys on the fly, letting each word attend only to its nearest neighbors in context—slashing complexity to $O(n^{1.5}d)$ while beating dense Transformers on language modeling by a wide margin.

---

1. Executive Summary

This paper proposes the Routing Transformer, a sparse attention model that learns dynamic, content-based sparsity patterns by routing queries and keys through online k-means clustering—assigning each token to attend only within its cluster rather than across the full sequence—while reducing attention complexity from $O(n^2 d)$ to $O(n^{1.5} d)$. Evaluated on language modeling (Wikitext-103, enwik-8, PG-19) and image generation (CIFAR-10, ImageNet-64), the Routing Transformer combines cluster-routed attention heads with classical local attention heads, achieving state-of-the-art results including 15.8 test perplexity on Wikitext-103 (versus 18.3 from Transformer-XL), 33.2 on PG-19 (versus 33.6 from Compressive Transformer), and 3.43 bits/dim on ImageNet-64 (versus 3.44 from Sparse Transformer), while using fewer self-attention layers and no segment-level recurrence. The paper further demonstrates that local attention alone constitutes a surprisingly strong baseline across all benchmarks, establishing that content-based routing provides complementary attention patterns that capture global consistency only when local representations have already been built over several layers.

2. Context and Motivation

The Core Problem: Quadratic Complexity Makes Self-Attention Impractical for Long Sequences

The fundamental problem this paper tackles is the quadratic computational and memory cost of self-attention with respect to sequence length. In a standard Transformer, for a sequence of length $n$, the attention matrix $A$ is $n \times n$. Computing it requires $O(n^2 d)$ operations where $d$ is the hidden dimension, and storing it requires $O(n^2)$ memory. This scaling behavior creates a hard practical ceiling: sequences longer than a few thousand tokens become prohibitively expensive or simply infeasible to process.

The paper quantifies this precisely in Section 1: for every position $i \leq n$, self-attention computes weights over its entire context of length $i$, yielding a total complexity of $\sum_{i \leq n} i = n(n-1)/2$. While $d$ is a constant factor, the $n^2$ term dominates rapidly as sequence length grows.

This is not merely a theoretical inconvenience. The paper argues that long sequences are "the norm in many domains" (Section 1), listing music generation, image generation (where images are flattened into pixel sequences), speech recognition, video modeling, and document-level machine translation. In each case, the raw data naturally produces sequences of thousands to tens of thousands of tokens. For language modeling specifically, the paper highlights datasets like PG-19, where the average document contains roughly 69,000 words — far beyond what quadratic attention can handle without some form of compromise.

The asymmetry between self-attention's modeling power and its computational cost creates a tension: self-attention gives models the unique ability to directly attend to any part of the previous context at every time step, regardless of temporal distance. RNNs and CNNs, by contrast, have direct interactions only within a local neighborhood. This global receptive field is precisely what makes Transformers so effective — but it is also what makes them so expensive.

Why This Problem Matters: Scaling Laws Meet Real-World Constraints

The significance of this problem has grown in direct proportion to the ambition of sequence modeling. When Vaswani et al. (2017) introduced the Transformer, typical sequence lengths in machine translation were on the order of 50–100 tokens, and $O(n^2)$ was manageable. But the subsequent years saw Transformers applied to increasingly longer sequences: language models operating over thousands of tokens of context (Radford et al., 2018; Dai et al., 2019), image generation treating images as sequences of 12,288+ pixels (Child et al., 2019), and music generation spanning minutes of audio (Huang et al., 2018).

Each of these applications hits the quadratic wall differently. For image generation on ImageNet-64, a single image flattened in raster-scan RGB order produces 12,288 tokens. The full attention matrix for this sequence has ~151 million entries — per head, per layer, per sample. For language modeling on PG-19, the average document length of ~69,000 words would produce an attention matrix with ~4.8 billion entries. These are not hypothetical limits; they are the exact regimes the paper reports results in (Sections 5.4 and 5.5).

The practical consequence is that researchers and practitioners face a forced choice: either truncate sequences (losing long-range dependencies), use smaller models (sacrificing capacity), or develop more efficient attention mechanisms. The paper's work on sparse attention directly addresses this third path, aiming to preserve the modeling benefits of attention while reducing its computational footprint.

There is also a subtler, more conceptual motivation. Local attention approaches — attending only to a fixed window of recent tokens — impose an inductive bias that correlation between observations naturally decreases with temporal distance. This is true for many problems but not all. The paper argues (Section 4) that some linguistic phenomena require arbitrarily long-range consistency: pronouns referring to entities introduced thousands of tokens earlier, document-level topic coherence, and stylistic consistency across paragraphs. A model restricted to local windows can never directly capture these dependencies; it must rely on information propagating through intermediate layers, which is lossy and indirect. The paper's stated goal is to preserve the ability to attend over long distances while paying only for the attention entries that actually matter.

Prior Approaches: Two Families of Sparse Attention

The paper situates itself at the intersection of two distinct lines of research on efficient attention (Section 2), and its primary contribution is bridging a gap between them.

Family 1: Attention with Temporal Sparsity (Fixed Patterns)

The dominant approach to scaling attention, at the time of this paper, was to restrict each query to attending to a fixed, position-based subset of keys — typically a local sliding window or a strided pattern. These methods are often called "data-independent" because the sparsity pattern is determined solely by positions, not by the actual content of the tokens.

Local attention (Luong et al., 2015) is the simplest form: each query attends only to the $k$ most recent keys, i.e., $S_i = \{j \mid i-k \leq j < i\}$. The attention matrix becomes a banded matrix with non-zero entries only near the diagonal. Complexity drops to $O(n k d)$, which is linear in $n$ when $k$ is fixed.

Block-sparse and strided attention (Child et al., 2019) generalizes this by using different patterns for different heads. In the Sparse Transformer, half the heads perform local attention, and the other half perform strided attention where $S_i = \{j \mid i - j \equiv 0 \pmod{k}, j < i\}$ — attending to every $k$-th token in the past. The combination of local and strided patterns approximates full attention by ensuring that information can propagate through the network via a "convolution-like" receptive field expansion.

Adaptive-span attention (Sukhbaatar et al., 2019) learns the window size $k$ per attention head using a soft masking function with an $L_1$ penalty, allowing the model to allocate longer attention spans to heads that benefit from them.

The paper acknowledges several strengths of these approaches (Section 4): they are "natively sparse," meaning the attention matrix is never instantiated in dense form; the sparsity patterns are regular and contiguous in memory, making them efficient on parallel hardware like GPUs and TPUs; and they have achieved strong empirical results, including the Sparse Transformer's state-of-the-art on ImageNet-64 and competitive language modeling results.

However, the paper identifies a fundamental limitation (Section 4):

"fixing the sparsity pattern of a content based mechanism such as self-attention can limit its ability to pool in information from large contexts."

In other words, temporal sparsity imposes a prior that relevant information is nearby. When the relevant information is not nearby — a pronoun referencing a named entity from paragraphs ago, or a pixel in an image depending on a semantically related but spatially distant pixel — local attention fails. The model can eventually propagate information through intermediate layers, but this requires depth, loses fidelity, and forces every layer to serve as a relay rather than directly attending to relevant context.

Family 2: Attention with Content-Based Sparsity (Flexible Patterns)

A second line of work learns sparsity patterns that depend on the actual content of queries and keys — not just their positions. This is the approach the paper ultimately adopts, but it argues that existing implementations had a critical flaw.

Sparsemax and entmax-based attention (Martins and Kreutzer, 2017; Malaviya et al., 2018; Correia et al., 2019) replaces the softmax operator with alternatives that produce exactly sparse probability distributions — many entries are driven to exactly zero. Correia et al. (2019) applied this to every layer in a Transformer, allowing the model to learn which query-key pairs to ignore based on their content.

The key advantage is modeling flexibility: the sparsity pattern is arbitrary and data-driven, not constrained by position. If a query is semantically related to a key 5,000 positions in the past, content-based sparsity can decide to attend to it.

The key disadvantage, and the gap the paper aims to fill (Section 2):

"sparsity here cannot be leveraged to improve space and time complexity since sparsemax/entmax formulations require instantiating the full attention matrix prior to sparsification."

This is a crucial distinction. With sparsemax, you still need to compute all query-key dot products to determine which ones to keep. The $O(n^2)$ cost is paid upfront; sparsity only reduces the downstream computation (the weighted sum over values). This makes content-based sparsity as expensive as full attention in practice — it offers no computational benefit for long sequences, only a modeling regularization effect. The paper frames this as a gap: content-based sparsity has the right idea (learning what to attend to) but the wrong implementation (dense computation followed by sparsification).

A Gap Between Two Worlds

The paper's central motivation is precisely this gap (Section 2):

"Our work is motivated by bridging this gap and allows for arbitrary sparsity patterns while avoiding having to instantiate non-zero entries of attention matrices."

The goal is to combine the modeling flexibility of content-based sparsity (attend to anything, anywhere, based on learned similarity) with the computational efficiency of natively sparse attention matrices (never compute the full $n \times n$ matrix). This is the specific technical challenge the paper addresses.

To achieve this, the paper draws on an analogy from the Maximum Inner Product Search (MIPS) literature. In a dot-product attention setting, the importance of key $K_j$ to query $Q_i$ is directly proportional to $Q_i^\top K_j$. The ideal sparse attention would, for each query, select the $k$ keys that maximize this dot product — exactly the MIPS problem. If these top-$k$ keys can be found efficiently (without computing all $n$ dot products), then sparse attention can be both content-based and efficient.

This is where clustering enters as the solution: group keys and queries into clusters, and let queries attend only to keys in the same cluster. If the clustering is good (i.e., intra-cluster dot products are high), this approximates MIPS while reducing the number of dot products from $n^2$ to roughly $n \times (n/k) = n^2/k$.

Contemporaneous Work: The Reformer

The paper explicitly positions itself relative to Kitaev et al. (2020)'s Reformer, which was developed contemporaneously and uses Locality-Sensitive Hashing (LSH) — a conceptually related but technically different approach to content-based sparse attention.

In the Reformer, queries and keys are hashed into buckets using random hyperplane projections, and tokens in the same bucket attend to each other. The hashing is performed with fixed, randomly initialized hyperplanes that do not change during training.

The Routing Transformer takes a similar high-level approach (cluster, then attend within clusters) but makes a different engineering choice: learned clustering centroids via online mini-batch spherical k-means, rather than fixed random hashes. The paper argues (Section 2) that this choice is motivated by the MIPS literature, where "typically spherical k-means is known to outperform LSH for MIPS" — citing Auvolat et al. (2015). The paper claims this theoretical advantage is empirically validated:

"This is borne out in the common task of Imagenet-64 generation, where Reformer gets around 3.65 bits/dim, while the Routing Transformer gets 3.43 bits/dim."

This is a meaningful gap (0.22 bits/dim) on a major benchmark, suggesting that learned space partitioning provides better MIPS approximations than random projections. The paper doesn't present head-to-head ablations of LSH vs. k-means within its own framework, but the cross-paper comparison establishes the motivation for the design choice.

How This Paper Positions Itself

The paper frames its contribution not as inventing content-based sparsity (which Correia et al., 2019 already did) or temporal sparsity (which Child et al., 2019 and others established), but as the first approach that successfully combines the efficiency of natively sparse attention with the flexibility of content-based routing. The key innovation is the clustering mechanism: by using online k-means to partition the key-query space, the model learns which tokens should attend to each other without ever computing the full attention matrix.

Equally important to the paper's positioning is a pragmatic architectural choice: half the heads perform local attention, half perform routing attention. The paper does not attempt to replace local attention with routing. Instead, it argues (Section 6.1) that local attention provides essential local consistency and fluency, while routing attention provides global consistency:

"we hypothesize that the reason for the strong performance of the Routing Transformer is due to the fact that it combines building local representations over several layers, together with enforcing global consistency for every token."

This hybrid design reflects a nuanced understanding: long-range attention is not uniformly necessary across all layers and all heads. The ablations on CIFAR-10 (Table 1) show that using only routing heads without local attention severely degrades performance (3.400 bits/dim at 512 window vs. 2.975 for the best hybrid), and that adding routing heads only in the top layers is often sufficient. On PG-19, the best model uses routing heads only in the last 2 of 22 layers. This aligns with the finding of Rae and Razavi (2020) that long-range attention is needed primarily in later, more abstract layers — the paper cites this as a motivating empirical observation.

Finally, the paper positions sparse attention as an orthogonal approach to the segment-level recurrence used in Transformer-XL (Dai et al., 2019) and Compressive Transformer (Rae et al., 2020). Those models train on short sequences (e.g., chunks of 512 tokens for PG-19) and maintain a memory cache of previous segments to extend the effective context length. Sparse attention instead trains directly on long sequences:

"The benefit of the Transformer-XL like approach is that it is less memory consuming... Sparse attention... on the other hand is more memory expensive since it trains directly on long sequences and therefore can scale to fewer layers for the same problem. However, as we demonstrate, it is competitive with the Transformer-XL like approaches even when using fewer layers and is guaranteed to generalize to the long sequence length that it was trained on."

This distinction is important: recurrence-based approaches approximate long-context modeling by bridging short segments; sparse attention directly models long contexts. The paper's results show that the direct approach can be competitive or superior, but with the tradeoff of requiring more memory per training step (since the entire long sequence must be processed at once) and therefore scaling to shallower models.

3. Technical Approach

3.1 Reader Orientation (Approachable Technical Breakdown)

The Routing Transformer is a modified Transformer architecture where, instead of every token attending to every previous token, each token only attends to a small subset of tokens selected by a learned clustering module — specifically, tokens that end up in the same cluster based on an online k-means assignment of their key and query vectors. The core idea is to solve a long-standing efficiency problem for self-attention: standard attention is $O(n^2)$ in sequence length, making it prohibitively expensive for documents, images, or any long-form sequential data, whereas the Routing Transformer reduces this to $O(n^{1.5})$ by using content-based clustering to select attention targets before computing dot products, so the $n \times n$ attention matrix is never instantiated.

3.2 Big-Picture Architecture (Diagram in Words)

The Routing Transformer has five major components interacting during a forward pass:

1. Linear projections: The input sequence $X \in \mathbb{R}^{n \times d}$ is projected into queries $Q$, keys $K$, and values $V$ via standard learned weight matrices — exactly as in a vanilla Transformer.
2. Cluster centroid parameters: A small set of learnable centroid vectors $\mu = (\mu_1, \dots, \mu_k) \in \mathbb{R}^{k \times d}$ is maintained as model parameters, shared across all sequences in a batch. These centroids partition the embedding space into $k$ regions.
3. Online mini-batch k-means routing: Both queries and keys are first normalized to the unit ball (via LayerNorm without scale/bias). Each query and each key is assigned to its nearest centroid. Then, for each centroid, the top-$w$ closest tokens are selected (where $w = n/k$, ensuring balanced cluster sizes). This assignment determines the sparsity pattern: token $i$ attends to token $j$ only if $j$ was assigned to the same centroid cluster as $i$.
4. Sparse attention within clusters: The actual dot-product attention is computed only within each cluster — a $k \times w \times w$ operation rather than $n \times n$. The values are gathered, multiplied by attention weights (with causal masking as needed), and scattered back to their original positions.
5. Centroid update: After each batch, the centroids are updated via an exponential moving average of the queries and keys assigned to them — an online k-means update interleaved with the main training objective.

These components are embedded in a standard Transformer layer stack, but with a design choice: in most configurations, half the attention heads use this routing mechanism and half use ordinary local (sliding-window) attention, giving the model both local fluency and global consistency.

3.3 Roadmap for the Deep Dive

• First, I'll examine the standard self-attention formulation and its complexity, to establish exactly where the quadratic cost comes from and what must be changed to reduce it.
• Second, I'll explain the $S_i$ sparsity-pattern abstraction — the formal way to think about "which keys can this query attend to?" — and contrast position-based sparsity with content-based sparsity.
• Third, I'll walk through the clustering-based routing mechanism in detail: how centroids work, how queries and keys are normalized and assigned, how balanced clusters are enforced, and how within-cluster attention is computed.
• Fourth, I'll derive the complexity analysis, showing why the optimal number of clusters is $\sqrt{n}$ and how this yields $O(n^{1.5} d)$.
• Fifth, I'll cover the online k-means update rule for centroids and the special handling of causal (left-to-right) masking in language models.
• Sixth, I'll explain the full architectural configuration — how routing heads are combined with local heads, where they're placed in the layer stack, and how this varies across tasks.

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily a methodological paper whose core contribution is a new attention mechanism that combines two previously separate ideas — content-based sparse attention and natively sparse computation — by routing tokens through a learned clustering before computing attention weights, thereby avoiding instantiation of the full $n \times n$ attention matrix.

---

Standard Self-Attention and Its Quadratic Complexity

The Routing Transformer operates within the standard Transformer architecture, which the paper reviews in Section 3. Starting from this foundation is essential because every design choice in the routing mechanism is motivated by eliminating specific computational steps that dominate the cost.

Input representation. For a single attention module, the input is a sequence of $n$ vectors of dimension $d$, represented as a matrix $X \in \mathbb{R}^{n \times d}$. Each row $x_i$ is the representation of the $i$-th token at that layer.

Linear projections. Three learned weight matrices — $W_Q$, $W_K$, $W_V$, each in $\mathbb{R}^{d \times d}$ — project the input into queries, keys, and values:

$Q = XW_Q, \quad K = XW_K, \quad V = XW_V$

These are exactly the standard Transformer projections; the Routing Transformer does not modify them.

The attention matrix. The core of the computation is the $n \times n$ attention matrix $A$, computed as:

$A = \text{softmax}\left(\text{ltr}\left(QK^\top\right)\right)$

where $\text{ltr}$ is the lower-triangular operator that masks out future positions (for auto-regressive models), and the softmax is applied row-wise.

Output computation. The next-layer representation $X'$ is then:

$X'_i = \sum_{j < i} A_{ij} V_j$

for each position $i$, summing over all previous keys $j$.

Where the cost comes from. The paper is explicit about the source of the quadratic complexity:

"for every position $i \leq n$, self-attention computes weights for its whole context of length $i$, which induces a complexity of $\sum_{i \leq n} i = n(n-1)/2$."

In practical terms: computing $QK^\top$ requires multiplying an $n \times d$ matrix by a $d \times n$ matrix, which costs $O(n^2 d)$ operations. The resulting $n \times n$ matrix must be stored, requiring $O(n^2)$ memory. For a sequence of length $n = 8192$ (as used in the enwik-8 and PG-19 experiments), this means $67$ million attention weights per head per layer — and Transformer models typically have 8–16 heads across 10–36 layers.

Multi-head attention and post-processing. A standard Transformer layer uses multiple attention heads (each with their own $W_Q$, $W_K$, $W_V$), concatenates their outputs, and passes the result through a feedforward network with a residual connection:

$X' = \text{layernorm}(X' + X)$ $X'' = \text{layernorm}(\text{mlp}(X') + X')$

The Routing Transformer preserves all of these elements unchanged. The only modification is how the attention matrix $A$ is computed — specifically, making it sparse before the quadratic operations.

---

The Sparse Attention Abstraction: $S_i$ Sets

Before introducing clustering, the paper formalizes sparse attention using a clean abstraction (Section 4). Every sparse attention model can be described by defining, for each query position $i$, a set $S_i$ of key positions that query $i$ is permitted to attend to:

$X'_i = \sum_{j \in S_i} A_{ij} V_j$

The set $S_i$ defines the sparsity pattern. The crucial property is that the number of keys attended to — $|S_i|$ — is much smaller than $n$, so the summation is cheap.

Examples given in the paper:

• Full causal attention: $S_i = \{j \mid j < i\}$ for every $i$. Every prior key is attended to; $|S_i|$ grows linearly with $i$, yielding quadratic total cost.
• Local attention with window $k$: $S_i = \{j \mid i - k \leq j < i\}$. Each query attends to at most $k$ keys; total cost is $O(n k d)$ — linear in $n$.
• Strided attention (Child et al., 2019): $S_i = \{j \mid i - j \equiv 0 \pmod{k}, j < i\}$. Attends to every $k$-th past position; also $O(n \cdot n/k \cdot d)$.

The key difference between families of methods is how $S_i$ is chosen:

• Position-based sparsity: $S_i$ depends only on the indices $i$ and $j$ — e.g., local windows produce the same pattern regardless of whether the tokens are "the cat sat" or "quantum entanglement dynamics."
• Content-based sparsity: $S_i$ depends on the actual vectors $Q_i$ and $K_j$ — a query about physics can attend to a key about physics even if they're thousands of positions apart.

The Routing Transformer's $S_i$ is defined as:

$S_i = \{j \mid \mu(K_j) = \mu(Q_i), \; j < i\}$

where $\mu(\cdot)$ returns the nearest centroid to a vector. In plain language: query $i$ attends to key $j$ if and only if they are assigned to the same cluster. The clustering itself is learned from the data, making this a content-based sparsity pattern.

What makes this natively sparse. The paper emphasizes a critical implementation distinction (Section 4):

"Content-based sparse attention should however be carefully implemented if we need to avoid instantiating full attention matrices at any point in time. For instance, Correia et al. (2019) infer sparsity from data but their formulation instantiates a full attention matrix before finding its sparse counterpart."

The Routing Transformer never computes $QK^\top$ as a single $n \times n$ operation. Instead, it first partitions tokens into clusters, then computes attention only within each cluster — a $k \times w \times w$ operation where $k \times w \approx n$. This is what makes it "natively sparse": the sparsity pattern is determined before the expensive dot products, not after.

---

Clustering-Based Routing: How Tokens Find Their Attention Partners

This is the central mechanism of the paper. The paper describes it in Section 4.1 and provides complete pseudocode in Algorithm 1.

Step 1: Normalize queries and keys to the unit ball.

Before clustering, both $Q$ and $K$ are projected onto the unit sphere:

$Q \leftarrow \text{LayerNorm}(Q), \quad K \leftarrow \text{LayerNorm}(K)$

where LayerNorm is applied with the scale and bias parameters disabled. This means each query and key vector is normalized to have zero mean and unit variance along its feature dimension, which projects it onto (approximately) the unit $d$-ball. The paper explains this choice (Section 4.1):

"In practice, instead of normalizing by the $\ell_2$ norm, we use Layer Normalization (Ba et al., 2016) with the scale and bias terms disabled. This has the benefit of projecting vectors in $\mathbb{R}^d$ to the $d$-ball and prevents its entries from becoming too small."

Why normalization matters for MIPS. The normalization step is not cosmetic — it is what makes the clustering approach a valid approximation to Maximum Inner Product Search (MIPS). The key insight is a mathematical equivalence:

When vectors are unit-normalized, Euclidean distance and dot product are inversely related:

$\|Q_i - K_j\|^2 = \|Q_i\|^2 + \|K_j\|^2 - 2Q_i^\top K_j = 2 - 2(Q_i^\top K_j)$

since $\|Q_i\| = \|K_j\| = 1$. Rearranging:

$Q_i^\top K_j = 1 - \frac{1}{2}\|Q_i - K_j\|^2$

What this equation means. The dot product between a query and a key — which is exactly what determines their attention weight — is a decreasing function of their Euclidean distance. The closer they are in Euclidean space, the larger their dot product, and therefore the larger their attention weight.

Why this form is critical. Without normalization, large dot products could arise simply from vectors with large norms rather than genuine semantic similarity — a key with magnitude 100 would dominate attention regardless of its direction. With normalization, the attention weight depends purely on angular similarity (since $\|Q_i\| = \|K_j\| = 1$, the dot product equals the cosine similarity). This is the property that spherical k-means exploits: clusters group vectors by their direction, which directly corresponds to grouping by likely attention weight.

Step 2: Assign queries and keys to centroids.

The model maintains $k$ centroid vectors $\mu = (\mu_1, \dots, \mu_k) \in \mathbb{R}^{k \times d}$ as learnable parameters, shared across all sequences. For each query $Q_i$ and each key $K_j$, we compute the nearest centroid:

$\mu(Q_i) = \arg\min_{c \in \{1,\dots,k\}} \|Q_i - \mu_c\|, \quad \mu(K_j) = \arg\min_{c \in \{1,\dots,k\}} \|K_j - \mu_c\|$

In practice, this is done by computing the $k \times n$ product matrix $\mu Q^\top$ (and similarly $\mu K^\top$), then finding the maximum along the $k$-dimension — since minimizing Euclidean distance is equivalent to maximizing dot product when all vectors are unit-normalized and centroids are approximately unit-length.

Computational cost of assignment. Computing $\mu Q^\top$ costs $O(n k d)$ — $n$ tokens times $k$ centroids times $d$ dimensions. This is the first term in the overall complexity.

Step 3: Enforce balanced clusters via top-$w$ selection.

A direct nearest-neighbor assignment could produce wildly unbalanced clusters — some centroids might attract hundreds of tokens while others attract none. This would be disastrous for parallel hardware because different clusters would require different amounts of computation and memory.

The paper's solution is simple and hardware-friendly (Section 4.1):

"for every centroid $\mu_i$ we sort tokens by distance to $\mu_i$ and cluster membership is determined by this threshold (top-$k$)."

Specifically: for each of the $k$ centroids, take the $w$ closest tokens, where $w = n/k$ (rounding as needed). This guarantees exactly $w$ tokens per cluster, for a total of $k \times w = n$ tokens assigned (with possible overlaps — a token might appear in multiple clusters' top-$w$ lists).

The paper acknowledges a subtlety:

"This adds an additional $O(n \log n)$ term to the cost, however note that this is eclipsed by the dominating term of $O(n^{1.5}d)$."

The sorting cost is explicitly accounted for but dismissed as negligible compared to the attention computation that follows.

The paper also notes a limitation:

"As a downside, this assignment does not guarantee that each point belongs to a single cluster. In the future, we want to investigate using balanced variants of k-means which is not common in an online setting."

This means a token could theoretically attend to tokens from multiple clusters if it appears in multiple top-$w$ lists. The paper treats this as acceptable for the current work while flagging it for future improvement.

Step 4: Gather tokens by cluster and compute attention within clusters.

Once the top-$w$ assignments are determined, the operations are:

• Gather: For each centroid $c$, collect the $w$ queries $Q_c'$ and $w$ keys $K_c'$ (and corresponding values $V_c'$) that were assigned to it. This produces $k$ groups of size $w \times d$.
• Compute attention: Within each cluster $c$, compute the $w \times w$ attention matrix $A_c = \text{softmax}(Q_c'(K_c')^\top)$, apply causal masking if needed, and compute the output $V_c' = A_c V_c'$.
• Scatter: Write the $w$ output vectors back to their original positions in the $n \times d$ output matrix.

The within-cluster attention costs $k \times O(w^2 d) = k \times O((n/k)^2 d) = O(n^2 d / k)$ operations. This is the second term in the overall complexity.

The causal masking nuance (Section 4.1). For auto-regressive language models, attention must be causal — a query at position $i$ cannot attend to keys at positions $j > i$. The routing mechanism introduces an additional complication:

"When grouping queries and keys belonging to a certain cluster centroid $\mu$, we may get as members queries $Q_i$ for keys $K_j$ where time-step $i \leq j$. This therefore requires an additional masking strategy in addition to the lower triangular mask used for causal attention."

In plain language: even within a cluster, a query at position $i$ might be grouped with a key at position $j > i$ (which represents a future token the model shouldn't see yet). The standard causal mask handles this within the cluster's $w \times w$ attention matrix, but there's a subtle pre-condition: sharing keys and queries prevents a mismatch where queries and keys from the same position end up in different clusters.

The paper's solution (shown in Algorithm 1, line 5):

"One solution that avoids having to use an additional mask, is to simply share keys and queries. Empirically, we have found that this works at par or better than separate keys and queries together with an additional masking strategy in the causal attention setting."

This means: in the causal (language modeling) setting, instead of having separate $W_Q$ and $W_K$ projections, the model uses the same projection for both, so $Q = K$. This guarantees that every token is in the same cluster as itself, which simplifies causal masking. For non-causal settings (encoder self-attention, encoder-decoder cross-attention), the paper notes that this issue doesn't arise.

---

Complexity Analysis: Why $k = \sqrt{n}$ Gives $O(n^{1.5} d)$

The total computational cost of the routing attention mechanism has two dominant terms (Section 4.1):

Term 1: Centroid assignment. Computing which centroid is closest to each of the $n$ queries and $n$ keys costs $O(n k d)$ — we compute $n$ dot products against $k$ centroids, each in $d$ dimensions.

Term 2: Within-cluster attention. Once tokens are partitioned into $k$ clusters of size $w = n/k$ each, computing attention within each cluster costs $k \times O(w^2 d) = O(n^2 d / k)$.

The total complexity is therefore:

$O(n k d) + O(n^2 d / k)$

Finding the optimal $k$. The two terms have opposite dependencies on $k$: the assignment cost grows with $k$ (more centroids to compare against), while the attention cost shrinks with $k$ (each cluster becomes smaller). Setting the derivative with respect to $k$ to zero (or, equivalently, setting the two terms equal) gives:

$n k d = n^2 d / k \implies k^2 = n \implies k = \sqrt{n}$

The resulting complexity. Substituting $k = \sqrt{n}$:

$O(n \sqrt{n} d) + O(n^2 d / \sqrt{n}) = O(n^{1.5} d) + O(n^{1.5} d) = O(n^{1.5} d)$

This is the paper's headline complexity claim: a reduction from $O(n^2 d)$ for full attention to $O(n^{1.5} d)$ for routing attention.

What this means concretely. For a sequence of length $n = 8192$, full attention requires approximately $8192^2 = 67$ million operations (times $d$), while routing attention with $k = \sqrt{8192} \approx 90$ clusters requires approximately $8192^{1.5} = 741,000$ operations — roughly a $90\times$ reduction in the attention computation. The actual wall-clock speedup is smaller because of constant factors and the sorting overhead, but the asymptotic improvement is substantial.

Comparison to Sparse Transformer. The paper notes that Child et al. (2019) also derived $O(n^{1.5})$ complexity for their strided attention patterns, but through a different mechanism (fixed strided patterns rather than learned clustering). The Routing Transformer achieves the same asymptotic complexity with more flexible, content-dependent sparsity.

The sorting cost. The paper explicitly adds that top-$w$ selection requires sorting tokens by distance to each centroid, costing $O(n \log n)$:

"This adds an additional $O(n \log n)$ term to the cost, however note that this is eclipsed by the dominating term of $O(n^{1.5}d)$."

For $n = 8192$, the sorting term is $8192 \log(8192) \approx 8192 \times 13 = 106,000$ operations — indeed negligible compared to $n^{1.5} d$ when $d$ is in the hundreds.

---

Online Mini-Batch k-Means: How Centroids Are Learned

The clustering centroids are not fixed — they are model parameters learned jointly with the rest of the network. The paper uses an online exponential moving average update (Section 4.1):

$\mu \leftarrow \lambda \mu + \frac{1 - \lambda}{2} \sum_{i: \mu(Q_i) = \mu} Q_i + \frac{1 - \lambda}{2} \sum_{j: \mu(K_j) = \mu} K_j$

where:

• $\lambda$ is a decay parameter, typically set to $0.999$.
• The first sum averages all queries assigned to this centroid in the current batch.
• The second sum averages all keys assigned to this centroid in the current batch.
• The factor $1/2$ gives equal weight to queries and keys.

What this update does operationally. After each mini-batch, each centroid is moved slightly toward the mean of the queries and keys that were assigned to it, weighted by $1 - \lambda = 0.001$ (so it moves 0.1% of the way toward the current batch's mean). Over many batches, the centroids converge to cluster centers that partition the embedding space into regions of high query-key similarity.

Why an exponential moving average? Standard batch k-means recomputes centroids as exact means of all assigned points, which requires storing all assignments from the entire dataset. Online k-means with an EMA allows the centroids to adapt continuously during training without requiring a separate clustering phase or storing historical assignments. This is the standard approach for embedding online clustering in neural network training, following Bottou and Bengio (1995).

The decay parameter choice. $\lambda = 0.999$ means the effective memory of the EMA is approximately $1/(1-\lambda) = 1000$ batches. This is long enough to smooth out batch-to-batch noise but short enough to track the evolving embedding space as the model's parameters change during training.

Exclusion of padding tokens. The paper notes that padding tokens — used to make all sequences in a batch the same length — are excluded from centroid updates:

"Additionally, we also exclude padding tokens from affecting the centroids."

This prevents padding positions (which contain no meaningful content) from distorting the cluster centers.

---

Architectural Integration: Combining Routing Heads with Local Heads

The Routing Transformer is not a pure routing-attention model. Across all experiments except PG-19, the architecture follows a hybrid design (Section 5):

"In all our models except the one used for PG-19, we allocate half the heads to do local attention and the other half to route attention."

Why hybrid? The paper's ablations on CIFAR-10 (Table 1) provide clear evidence. A model with only routing attention heads (8 routing layers, 8 routing heads) achieves 3.190 bits/dim with a 512 attention window, dramatically worse than the best hybrid model with 4 routing layers and 4 routing heads (2.975 bits/dim). The pure-local baseline achieves 3.009 bits/dim, and the pure-full-attention baseline achieves 2.983 bits/dim.

The interpretation (Section 6.1):

"we hypothesize that the reason for the strong performance of the Routing Transformer is due to the fact that it combines building local representations over several layers, together with enforcing global consistency for every token."

In other words, local attention builds up fluent, grammatically coherent local representations; routing attention adds long-range consistency (coreference, topic coherence, style) on top. Neither alone is sufficient.

Layer placement of routing heads (Section 5.5, PG-19 experiments). On the PG-19 dataset, the best model deviates from the "half local, half routing" pattern in two ways:

1. Only 2 of 8 heads are routing heads (not 4).
2. Routing heads appear only in the last 2 of 22 layers.

The paper explains:

"This is motivated by our empirical finding that long range attention is only needed in the last few layers - see also Rae and Razavi (2020)."

This is an important architectural insight: early layers benefit most from local attention (building up low-level features from nearby context), while later layers — where representations are more abstract and semantic — benefit from the global consistency that routing attention provides. This aligns with the intuition that a pronoun-resolution head needs high-level semantic representations of both the pronoun and its antecedent, which emerge only in deeper layers.

Specifications for each experiment:

• CIFAR-10 (Table 1): 12 layers, 8 heads. Routing layers are always "added at the top of the model." The remaining layers and heads use local attention. Configurations vary from 2 to 12 routing layers and 2 to 8 routing heads.
• Wikitext-103 (Section 5.2): 10 layers, 16 heads. Half the heads are routing (8 heads), half local (8 heads). Routing heads present in every layer. $k = 16$ clusters, attention window = 256.
• enwik-8 (Section 5.3): 24 layers, 8 heads. $k = 32$ clusters, attention window = 256.
• ImageNet-64 (Section 5.4): 24 layers, 16 heads. Half routing, half local. $k = 8$ clusters, attention window = 2048, batch size = 1.
• PG-19 (Section 5.5): 22 layers, 8 heads. Only 2 routing heads, only in the last 2 layers. Remaining 6 heads use local attention. $k$ not explicitly stated for this configuration; attention window = 512.

The "attention window" parameter. In both local and routing heads, the attention is further restricted by a window size — the maximum number of tokens any query can attend to. For local heads, this is the sliding window size (e.g., 512 means attend to the most recent 512 tokens). For routing heads, this is the cluster size $w = n/k$ (e.g., with $n = 8192$ and $k = 32$, each cluster contains $w = 256$ tokens). The paper sweeps window sizes 512 and 1024 in the CIFAR-10 ablations.

Relative position encodings. All language modeling experiments (Wikitext-103, enwik-8, PG-19) use the relative position encoding scheme of Shaw et al. (2018), which modifies attention scores to depend on the relative distance between query and key positions rather than absolute positions. This is standard practice for Transformer language models and is orthogonal to the routing mechanism.

Causal attention handling in Algorithm 1. The pseudocode in Algorithm 1 shows the complete routing attention computation:

• Lines 7–8: Normalize Q and K with LayerNorm (scale, bias disabled).
• Line 9: Compute query-centroid dot products $Q_\text{prod} = \mu Q^\top$ (size $k \times n$).
• Lines 13–14: Take top-$w$ along the $n$-dimension for each centroid, then sort to preserve token order.
• Line 15: In the causal setting, key assignments use the same indices as query assignments ($K_\text{idx} = Q_\text{idx}$, line 15) — this is the key-query sharing described earlier.
• Lines 19–21: Gather queries, keys, and values according to the index tensors (producing $k \times w \times d$ tensors).
• Line 22: Compute $A = Q'(K')^\top$, producing a $k \times w \times w$ attention tensor.
• Lines 23–24: Apply lower-triangular masking if in causal mode.
• Lines 25–26: Apply softmax and weighted sum over values.
• Line 27: Scatter the $k \times w \times d$ output back to $n \times d$ using the key indices.
• Lines 28–29: Compute one-hot cluster assignments for centroid updates (not used in the forward pass).
• Line 31: Update centroids via EMA.

---

The MIPS Connection: A Deeper Look

The paper frames the clustering mechanism as solving an approximate Maximum Inner Product Search (MIPS) problem, which is worth examining in detail because it provides the theoretical justification for why the method works.

The MIPS problem in the attention context (Section 4.1). Given a large collection of keys $\mathcal{K} = \{K_1, \dots, K_n\}$ and a query $Q_i$, the MIPS problem is to find:

$K^*_j = \arg\max_{x \in \mathcal{K}} \; Q_i^\top x$

In other words: find the key with the highest dot product with the query.

Why MIPS matters for attention. In dot-product attention, the softmax function assigns high weight to keys with high dot products and near-zero weight to keys with low dot products. If we restrict each query to attending to only $w$ keys (for efficiency), the optimal choice — the one that minimizes information loss — is to select the $w$ keys with the highest dot products. This is exactly a MIPS problem.

Why clustering approximates MIPS. The paper provides a geometric argument (Equations 10–13 in Section 4.1). When queries and keys are unit-normalized:

$\|Q_i - K_j\|^2 = 2 - 2Q_i^\top K_j$

If $Q_i$ and $K_j$ are both assigned to the same centroid $\mu$, then by the triangle inequality:

$\|Q_i - K_j\| \leq \|Q_i - \mu\| + \|K_j - \mu\|$

If both distances to the centroid are small (say, less than $\varepsilon$), then $\|Q_i - K_j\| < 2\varepsilon$, which implies:

$Q_i^\top K_j > 1 - 2\varepsilon^2$

$\text{Therefore, when two time steps $i > j$ are assigned the same cluster due to a small $\|Q_i - \mu\|, \|K_j - \mu\|$ distance, it also means that their attention weight $Q_i^\top K_j$ is high, i.e., $K_j$ is an approximate solution to the MIPS objective of Equation 9 for query $Q_i$.}$

What this derivation establishes. It proves that the clustering-based routing preserves large attention weights — the pairs that would have received high attention in full attention are exactly the pairs that end up in the same cluster. The pairs that are not preserved are those with low dot products, which would have received near-zero attention anyway. In principle, the method discards only the negligible attention weights.

Why this is not perfect in practice. The derivation assumes both $Q_i$ and $K_j$ are close to their shared centroid, which is true on average for well-clustered data but not guaranteed for every individual pair. Some high-dot-product pairs might be assigned to different centroids (a false negative — missed attention), and some low-dot-product pairs might end up in the same cluster (a false positive — attending to irrelevant tokens). The quality of the clustering directly determines how close the sparse attention is to full attention.

Comparison to LSH (Reformer). Kitaev et al. (2020)'s Reformer uses Locality-Sensitive Hashing, which provides a similar guarantee probabilistically: with high probability, nearby vectors hash to the same bucket. The difference is that LSH uses random projections while the Routing Transformer uses learned projections (the centroids). The paper cites Auvolat et al. (2015) for the claim that spherical k-means typically outperforms LSH for MIPS, and points to the 3.65 vs. 3.43 bits/dim ImageNet-64 comparison as empirical evidence. The intuition is that learned centroids can adapt to the data distribution, placing cluster boundaries in low-density regions where they minimize false splits of semantically similar tokens, whereas random projections are data-agnostic.

---

Training and Optimization Details

The paper provides specific training configurations for each experiment, which illuminate how the routing mechanism integrates into the full training pipeline.

Optimizer. All experiments except PG-19 use Adam (Kingma and Ba, 2015) with:

• Learning rate: $2 \times 10^{-4}$
• $\beta_1 = 0.9$, $\beta_2 = 0.98$
• Learning rate schedule: the warmup-then-decay schedule from Vaswani et al. (2017)

For PG-19, the paper uses the Adafactor optimizer (Shazeer and Stern, 2018) instead, with:

• Learning rate constant: $0.01$
• Linear warmup over 10,000 steps
• Followed by rsqrt_normalized_decay

The switch to Adafactor is motivated by memory efficiency for training larger models on the larger PG-19 dataset.

Dropout. Varies by experiment:

• Wikitext-103: attention dropout 0.3, ReLU dropout 0.3
• enwik-8: attention dropout 0.4, ReLU dropout 0.4
• PG-19: no dropout, no weight decay

Batch sizes and sequence lengths. These reflect the memory constraints of training on long sequences:

• CIFAR-10: batch size 32, sequence length 3072
• Wikitext-103: batch size not explicitly stated, but sequence length is the full article length (up to the model's context window)
• enwik-8: sequence length 8192
• ImageNet-64: batch size 1, sequence length 12,288
• PG-19: sequence length 8192, batch size 8192 tokens (likely meaning effective batch size in tokens, not sequences)

Hardware. All models are trained on 128 TPUv3 cores. The step-time comparisons in Tables 1 and 7 are reported on TPUv3.

Training duration. CIFAR-10 models train for 200,000 steps. ImageNet-64 models train for "roughly 70 epochs" (Child et al., 2019 used a similar duration).

Centroid initialization. The paper does not explicitly describe how centroids are initialized, but the standard approach for online k-means in neural networks is random initialization (e.g., sampling from a normal distribution, or initializing from the first batch's data). Given the $O(n^{1.5}d)$ complexity, the centroids adapt quickly as training progresses.

4. Key Insights and Innovations

Innovation 1: Content-Based Sparsity Can Be Natively Sparse — Not a Post-Hoc Filter

Prior to this work, there was an unspoken assumption in the efficient attention literature that content-based sparsity and computational efficiency were mutually exclusive. If you wanted the sparsity pattern to depend on what the tokens actually said (rather than just where they appeared), you had to compute all pairwise dot products first, then zero out the ones you didn't want. Correia et al. (2019) embodied this assumption: their entmax-based adaptively sparse Transformers learned highly flexible, data-dependent sparsity patterns, but required instantiating the full $n \times n$ attention matrix before sparsifying it. The computational savings from sparsity were purely downstream (fewer operations in the weighted sum over values), but the $O(n^2 d)$ bottleneck in computing query-key dot products remained untouched.

The Routing Transformer's core conceptual move is to break this coupling. It asks: can you decide which query-key pairs to attend to without computing all query-key dot products? The answer it provides — yes, via online spherical k-means clustering — is far from obvious because clustering introduces its own cost ($O(nkd)$ for centroid assignment plus $O(n \log n)$ for sorting) and its own approximation error (some high-dot-product pairs may be assigned to different clusters, and some low-dot-product pairs may end up together). The paper's contribution is demonstrating that this approximation is not just theoretically sound (via the MIPS equivalence argument connecting Euclidean distance to dot product for unit-normalized vectors) but empirically viable — the model outperforms full attention in some configurations (e.g., CIFAR-10 at 2.971 bits/dim for the best routing model versus 2.983 for full attention, Table 1), even though it computes only a fraction of the dot products.

This is not an incremental refinement of prior content-based sparsity methods. It is a fundamental reframing of the problem: from "compute everything, then prune" to "route first, compute locally." The shift is architectural — the routing module becomes a learned gate that determines the flow of information before the expensive operations, rather than a regularization term applied after them. This is why the paper's contribution is not merely "we use clustering for attention" (which would be a direct application of known MIPS techniques), but rather "we show that a learned routing module can replace dense query-key comparison without degrading model quality, thereby achieving for content-based sparsity what local attention achieved for position-based sparsity."

The MIPS framing itself is diagnostic: it reframes the attention sparsity problem as an approximate nearest-neighbor search problem. This is a conceptual bridge between two historically separate literatures — efficient Transformers and large-scale similarity search — that opens the door for future work to import techniques from the latter (product quantization, hierarchical navigable small world graphs, etc.) into the former. The paper's specific choice of spherical k-means is almost incidental to this larger move; what matters is that the problem has been recast in a way that makes a whole toolkit of approximate search methods applicable.

The empirical anchoring for this innovation is not a single number but a pattern across experiments. The fact that routing attention matches or exceeds full attention (CIFAR-10, Table 1) while computing only $O(n^{1.5} d)$ operations demonstrates that the clustering approximation is sufficiently accurate. The fact that it outperforms random routing by a wide margin (2.971 vs. 3.076 bits/dim, Table 1) demonstrates that the learned clustering is genuinely solving the MIPS problem — routing tokens with high dot products together — rather than just providing a regularizing noise source. And the fact that it underperforms local attention when local heads are removed (3.400 bits/dim with 8 routing heads and no local heads, Table 1) demonstrates that routing alone is not a complete replacement for dense attention, but rather a complementary mechanism.

---

Innovation 2: Local and Global Attention Are Complementary Mechanisms, Not Competing Ones

When the Sparse Transformer (Child et al., 2019) introduced strided attention as an alternative to local attention, it framed the two as alternative sparsity patterns — different ways to approximate full attention, to be used in different heads of the same model. The Routing Transformer takes this observation and elevates it to a principle about representation learning: local attention and long-range attention serve fundamentally different functions in a deep network, and a well-designed architecture should provide both, not choose between them.

The paper's evidence for this claim is both quantitative and qualitative. Quantitatively, the CIFAR-10 ablations (Table 1) show that:

• Local-only: 3.009 bits/dim (strong baseline, close to full attention's 2.983)
• Routing-only: 3.400 bits/dim at 512 window, 3.291 at 1024 window (much worse than local-only)
• Hybrid (4 routing layers, 4 routing heads, rest local, 1024 window): 2.958 bits/dim (better than full attention)

The pattern is clear: routing attention alone is worse than local attention alone, but combining them is better than either. This is not what one would expect if routing were simply a better approximation of full attention that could replace local attention. Instead, it suggests that the two mechanisms are capturing different types of dependencies, and the model needs both.

The qualitative evidence comes from the Jensen-Shannon divergence analysis (Table 6). The paper computes the divergence between attention distributions of different heads on Wikitext-103 and finds that:

• JSD between local heads is consistently very low (0.0038 to 0.3071 across layers).
• JSD between local and routing heads is almost always near the theoretical upper bound of 0.6931 (0.4706 to 0.6674).
• JSD between routing heads is intermediate (0.1579 to 0.5820).

What this pattern means: Local attention heads all produce similar attention distributions — they're all looking at recent context in broadly the same way. Routing attention heads produce attention distributions that are radically different from local heads and substantially different from each other. This is direct evidence that the routing mechanism is enabling heads to specialize in attending to semantically related but spatially distant tokens — a capability that local heads fundamentally lack. The near-maximum JSD between local and routing heads means the routing attention distribution is essentially non-overlapping with the local attention distribution, confirming that routing is not just "local attention with a larger window" but a qualitatively different attention mechanism.

This insight connects to a broader architectural principle that was emerging in the field around the same time: different layers and different heads should have different receptive field characteristics. Sukhbaatar et al. (2019) learned per-head attention spans, finding that some heads naturally learn short spans and others learn long ones. Rae and Razavi (2020) found that long-range attention is only needed in the last few layers. The Routing Transformer synthesizes these observations into a concrete design principle: the model should contain dedicated "local" heads for building fluent, low-level representations, and dedicated "routing" heads for enforcing global consistency, with the routing heads concentrated in deeper layers where representations are sufficiently abstract for semantic matching to be meaningful. The PG-19 experiment operationalizes this directly: only 2 of 22 layers contain routing heads, and they appear only at the top.

The innovation here is not the hybrid architecture itself (Child et al., 2019 already used half local, half strided heads) but the diagnostic demonstration that the two mechanisms are complementary rather than substitutable, backed by both performance trends and attention distribution analysis. This finding has implications for architecture design: it suggests that future work on efficient attention should not look for a single sparsity pattern that does everything, but rather for ways to combine specialized attention mechanisms that each serve a distinct representational purpose.

---

Innovation 3: Local Attention Is a Surprisingly Strong Baseline — A Diagnostic Contribution

One of the paper's most striking findings is buried in the results tables rather than highlighted as a headline claim, but it has substantial implications for how the field evaluates efficient attention methods. Across every benchmark tested, a scaled-up version of pure local attention — Transformer with relative position encoding and a sliding window, with no recurrence, no dynamic sparsity, and no learned routing — achieves results that are competitive with or close to state-of-the-art.

The paper reports:

• Wikitext-103: Local Transformer achieves 19.8 perplexity (Table 2), compared to 18.3 for Transformer-XL and 15.8 for Routing Transformer.
• enwik-8: Local Transformer achieves 1.10 bits per byte (Table 3), compared to 0.99 for Transformer-XL and Sparse Transformer, and 0.98 for Adaptive Transformer.
• ImageNet-64: Local Transformer (scaled ImageTransformer) achieves 3.48 bits/dim (Table 4), compared to 3.44 for Sparse Transformer and 3.43 for Routing Transformer.
• PG-19: Local Transformer achieves 39.3 perplexity (Table 5), compared to 36.3 for Transformer-XL and 33.6 for Compressive Transformer.

These numbers tell a consistent story: a well-tuned local attention model gets 80–95% of the way to state-of-the-art performance on long-sequence benchmarks. The gap between local attention and full attention is often smaller than the gap between competing efficient attention methods. On ImageNet-64, the difference between local attention (3.48) and Routing Transformer (3.43) is 0.05 bits/dim — a fraction of the gap between local attention and the earlier SPN (3.52) or PixelSNAIL (3.52).

Why this is a diagnostic contribution rather than just a baseline. Prior work on efficient attention (Child et al., 2019; Sukhbaatar et al., 2019; Dai et al., 2019) had generally reported local attention as a weak baseline, motivating the need for more sophisticated sparsity patterns or recurrence mechanisms. The Routing Transformer paper shows that this characterization was largely an artifact of undertuned local attention baselines — using too few layers, too few heads, or suboptimal training configurations. When local attention is scaled up to comparable model sizes and properly tuned, it performs substantially better than previously reported.

This finding reframes the burden of proof for new efficient attention methods. If a proposed method improves over local attention by only a small margin, the improvement may be attributable to better tuning, larger model capacity, or training for more steps rather than to the mechanism itself. The paper implicitly establishes a higher bar: new methods should demonstrate substantial gains over a well-tuned local attention baseline at matched computational budget, not just over whatever local attention configuration was convenient to run.

The paper also uses this finding to contextualize its own contributions honestly. The gains from routing attention are real (15.8 vs. 19.8 on Wikitext-103 is substantial; 33.2 vs. 39.3 on PG-19 is meaningful), but the paper does not claim that local attention is a straw man. Instead, it argues (Section 6.1) that local attention's strength actually supports the hybrid design: local attention is good at what it does (local fluency and consistency), and routing attention adds complementary value (global coherence) on top. The diagnostic contribution thus serves double duty: it calibrates expectations for future work while also providing evidence for the complementarity hypothesis.

---

Innovation 4: Recurrence-Based and Sparse-Attention-Based Long-Context Modeling Are Orthogonal, Not Rival, Paradigms

By the time the Routing Transformer was published, two dominant paradigms had emerged for handling long sequences in Transformers. Recurrence-based methods (Transformer-XL, Compressive Transformer) trained on short segments but maintained a memory cache of previous segments, performing cross-attention between the current segment and the cached history to approximate longer context. Sparse attention methods (Sparse Transformer, Adaptive Span) trained directly on long sequences by restricting attention to subsets of positions.

The literature largely treated these as competing approaches — different solutions to the same problem, to be compared on the same benchmarks. The Routing Transformer paper makes a subtle but important conceptual move by presenting them as orthogonal design decisions with different tradeoffs, not as rivals to be ranked (Section 6.2):

"The benefit of the Transformer-XL like approach is that it is less memory consuming and thus is able to scale to 36 layers. Sparse attention (including local attention) on the other hand is more memory expensive since it trains directly on long sequences and therefore can scale to fewer layers for the same problem. However, as we demonstrate, it is competitive with the Transformer-XL like approaches even when using fewer layers and is guaranteed to generalize to the long sequence length that it was trained on."

This framing does three things. First, it identifies the specific tradeoff: recurrence-based methods trade depth for memory efficiency (you can train deeper models because each training segment is short), while sparse attention trades depth for direct long-range modeling (you train on the full sequence, so memory limits force shallower models, but the model actually sees long-range dependencies during training rather than approximating them through a cache). Second, it establishes a bound on the comparison: since the two families operate under different constraints, direct perplexity comparisons are valid but should be interpreted with the tradeoff in mind — the Routing Transformer's 22 layers vs. Compressive Transformer's 36 layers on PG-19, for example. Third, it suggests a future direction the paper does not explore: combining the two approaches, using sparse attention within long training segments and a recurrence cache across segments, potentially capturing the benefits of both.

The "guaranteed to generalize" claim is particularly important. Recurrence-based methods are trained on short segments (e.g., 512 tokens for PG-19 in Compressive Transformer) and evaluated on long sequences by chaining the cache forward. This means they are evaluated in a regime they never saw during training — the model must generalize from 512-token training segments to 8192+-token evaluation sequences. Sparse attention models, by contrast, are trained on full-length sequences from the start. Whatever performance they achieve at evaluation length, they achieve because they were explicitly optimized for it. The paper doesn't argue that this makes sparse attention inherently better, but it identifies the generalization assumption in recurrence-based methods as an implicit weakness that sparse attention avoids by construction.

The significance of this framing goes beyond the specific comparison. It establishes a taxonomy of efficient attention methods organized by what kind of approximation is being made (temporal locality vs. segment-level recurrence vs. content-based routing) and what is being traded off (memory vs. depth vs. generalization guarantees). This makes the space of possible designs more navigable and helps future work identify which combinations of approximations might be complementary.

5. Experimental Analysis

Evaluation Methodology

• Dataset. The paper evaluates on five benchmarks spanning text and image generation. For language modeling: Wikitext-103 (Merity et al., 2017), containing over 100M tokens from ~28K Wikipedia articles with an average of 3.6K tokens per article; enwik-8 (Mahoney, 2011), the first 100M bytes of unprocessed Wikipedia typically used for character-level language modeling; and PG-19 (Rae et al., 2020), approximately 28,000 Project Gutenberg books published before 1919, consisting of 1.9B tokens with an average context size of ~69,000 words. For image generation: CIFAR-10, 60,000 32×32 color images (sequence length 3,072 when flattened) used primarily for ablation studies; and ImageNet 64×64 (Child et al., 2019), images of 64×64×3 bytes represented as sequences of length 12,288 in raster-scan RGB order. The test splits used are the standard ones released with each dataset.

• Base model(s). All experiments use variants of the Transformer architecture (Vaswani et al., 2017) as the base model, with standard components: multi-head self-attention, feedforward networks, residual connections, and layer normalization. No external pretrained model is used — all models are trained from scratch on each benchmark. Model scales vary by experiment: 10–24 layers, 8–16 attention heads, hidden dimensions appropriate for each task. The paper explicitly scales up local attention baselines to match the capacity of Routing Transformer models (e.g., a 24-layer, 16-head Local Transformer for ImageNet-64, matching the 24-layer, 16-head Routing Transformer). Relative position encodings (Shaw et al., 2018) are used for all language modeling experiments. The choice of model sizes is motivated by computational feasibility on long sequences — the paper trains on full-length sequences (8,192 tokens for enwik-8 and PG-19) rather than using segment-level recurrence.

• Metrics. For language modeling: test perplexity (lower is better), computed as $\exp(\text{negative log-likelihood per token})$. For Wikitext-103 and PG-19, perplexity is normalized by token counts as reported in the respective dataset papers. For enwik-8: bits per byte (bpc), the standard compression metric computed as negative log-likelihood per byte. For image generation: bits per dimension (bits/dim), computed as negative log-likelihood of pixel values normalized by the total number of pixel dimensions (height × width × channels). Bit/dim and bits per byte are directly comparable metrics for density estimation quality.

• Baselines. The paper compares against several strong baselines from concurrent and prior work:

  • Local Transformer: Standard Transformer with relative position encodings and local (sliding window) attention, scaled up to match the Routing Transformer's layer count and head count for fair comparison. This is the paper's own baseline, reported in Table 2, 3, 4, and 5.
  • Transformer-XL (Dai et al., 2019): Uses segment-level recurrence with a memory cache to extend effective context beyond training segment length. Reported on Wikitext-103 (18.3 perplexity, 18 layers), enwik-8 (0.99 bpc, 24 layers), and PG-19 (36.3 perplexity, 36 layers).
  • Sparse Transformer (Child et al., 2019): Uses half local and half strided attention heads. Reported on enwik-8 (0.99 bpc, 30 layers) and ImageNet-64 (3.44 bits/dim, 48 layers for the strided variant).
  • Adaptive Transformer (Sukhbaatar et al., 2019): Learns per-head attention spans via L1-penalized soft masking. Reported on Wikitext-103 (20.6 perplexity, 36 layers) and enwik-8 (0.98 bpc, 24 layers).
  • Compressive Transformer (Rae et al., 2020): Extends Transformer-XL with a compressive memory that stores compressed representations of old activations. Reported on PG-19 (33.6 perplexity, 36 layers, trained on 512-token chunks).
  • Reformer (Kitaev et al., 2020): Uses Locality-Sensitive Hashing for content-based sparse attention with fixed random projections. Reported on ImageNet-64 (3.65 bits/dim).
  • Adaptive Input (Baevski and Auli, 2019): Uses adaptive input embeddings tied to adaptive softmax. Reported on Wikitext-103 (18.7 perplexity, 16 layers).
  • Additional baselines specific to each benchmark include LSTMs, QRNNs, PixelCNN, Glow, SPN, PixelSNAIL, ImageTransformer, and T64, as listed in Tables 2–5.

• Generation budget / compute accounting. The paper measures computational cost primarily through asymptotic complexity ($O(n^{1.5}d)$ vs. $O(n^2 d)$ for full attention) and wall-clock step time on TPUv3 hardware (Tables 1 and 7). For CIFAR-10 ablations, step times are reported in steps per second, allowing direct efficiency comparisons. The "attention window" parameter controls the number of tokens each query attends to: for local heads this is the sliding window size; for routing heads this is the cluster size $w = n/k$. Window sizes are fixed at 512 or 1024 for CIFAR-10, 256 for Wikitext-103 and enwik-8, 2048 for ImageNet-64, and 512 for PG-19. The paper does not use a unified "generation budget" metric across experiments (unlike test-time compute scaling papers) — instead, efficiency is assessed by comparing models at similar architectural scales (layers, heads) and reporting step times. All training runs use 128 TPUv3 cores with the same number of cores and batch sizes within each comparison to ensure fair timing.

• Cross-validation / statistical protocol. The paper does not employ cross-validation or report confidence intervals. For language modeling, results are reported as single test-set perplexity values. The CIFAR-10 ablation study (Table 1) sweeps hyperparameters (number of routing layers, number of routing heads, attention window size) across 24 configurations, but each is a single training run. The Jensen-Shannon divergence analysis (Table 6) reports means and standard deviations over 10 runs with different random subsets of heads, but this is a post-hoc analysis of a trained model rather than a cross-validation procedure. The paper acknowledges random variation implicitly (e.g., calling the best model's improvement over full attention "not by a large enough amount to rule out noise," Section 6.1), but provides no formal statistical quantification.

Main Quantitative Results

CIFAR-10 Ablations: Understanding the Local-Routing Tradeoff

The CIFAR-10 experiments (Table 1) serve as the paper's primary ablation platform, systematically varying three hyperparameters to characterize the routing mechanism's behavior: (1) number of routing attention heads, (2) number of routing attention layers, and (3) attention window size. All models have 12 layers and 8 total heads; remaining heads and layers use local attention.

Headline pattern. The best Routing Transformer configuration (4 routing layers, 4 routing heads, 1024 attention window) achieves 2.958 bits/dim, outperforming both full attention (2.983 bits/dim) and local attention (3.009 bits/dim at 512 window; 3.048 or similar extrapolated at 1024). The worst routing-only configuration (8 routing layers, 8 routing heads, 512 window) achieves only 3.190 bits/dim, substantially worse than local-only.

Local attention as a strong baseline. Local attention alone (0 routing heads, 512 window) achieves 3.009 bits/dim, only 0.026 bits/dim worse than full attention (2.983). This 0.9% degradation establishes the surprising strength of local attention that the paper emphasizes throughout. Full attention at 3072 window achieves 2.983 bits/dim with 5.608 steps/sec — the slowest configuration in the table.

Effect of adding routing heads to a fixed number of routing layers. Holding routing layers fixed and increasing routing heads reveals a concave relationship — too few or too many routing heads hurts performance:

• With 2 routing layers: 2 heads → 3.005, 4 heads → 2.986, 8 heads → 2.992 bits/dim (512 window)
• With 4 routing layers: 2 heads → 2.995, 4 heads → 2.975, 8 heads → 2.991 bits/dim (512 window)
• With 8 routing layers: 2 heads → 2.995, 4 heads → 2.971, 8 heads → 3.190 bits/dim (512 window)

The sharp degradation at 8 routing heads × 8 routing layers (3.190, nearly as bad as the Random Transformer's 3.076) suggests that replacing too many local heads with routing heads destroys essential local processing capacity. The optimal configuration consistently uses 4 routing heads across 4–8 routing layers.

Effect of adding routing layers at fixed routing heads. The relationship is similarly non-monotonic. With 4 routing heads (512 window): 2 layers → 2.986, 4 layers → 2.975, 8 layers → 2.971, 12 layers → 2.994. The optimum sits at intermediate depth (4–8 layers); moving routing heads to all 12 layers slightly degrades performance. This supports the paper's later finding on PG-19 that routing is most beneficial in deeper layers, not throughout the entire network.

Effect of increasing attention window from 512 to 1024. Across all configurations, the larger window uniformly improves performance (Table 1, comparing top and bottom halves). The best 512-window model (4 layers, 4 routing heads) achieves 2.975; the best 1024-window model (also 4 layers, 4 routing heads) achieves 2.958 — a 0.017 bits/dim improvement. The improvement is consistent across configurations: 2 layers × 2 heads improves from 3.005 to 2.975; 8 layers × 4 heads improves from 2.971 to 2.983. This demonstrates that the clustering mechanism scales beneficially with window size — larger clusters provide more candidate tokens per query without degrading the MIPS approximation quality.

Step-time efficiency tradeoffs. Step times (Table 1, final column) reveal the computational cost of routing: local attention (9.023 steps/sec) is 1.6× faster than full attention (5.608 steps/sec). The best Routing Transformer at 512 window (4 routing layers, 4 routing heads) achieves 5.140 steps/sec — slightly slower than full attention but with better accuracy (2.975 vs. 2.983 bits/dim). A lighter routing configuration (2 routing layers, 4 routing heads) achieves 7.409 steps/sec while still substantially outperforming local attention (2.986 vs. 3.009). This demonstrates that modest routing budgets can capture most of the benefit at lower computational cost.

Random Transformer ablation. Replacing routing heads with random attention heads (same window size, but randomly selected keys instead of cluster-routed keys) achieves 3.076 bits/dim — dramatically worse than the best routing model (2.971) and even worse than local attention (3.009). This ablation confirms that the clustering mechanism is genuinely selecting informative keys, not merely providing a regularizing noise source or increased diversity. Random routing loses information relative to local attention; learned routing gains it.

Language Modeling: Wikitext-103

Headline result (Table 2). The Routing Transformer achieves 15.8 test perplexity, improving substantially over Transformer-XL (18.3, 18 layers), Adaptive Input (18.7, 16 layers), Local Transformer (19.8, 16 layers), and Adaptive Transformer (20.6, 36 layers). The Routing Transformer uses only 10 layers and 16 heads — fewer layers than all baselines except the two 16-layer models. It achieves this without segment-level recurrence (used by Transformer-XL) or adaptive mechanisms (used by Adaptive Transformer and Adaptive Input).

Model configuration. 10 layers, 16 heads, half routing (8 heads) and half local (8 heads). Routing attention uses $k = 16$ clusters and attention window 256. Relative position encodings (Shaw et al., 2018) are used. Attention dropout and ReLU dropout are both 0.3.

Comparison to Transformer-XL. The 15.8 vs. 18.3 perplexity difference represents a 13.7% relative reduction in perplexity. The Routing Transformer uses 10 layers versus Transformer-XL's 18 — 44% fewer layers — and does not require the segment-level recurrence mechanism that Transformer-XL uses to extend effective context beyond training chunks. This is a strong result because Transformer-XL was the dominant long-context language model at the time, and the Routing Transformer achieves substantially better performance with a simpler design (no separate memory cache, no cross-segment attention).

Comparison to Local Transformer. The Routing Transformer's 15.8 vs. Local Transformer's 19.8 represents a 20.2% perplexity reduction. Since both models use the same relative position encodings and the Local Transformer baseline is scaled up to 16 layers (vs. Routing's 10), the improvement can be attributed specifically to the routing mechanism. The fact that a shallower model with routing outperforms a deeper model without it suggests that routing provides efficiency gains beyond what depth alone can achieve — the routing heads are learning to use their attention budget more effectively.

No segment-level recurrence. The paper notes explicitly that the Routing Transformer achieves these results "without the need for segment level recurrence" — a mechanism used by Dai et al. (2019) and Sukhbaatar et al. (2019) to maintain an attention cache that extends beyond training batch boundaries. This is significant because segment-level recurrence introduces engineering complexity (managing the cache, handling padding, ensuring gradient flow across segment boundaries) and a generalization assumption (the model trained on short segments must handle much longer sequences at evaluation). The Routing Transformer avoids both by training directly on full-length sequences.

Language Modeling: enwik-8

Headline result (Table 3). The Routing Transformer achieves 0.99 bits per byte on enwik-8, matching Transformer-XL (0.99, 24 layers) and Sparse Transformer (0.99, 30 layers), and coming close to Adaptive Transformer (0.98, 24 layers). The Routing Transformer uses only 12 layers and 8 heads — half the layers of the 24-layer baselines that achieve comparable performance.

Model configuration. 12 layers, 8 heads, with half routing and half local. Routing attention uses $k = 32$ clusters and attention window 256. Attention and ReLU dropout are both 0.4. Relative position encodings are used. Sequence length is 8,192.

Comparison to baselines. At 12 layers, the Routing Transformer achieves bpc equivalent to models with 2–2.5× more layers (Transformer-XL at 24 layers, Sparse Transformer at 30 layers). The Adaptive Transformer edges it out by 0.01 bpc (0.98 vs. 0.99) but uses 24 layers — twice the depth. The Local Transformer baseline (24 layers, 8 heads) achieves 1.10 bpc, which the Routing Transformer improves on by 0.11 bpc while using half the layers. This pattern reinforces the Wikitext-103 finding: routing enables substantially more parameter-efficient modeling than either local attention or recurrence-based approaches.

Why the improvement is smaller than on Wikitext-103. Unlike Wikitext-103, where the Routing Transformer achieved a clear state-of-the-art (15.8 vs. next best 18.3), on enwik-8 it ties with existing methods rather than surpassing them. This difference likely reflects the character-level nature of enwik-8: character-level sequences have different dependency structures than word-level sequences, with most meaningful dependencies falling within relatively short character spans (words are typically 5–10 characters). Long-range dependencies on enwik-8 are less semantically meaningful and more about spelling consistency, which local attention may already capture adequately. The benefit of content-based routing — capturing semantically meaningful long-range token-level dependencies — is inherently more valuable for word-level modeling.

Image Generation: ImageNet 64×64

Headline result (Table 4). The Routing Transformer achieves 3.43 bits/dim, improving on the previous state-of-the-art Sparse Transformer (3.44 bits/dim, 48 layers with strided sparsity) while using half the layers (24 vs. 48). It also substantially outperforms the Reformer (3.65 bits/dim), which uses LSH-based content-based sparsity.

Model configuration. 24 layers, 16 heads, half routing and half local. Routing attention uses $k = 8$ clusters and attention window 2,048. Batch size is 1, and the model is trained for roughly 70 epochs. Sequence length is 12,288 (64 × 64 × 3 pixels, flattened in raster-scan RGB order).

Comparison to Sparse Transformer. The 3.43 vs. 3.44 bits/dim improvement is small in absolute terms (0.01 bits/dim), but the Routing Transformer achieves it with half the layers. The Sparse Transformer used 48 layers with strided sparsity patterns, while the Routing Transformer uses 24 layers with learned content-based routing. This suggests that learned routing is significantly more parameter-efficient than fixed strided patterns — it achieves equivalent or better long-range modeling with less depth, because each routing head can directly attend to semantically related tokens regardless of their spatial position, rather than relying on strided patterns that propagate information through intermediate layers.

Comparison to Reformer (LSH-based routing). The Routing Transformer's 3.43 vs. Reformer's 3.65 represents a substantial 0.22 bits/dim gap. The paper attributes this to the superiority of learned spherical k-means over fixed random LSH projections for the MIPS approximation. Both methods use content-based clustering (hash buckets vs. k-means clusters) to determine attention patterns, but the learned centroids can adapt to the data distribution while the random LSH hyperplanes remain fixed. The gap provides empirical support for the paper's decision to use learned clustering rather than LSH.

Local Transformer baseline. The scaled-up Local Transformer (24 layers, 16 heads, local attention only) achieves 3.48 bits/dim — already competitive with earlier state-of-the-art methods like SPN (3.52) and PixelSNAIL (3.52), and only 0.05 bits/dim worse than the Routing Transformer. This reinforces the paper's finding that local attention is a strong baseline, and that the routing heads provide a meaningful but modest improvement over local attention alone on this task.

Cross-modal generalization. The strong ImageNet-64 result demonstrates that content-based routing generalizes beyond text: in image generation, "long-range dependencies" manifest as spatial relationships between semantically related but spatially distant pixels (e.g., the left and right edges of a large object, or repeated texture patterns). The fact that routing attention — which was designed using language-modeling intuitions about word-level semantic similarity — transfers effectively to pixel-level generation suggests the mechanism captures a genuinely modality-agnostic principle: cluster tokens by learned similarity and attend within clusters.

Language Modeling: PG-19

Headline result (Table 5). The Routing Transformer achieves 33.2 test perplexity on PG-19, setting a new state-of-the-art and improving on the Compressive Transformer (33.6, 36 layers) and Transformer-XL (36.3, 36 layers). The Routing Transformer uses only 22 layers and 8 heads — substantially fewer parameters than the 36-layer baselines. It trains on sequences of length 8,192 directly, while the Compressive Transformer trains on chunks of size 512 and relies on a compressive memory to extend context.

Model configuration. The PG-19 model uses a different architecture than other experiments, based on the finding (Section 5.5) that long-range attention is most beneficial in deeper layers. The configuration: 22 layers, 8 heads total, but only 2 routing heads and only in the last 2 layers. The remaining 6 heads across all layers use local attention. This is the most extreme example of the paper's "routing in late layers" principle. Training uses the Adafactor optimizer (more memory-efficient than Adam), learning rate constant 0.01 with linear warmup over 10,000 steps followed by rsqrt-normalized decay, no dropout or weight decay, hidden dimension 1032, batch size 8192 tokens.

Comparison to Compressive Transformer. The 33.2 vs. 33.6 perplexity improvement is modest (0.4 perplexity, a 1.2% relative reduction) but comes from a substantially simpler architecture: the Compressive Transformer uses 36 layers, segment-level recurrence, and a learned compressive memory that stores compressed representations of old hidden states. The Routing Transformer achieves better results with 22 layers and no recurrence or external memory — just sparse attention within the sequence. This suggests that direct long-context modeling (training on 8,192-length sequences) is at least as effective as sophisticated recurrence-based approximations, and potentially more parameter-efficient since it doesn't require learned compression modules.

Comparison to Local Transformer. The Local Transformer baseline (24 layers, 8 heads, local attention only) achieves 39.3 perplexity — substantially worse than the Routing Transformer's 33.2. This is the largest margin between local and routing attention across all benchmarks (6.1 perplexity, a 15.5% reduction), suggesting that PG-19's very long documents (average ~69,000 words) contain dependencies that local attention fundamentally cannot capture, making content-based routing particularly valuable. The Routing Transformer improves on Local Transformer by 6.1 perplexity while using 2 fewer layers, reinforcing the parameter-efficiency argument.

Why fewer routing heads work best here. The paper's decision to use only 2 routing heads (vs. half the heads in other experiments) and only in the last 2 layers (vs. all layers) is motivated by an empirical finding cited from Rae and Razavi (2020): "long range attention is only needed in the last few layers." On PG-19, with documents averaging 69,000 words, the model must balance the need for long-range coherence against the risk of diluting local representations. Using many routing heads throughout the network appears to harm performance (as seen in the CIFAR-10 ablations where 8 routing heads × 8 routing layers performed worst), and concentrating routing capacity in late layers provides the best tradeoff. The paper does not present PG-19 ablations varying this choice — the configuration is presented as the result of empirical tuning rather than a systematic sweep.

Sequence length considerations. The Routing Transformer trains on sequences of 8,192 tokens. Given PG-19's average document length of ~69,000 words and the model's subword vocabulary of ~98,000 tokens, an 8,192-token sequence covers only a fraction of a typical document. This means the model is still not capturing full-document dependencies — it operates on 8K-token windows, which is substantially longer than the 512-token chunks used by Compressive Transformer, but still a fraction of the full document. The paper's claim of "training directly on long sequences" is relative to the 512-token segments used by recurrence-based methods, not relative to full PG-19 documents.

Ablation Studies and Robustness Checks

Number of routing heads vs. number of routing layers (Table 1). Sweeping both dimensions across 24 configurations reveals that performance degrades at both extremes (too few or too many routing components) and peaks at intermediate values. The optimal configuration uses 4 routing heads across 4–8 routing layers on CIFAR-10. Pure routing (8 heads, 8–12 layers) severely underperforms local attention, while adding just 2 routing heads to 2 routing layers already closes most of the gap to full attention. This demonstrates that routing is beneficial in small doses and harmful in excess.

Attention window size (Table 1). Increasing the window from 512 to 1024 uniformly improves performance across all 12 routing configurations tested, with an average improvement of roughly 0.015–0.020 bits/dim. No configuration shows degradation at the larger window, suggesting that for this sequence length (3,072), the clustering mechanism scales gracefully — larger clusters provide more candidate keys without introducing excessive noise. The paper does not test windows beyond 1024 or systematically vary window size independently of other parameters.

Random vs. learned routing (Table 1). The Random Transformer — identical to the best routing model but with cluster assignments replaced by random key selection — achieves 3.076 bits/dim, compared to 2.971 for the learned routing model and 3.009 for local attention. The random model is worse than both, confirming that neither randomness alone nor the mere existence of long-range attention suffices; the attention must be selective, routing queries to semantically similar keys. This is the cleanest evidence that the k-means clustering is genuinely solving the MIPS problem rather than acting as a regularizer.

Pure routing vs. hybrid local+routing (Table 1). The configurations with 8 routing layers and 8 routing heads (i.e., all layers and all heads are routing; no local attention anywhere) achieve 3.190 bits/dim at 512 window and 3.131 at 1024 window — dramatically worse than the best hybrid configs (~2.97). This negative result is one of the paper's most informative findings: routing attention alone cannot replace local attention. The model needs dedicated local heads to build up fluent, low-level representations before routing heads can effectively impose global consistency. This undermines any claim that routing is a general-purpose replacement for dense attention and instead positions it as a specialized mechanism for long-range dependencies that complements local processing.

Placement of routing layers (Table 1, Section 5.5). On CIFAR-10, routing layers are "always added at the top of the model" — i.e., the last N layers are routing layers and the first (12 − N) layers remain local. The paper reports that this configuration outperforms interleaving or bottom-placement (though no direct ablation on placement order is presented). On PG-19, this principle is taken further: routing heads appear only in the last 2 of 22 layers, with the first 20 layers being entirely local. The strong PG-19 result (33.2 perplexity) validates the late-layer placement strategy at scale.

Jensen-Shannon divergence analysis between attention heads (Table 6). On Wikitext-103, the JSD between local attention heads is consistently low (0.0038 to 0.3071 across layers), indicating homogeneous local attention patterns. The JSD between local and routing heads is consistently near the theoretical maximum of 0.6931 (ranging from 0.4706 to 0.6674), confirming that routing heads attend to radically different parts of the sequence than local heads. The JSD between routing heads is intermediate (0.1579 to 0.5820), suggesting they specialize in different long-range dependency types. This analysis provides the paper's strongest evidence that routing and local attention are complementary rather than substitutable.

Number of routing heads on PG-19 (implied by configuration). The PG-19 model uses only 2 routing heads out of 8 total, and only in the last 2 layers. The paper states this configuration was "motivated by our empirical finding that long range attention is only needed in the last few layers," but does not present a sweep over head count or layer placement for PG-19 specifically. The finding is presented as an architectural choice rather than an ablation result.

Centroid decay parameter. The paper states $\lambda = 0.999$ is "usually set" to this value but does not present experiments varying $\lambda$. The choice corresponds to an effective memory of ~1,000 batches. No ablation on this hyperparameter is reported.

Number of clusters $k$. The paper sets $k = 6$ for CIFAR-10, $k = 16$ for Wikitext-103, $k = 32$ for enwik-8, and $k = 8$ for ImageNet-64. These values are chosen such that the cluster size $w = n/k$ is roughly the attention window size (256–2048 depending on the experiment). The paper does not present an ablation varying $k$ independently of window size. The theoretical optimal $k = \sqrt{n}$ analysis is presented but not empirically verified across different $k$ values.

Critical Assessment

Claim 1: Routing attention reduces complexity to $O(n^{1.5}d)$ while matching or exceeding full attention performance

What the experiments demonstrate. The CIFAR-10 results (Table 1) provide the cleanest support: the best Routing Transformer at 1024 window achieves 2.958 bits/dim, actually surpassing full attention's 2.983, while using the $O(n^{1.5}d)$ mechanism. On the larger-scale benchmarks, direct comparison to full attention is infeasible (the whole point is that full attention is too expensive), so the comparison is against local attention and recurrence-based methods. On Wikitext-103 (Table 2), routing achieves 15.8 vs. local attention's 19.8 — a large improvement that implies the $O(n^{1.5}d)$ mechanism is successfully capturing dependencies that local attention misses. On PG-19 (Table 5), routing achieves 33.2 vs. local's 39.3 — an even larger margin.

What the experiments do not demonstrate. The paper never compares the Routing Transformer to full attention on sequences long enough that the $O(n^{1.5})$ vs. $O(n^2)$ distinction actually matters. On CIFAR-10 (sequence length 3,072), full attention is tractable, and routing barely outperforms it (2.958 vs. 2.983). On the long-sequence benchmarks where the complexity reduction is actually necessary (Wikitext-103, enwik-8, ImageNet-64, PG-19), full attention is infeasible, so we cannot assess how much performance is lost relative to the "gold standard" of dense attention. The comparison is always routing vs. local attention or routing vs. recurrence — not routing vs. full attention. The paper is thus evaluating routing as an improvement over other efficient approximations, not as a lossless compression of full attention.

Conditional assessment. The claim holds in the regime where full attention is too expensive to run but local attention leaves performance on the table. The claim does not hold as "routing equals full attention" — it holds as "routing outperforms local attention and recurrence-based methods on long sequences, and on short sequences where full attention is feasible, routing roughly matches it." The $O(n^{1.5}d)$ complexity reduction is real and asymptotically significant, but the absolute speedup on the tested sequence lengths (8,192) is approximately 1.7× slower wall-clock time than local attention (Table 7) — not the 90× speedup over full attention that the asymptotic analysis would suggest for $n = 8192$ (since full attention is not run at that length for comparison).

Claim 2: Routing attention and local attention are complementary, providing different and both-necessary functions

What the experiments demonstrate strongly. This is the paper's most robustly supported claim. Three independent lines of evidence converge:

1. Pure routing (8 heads, 8+ layers, no local attention) performs dramatically worse than pure local attention on CIFAR-10 (3.190 vs. 3.009 at 512 window; Table 1). Neither alone is optimal.
2. The hybrid model (half local, half routing) outperforms both pure variants and, in the best configuration, full attention.
3. The JSD analysis (Table 6) shows that local and routing attention distributions are essentially non-overlapping (JSD near the theoretical maximum of 0.6931), while routing heads differ substantially from each other. This is direct quantitative evidence that routing heads are doing something qualitatively different from local heads, not just "attending farther."

What the experiments do not demonstrate. The complementarity claim is established on CIFAR-10 and (via JSD) on Wikitext-103, but is not systematically verified on the other benchmarks. The ImageNet-64 and PG-19 configurations use hybrid models (half routing or 2 routing heads), but there are no ablations on these benchmarks showing that removing either component degrades performance. The paper also does not systematically characterize what routing attention captures that local attention misses — the JSD shows the distributions differ, but provides no semantic interpretation of that difference.

Conditional assessment. Strongly supported for the claim that local and routing attention serve different functions, with the caveat that the optimal ratio of local to routing heads appears task-dependent (half-and-half for Wikitext-103 and ImageNet-64; only 2 of 8 heads and 2 of 22 layers for PG-19). The paper's architectural choices adapt to this implicitly, but it does not provide a principled way to determine the optimal ratio for a new task.

Claim 3: Local attention is a surprisingly strong baseline

What the experiments demonstrate comprehensively. Across all five benchmarks, scaled-up local attention achieves results within striking distance of state-of-the-art methods:

• Wikitext-103: 19.8 (local) vs. 18.3 (Transformer-XL) — 8% gap
• enwik-8: 1.10 (local) vs. 0.99 (Transformer-XL) — 11% gap in bpc
• ImageNet-64: 3.48 (local) vs. 3.44 (Sparse Transformer, 48 layers) — 1% gap
• PG-19: 39.3 (local) vs. 36.3 (Transformer-XL) — 8% gap

This is genuinely surprising: a model with no mechanism for long-range attention beyond the local window achieves ~90–99% of state-of-the-art performance. The claim is a diagnostic contribution that raises the bar for future efficient attention methods.

What the experiments do not demonstrate. The paper does not systematically analyze why local attention performs so well. Is it because most dependencies in these benchmarks are genuinely local? Because information propagates effectively through intermediate layers despite the lack of direct long-range connections? Because the benchmarks have been implicitly designed around models with limited context? The paper speculates on these questions but conducts no experiments to distinguish them. Additionally, the scaled-up local baselines are the paper's own implementations — prior work's local attention baselines may have been undertuned, but the paper does not replicate those undertuned versions to quantify the gap.

Conditional assessment. The claim is well-supported as an empirical observation. Its generalizability to other tasks, modalities, and sequence lengths is untested but plausible. The claim functions primarily as a calibration of expectations rather than a theoretical contribution.

Claim 4: Content-based routing outperforms LSH-based routing (Reformer)

What the experiments demonstrate. On ImageNet-64, Routing Transformer achieves 3.43 bits/dim versus Reformer's 3.65 bits/dim (Table 4) — a meaningful 0.22 bits/dim gap. This is the only direct comparison in the paper between k-means routing and LSH routing.

What the experiments do not demonstrate. The comparison is cross-paper — the Reformer result is from Kitaev et al. (2020), not from the authors' own replication. The two models may differ in aspects other than the clustering method: optimization details, hyperparameters, training duration, model architecture, data preprocessing, etc. The paper does not implement LSH routing within its own framework to perform a controlled ablation. The claim that "spherical k-means is known to outperform LSH for MIPS" (Section 2) cites prior work (Auvolat et al., 2015) but is not empirically verified within the attention context. A direct head-to-head comparison of k-means routing vs. LSH routing within the same architecture and training setup would substantially strengthen this claim.

Conditional assessment. The cross-paper comparison is suggestive but not conclusive. The claim that learned centroids outperform fixed random projections is theoretically plausible and consistent with the MIPS literature, but the experimental evidence provided is circumstantial — it relies on comparing two different papers' results on the same benchmark rather than a controlled experiment.

Claim 5: Routing attention with clustering approximates Maximum Inner Product Search

What the experiments demonstrate. The theoretical argument (Equations 10–13) establishes that, for unit-normalized vectors, assigning queries and keys to the same cluster via k-means preserves high dot products. The empirical evidence is indirect: the Random Transformer (Table 1) performs substantially worse than the Routing Transformer (2.971 vs. 3.076 bits/dim), confirming that the clustering is not random — it is selecting keys that are genuinely useful for attention. But this does not directly test the MIPS approximation quality.

What the experiments do not demonstrate. The paper does not measure how well the clustered attention approximates full attention. A direct test would be: for a model where full attention is computable (e.g., CIFAR-10), compare the actual attention weights assigned by routing attention to the attention weights that full attention would have assigned to the same query-key pairs. How often does routing attend to the keys that full attention would have given the highest weights? What is the correlation between routing attention weights and full attention weights for the pairs that routing selects? None of these analyses are presented. Without them, the MIPS framing remains a motivating analogy rather than a verified mechanism.

Conditional assessment. The MIPS analogy is theoretically sound and provides useful intuition, but the paper does not empirically validate that clustering approximates MIPS well in practice — it only shows that clustering outperforms random selection and enables strong downstream performance. These are different claims.

Genuine weaknesses in the experimental design

1. No direct comparison to full attention on long sequences. The paper's central motivation is the infeasibility of $O(n^2 d)$ attention for long sequences, but all comparisons on long-sequence benchmarks are against other efficient approximations (local attention, recurrence), not against full attention. We never learn how much performance is sacrificed to achieve the $O(n^{1.5}d)$ complexity — because we can't measure it. This is an inherent limitation of the problem setting, not a flaw in the paper, but it means the central claim about "matching or exceeding full attention" is only demonstrated on the one benchmark (CIFAR-10) where full attention is feasible, and there the margin is tiny (0.025 bits/dim).

2. No systematic hyperparameter sweeps on large benchmarks. The CIFAR-10 ablations (Table 1) sweep 24 configurations, but the large-scale benchmarks (Wikitext-103, enwik-8, ImageNet-64, PG-19) report single configurations with no ablations. This is understandable given computational constraints, but it means we don't know how sensitive the headline results are to the specific choice of $k$, routing head count, layer placement, or window size. The CIFAR-10 sweeps show that performance varies substantially across configurations — the optimal CIFAR-10 model required searching over 24 settings. The large-benchmark results may represent lucky single configurations rather than robust optima.

3. Limited analysis of what routing attention learns. The JSD analysis (Table 6) confirms that routing attention differs from local attention, but provides no semantic interpretation. What kinds of token pairs are routed together? Do routing heads specialize in particular dependency types (coreference, topic, syntax)? The paper hypothesizes about "gender, nouns, dates and names of places" being connected by routing, but provides no empirical evidence from attention pattern analysis to support this.

4. Single training run per configuration. All results are reported from single training runs. No error bars, confidence intervals, or multi-seed results are reported. The paper acknowledges noise implicitly (Section 6.1: "not by a large enough amount to rule out noise" for CIFAR-10 routing vs. full attention), but does not quantify it. With test sets of 500 (ImageNet-64) to thousands (Wikitext-103) of examples, and model sizes in the tens of millions of parameters, run-to-run variance could be meaningful.

5. The PG-19 model configuration is not ablated. The Routing Transformer's state-of-the-art PG-19 result (33.2 perplexity) uses an unusual configuration — only 2 routing heads, only in the last 2 of 22 layers — that differs from every other experiment. The paper does not present ablations showing that this specific configuration is optimal for PG-19. The configuration may reflect extensive trial-and-error tuning, and a different choice might yield substantially different results.

6. Step-time comparisons are hardware-specific. The step-time comparisons (Tables 1, 7) are reported on TPUv3, and the paper notes that "the lack of support for sparse operations on the TPU" limits routing attention's speed. It speculates that GPU sparse kernels (Gale et al., 2020) could significantly improve wall-clock performance, but provides no GPU benchmarks. The practical speedup of routing over local attention may be substantially different on different hardware.

7. Missing experiments that would have strengthened the paper. A controlled comparison of k-means routing vs. LSH routing within the same architecture; a direct measurement of MIPS approximation quality (how often does routing select the keys that full attention would have weighted most highly?); a systematic sweep of $k$ (number of clusters) on a large benchmark to validate the $\sqrt{n}$ theoretical optimum; an analysis of attention patterns showing what linguistic or visual phenomena routing heads capture; and multi-seed experiments to establish statistical reliability of the headline results.

6. Limitations and Trade-offs

Limitation 1: The Difficulty Estimation Cost Is Unaccounted for in Efficiency Claims

The assumption or constraint. The paper claims a complexity reduction from $O(n^2 d)$ to $O(n^{1.5} d)$ based on setting the number of clusters $k = \sqrt{n}$. The dominant cost terms are $O(n k d)$ for centroid assignment plus $O(n^2 d / k)$ for within-cluster attention, which balance at $k = \sqrt{n}$ to yield $O(n^{1.5} d)$. However, this analysis omits a mandatory cost that the paper acknowledges explicitly in Section 4.1:

"This adds an additional $O(n \log n)$ term to the cost, however note that this is eclipsed by the dominating term of $O(n^{1.5}d)$."

The paper is referring to the sorting step required to select the top-$w$ tokens per cluster: for each of the $k$ centroids, the $n$ tokens must be sorted by distance, costing $O(n \log n)$ per centroid, for a total of $O(k n \log n) = O(\sqrt{n} \cdot n \log n) = O(n^{1.5} \log n)$.

The consequence. The claimed $O(n^{1.5} d)$ complexity is only accurate if $d$ is large enough that the $O(n^{1.5} d)$ term dominates the $O(n^{1.5} \log n)$ sorting term. For the model configurations used in the paper, this holds: $d$ ranges from hundreds to over a thousand, while $\log n$ for $n = 8192$ is approximately 13. But the sorting term has a non-trivial constant factor — top-$k$ selection on TPU/GPU hardware requires either a full sort or a selection algorithm, both of which involve irregular memory access patterns that are substantially slower per-operation than dense matrix multiplications. The practical consequence is visible in the step-time measurements: the Routing Transformer is consistently slower than the Local Transformer at equivalent sequence lengths (5.140 vs. 9.023 steps/sec on CIFAR-10 at 512 window, Table 1; 0.7236 vs. 1.231 steps/sec on PG-19, Table 7). The asymptotic analysis says routing should be faster than full attention by a factor of $\sqrt{n} \approx 90$ for $n = 8192$, and faster than local attention for sufficiently long sequences; the wall-clock measurements show routing is 1.22–1.76× slower than local attention, and only marginally faster than full attention on CIFAR-10 (where full attention is runnable — 5.140 vs. 5.608 steps/sec).

What evidence exists in the paper. Table 1 provides direct step-time comparisons: every Routing Transformer configuration is slower than the Local Transformer baseline at the same window size, despite having the same or better asymptotic complexity. Table 7 shows the same pattern on PG-19: the Routing Transformer trains at 0.7236 steps/sec versus 1.231 for Local Transformer — a 1.7× slowdown. The paper does not decompose this slowdown into its components (sorting overhead vs. gather/scatter overhead vs. centroid update overhead), so we cannot attribute it specifically to the sorting term, but the aggregate timing data makes clear that the practical efficiency benefit over local attention is negative at the sequence lengths tested.

Mitigation status. The paper partially acknowledges this limitation in Section 6.3:

"This trade-off with respect to speed compared to the Local Transformer is due to the lack of support for sparse operations on the TPU; on the GPU various sparse kernels have been proposed which promise to significantly speed up training of these models."

The paper defers to future hardware/software improvements (Gale et al., 2020) for closing the gap between theoretical and practical efficiency, and explicitly states that "wall-clock time efficiency is only a secondary goal." This is an honest characterization, but it means a practitioner deciding whether to deploy the Routing Transformer must weigh a 1.2–1.8× training slowdown against the accuracy gains, with no guarantee that the asymptotic advantage materializes at current hardware-software maturity levels.

---

Limitation 2: The Claim That Routing Outperforms LSH-Based Hashing (Reformer) Rests on a Cross-Paper Comparison, Not a Controlled Experiment

The assumption or constraint. A central motivation for using learned spherical k-means over LSH with fixed random projections is the claim that "spherical k-means is known to outperform LSH for MIPS" (Section 2), supported by a citation to Auvolat et al. (2015) and the empirical observation that the Routing Transformer achieves 3.43 bits/dim on ImageNet-64 while the Reformer (Kitaev et al., 2020) achieves 3.65 bits/dim. The paper presents this comparison as evidence that k-means routing is superior to LSH routing for content-based sparse attention.

The consequence. The 3.43 vs. 3.65 bits/dim comparison (Table 4) is cross-paper, not cross-method within a single controlled framework. The two models differ along numerous dimensions beyond the clustering mechanism:

• Architecture: The Routing Transformer uses 24 layers with 16 heads (half local, half routing), while the Reformer configuration from Kitaev et al. (2020) likely differs in layer count, head count, and local/routing head ratio.
• Training recipe: Different optimizers, learning rates, batch sizes, training duration, dropout rates, and data preprocessing pipelines.
• Model-scale interactions: The Reformer's LSH mechanism is embedded in a model that also uses reversible residual layers and chunked feedforward computation to reduce memory — design choices the Routing Transformer does not share.
• Hardware and implementation: The two models were trained on potentially different hardware with different software implementations, making wall-clock comparisons even less reliable.

Any of these confounds could account for the 0.22 bits/dim gap independently of the clustering method. A practitioner choosing between k-means routing and LSH routing cannot determine from this evidence alone whether the performance difference is attributable to the routing mechanism or to other architectural and training choices.

What evidence exists in the paper. The only comparison between routing and LSH appears in Table 4 as two numbers drawn from two different papers. The paper does not implement LSH routing within the Routing Transformer architecture as an ablation, nor does it attempt to control for architecture, training, or implementation differences. The CIFAR-10 ablations (Table 1) do include a Random Transformer for comparison with routing, but the Random Transformer selects keys uniformly at random — it is not an LSH baseline that would directly test the k-means vs. LSH hypothesis. Section 2 states the theoretical motivation (spherical k-means outperforms LSH for MIPS "see e.g. Auvolat et al. (2015)" and "this is borne out in the common task of Imagenet-64 generation"), but the empirical citation merely observes the correlation without isolating the causal factor.

Mitigation status. The paper does not acknowledge this as a limitation. It presents the cross-paper comparison as if it were a direct ablation. The theoretical argument from prior MIPS literature provides some support for the superiority of k-means over LSH, but the paper does not supplement this with a controlled within-framework experiment. The claim remains plausible but unverified — a missed opportunity, since implementing LSH routing as an ablation within the Routing Transformer's own codebase would have been straightforward (replace learned centroids with random hyperplane projections, keep everything else identical). The paper's headline claim that routing outperforms LSH-based approaches should be considered suggestive rather than established.

---

Limitation 3: Local Attention Is Competitive With Routing on All Benchmarks, and the Relative Gain Shrinks on Larger-Scale Tasks

The assumption or constraint. The paper's central architectural claim is that routing attention provides complementary value to local attention, capturing long-range dependencies that local windows miss. This claim implies that routing should provide substantial performance improvements over local attention, and that this improvement should be larger on benchmarks with longer-range dependencies. The paper reports the relative performance in Tables 2–5.

The consequence. The gains from routing over local attention, while real, are surprisingly modest when both models are scaled comparably:

• Wikitext-103 (Table 2): Routing Transformer 15.8 vs. Local Transformer 19.8 — a 20.2% relative reduction in perplexity. This is the largest gain across all benchmarks.
• enwik-8 (Table 3): Routing Transformer 0.99 bpc vs. Local Transformer 1.10 bpc — a 10.0% relative reduction in bits per byte. Modest.
• ImageNet-64 (Table 4): Routing Transformer 3.43 bits/dim vs. Local Transformer 3.48 bits/dim — a 1.4% relative reduction. Very small.
• PG-19 (Table 5): Routing Transformer 33.2 vs. Local Transformer 39.3 — a 15.5% relative reduction. Substantial, but note that this PG-19 Local Transformer baseline (24 layers, 8 heads, 39.3 perplexity) is actually better than Transformer-XL (36 layers, 36.3 perplexity) when normalized by the token count metric from Rae et al. (2020).

The pattern across benchmarks is telling: on tasks where the Local Transformer baseline is already strong (ImageNet-64: 3.48 vs. Sparse Transformer's 3.44), routing adds marginal value. On tasks where local attention struggles (PG-19 with documents averaging 69,000 words, far exceeding the local window), routing provides more substantial gains. But even on PG-19, the 15.5% improvement over local attention is comparable to the 8.5% improvement of Compressive Transformer (33.6) over Transformer-XL (36.3) — and the Compressive Transformer uses a completely different mechanism (compressive memory with recurrence) rather than sparse attention.

A practitioner evaluating whether to adopt routing attention should consider: the implementation complexity of online k-means clustering, centroid updates, gather/scatter operations, and balanced top-$w$ selection is substantially higher than implementing local attention (which is just a banded matrix). The performance gain must justify this complexity. On ImageNet-64, a 0.05 bits/dim improvement may not justify the engineering burden. On PG-19, a 6.1 perplexity improvement likely does — but only if PG-19-style very-long-document modeling is the target application.

What evidence exists in the paper. The numbers cited above are drawn directly from Tables 2–5. The paper itself emphasizes the strength of local attention as a baseline (Section 6.1):

"local attention is slightly worse than full attention - 3.009 vs 2.983 bits per dim. Adding 2 routing layers with 4 heads almost closes the gap with the performance of full attention, achieving 2.986 bits per dim."

The paper is transparent about this, but the broader implication — that routing's value is highly task-dependent and sometimes marginal — is left implicit rather than discussed as a deployment consideration.

Mitigation status. The paper does not present this as a limitation; instead, it uses the strength of local attention to argue for the complementarity hypothesis (routing adds value on top of local, rather than replacing it). This is a valid framing, but it obscures a practical question: for tasks where local attention already achieves 95%+ of routing performance, is the additional complexity worth it? The paper provides no guidance on when routing is likely to help beyond the qualitative observation that "text that is 10× longer in context than all prior data-sets" (PG-19) benefits more. A cost-benefit analysis — in terms of implementation complexity, training slowdown, and accuracy gain — would help practitioners decide.

---

Limitation 4: No Empirical Validation That $k = \sqrt{n}$ Is the Optimal Number of Clusters for Practical Sequence Lengths

The assumption or constraint. The paper derives $k = \sqrt{n}$ as the optimal number of clusters by balancing two cost terms: centroid assignment $O(n k d)$ and within-cluster attention $O(n^2 d / k)$. Setting $k = \sqrt{n}$ equalizes these terms and yields the $O(n^{1.5} d)$ headline complexity. This derivation is purely asymptotic — it ignores constant factors, the sorting cost, and hardware-specific properties of the operations.

The consequence. The optimal $k$ in practice may differ substantially from $\sqrt{n}$ for several reasons:

1. Constant factors differ dramatically between operations. Dense matrix multiplication (the within-cluster attention, implemented as batched $w \times w$ matrix multiplies) is one of the most optimized operations on modern accelerators. Nearest-neighbor search and sorting (the centroid assignment and top-$w$ selection) involve irregular memory access and are comparatively slow per FLOP. This means the centroid assignment term $O(n k d)$ may have a much larger constant factor than the attention term $O(n^2 d / k)$, which would shift the optimal $k$ downward (fewer, larger clusters).

2. Hardware constraints on cluster sizes. The within-cluster attention requires clusters of size $w = n/k$. If $w$ is too small, the matrix multiplications become too small to saturate GPU/TPU parallelism. If $w$ is too large, the attention itself becomes the bottleneck. The optimal $w$ from a hardware utilization perspective is independent of the asymptotic analysis.

3. Cluster quality degrades as $k$ increases. Spherical k-means with more clusters partitions the space more finely. For very large $k$, some clusters may contain few semantically related tokens and many "leftover" tokens that happen to be closest to that centroid. This degrades the MIPS approximation quality — tokens that should attend to each other are separated, and tokens that shouldn't attend are forced together. The asymptotic analysis assumes balanced, high-quality clusters regardless of $k$, which is unrealistic.

4. The paper's own $k$ choices deviate from $\sqrt{n}$ in practice:

   • CIFAR-10: $n = 3072$, $\sqrt{n} \approx 55$, paper uses $k = 6$
   • Wikitext-103: $n = 4096$ (implied by window 256 and the fact that routing clusters span the full sequence), $\sqrt{n} = 64$, paper uses $k = 16$
   • enwik-8: $n = 8192$, $\sqrt{n} \approx 90$, paper uses $k = 32$
   • ImageNet-64: $n = 12288$, $\sqrt{n} \approx 111$, paper uses $k = 8$

In every case, the paper uses $k \ll \sqrt{n}$, by factors ranging from 3× to 14×. This is not a flaw — it reflects practical engineering judgment — but it means the $O(n^{1.5} d)$ complexity derived under the $k = \sqrt{n}$ assumption is not the complexity of the actual deployed models. The actual models have $k$ closer to a small constant (6–32), making the dominant cost term $O(n^2 d / k)$ — which is $O(n^2 d)$ with a smaller constant factor rather than asymptotically $O(n^{1.5} d)$. The throughput advantage over full attention at the tested sequence lengths comes from the small $w = n/k$, not from the $n^{1.5}$ asymptotic.

What evidence exists in the paper. The paper reports the specific $k$ values used in each experiment (Sections 5.1–5.5) but does not present an ablation over $k$. The CIFAR-10 ablations (Table 1) vary attention window size (512 vs. 1024) but hold $k = 6$ fixed, so we see the effect of varying $w$ given fixed $k$, not the effect of varying $k$ given fixed $w$ or fixed budget. The step-time measurements (Tables 1, 7) reflect the actual $k$ choices, not the theoretical optimum. A practitioner cannot determine from this paper whether different $k$ values would yield better accuracy-efficiency tradeoffs, or whether the $k = \sqrt{n}$ optimum is ever practically relevant.

Mitigation status. The paper does not acknowledge the gap between theoretical $k = \sqrt{n}$ and practical $k \ll \sqrt{n}$ as a limitation. The asymptotic analysis is presented as the theoretical motivation, and the actual $k$ values are presented as engineering choices without reconciliation. A practitioner should understand that the $O(n^{1.5} d)$ claim describes an idealized scaling regime that the paper's own experiments do not operate in — the actual models achieve efficiency through small constant-factor improvements (attending to $n/k \approx 256$–$512$ tokens instead of $n \approx 3000$–$12000$ tokens) rather than through the asymptotic scaling advantage.

---

Limitation 5: No Analysis of What Routing Attention Actually Learns — The MIPS Claim Remains Unverified Mechanistically

The assumption or constraint. The paper frames routing attention as an approximation to Maximum Inner Product Search (MIPS): by clustering unit-normalized queries and keys and attending within clusters, the model preserves high-dot-product pairs while discarding low-dot-product pairs. The theoretical derivation (Equations 10–13) establishes that if $Q_i$ and $K_j$ are both assigned to the same centroid, then $Q_i^\top K_j > 1 - 2\varepsilon^2$, i.e., their dot product is high. This is presented as the mechanism by which routing works: "This analysis shows that our clustering routing strategy preserves large attention weights as non-zero entries" (Section 4.1).

The consequence. The deductive chain from "same cluster" to "high dot product" to "preserved attention weight" has a critical missing empirical link: the paper never measures whether the tokens that routing attention actually attends to are the tokens that full attention would have weighted most highly. There are several ways this could fail:

1. False negatives: A query-key pair could have a high dot product but be assigned to different clusters. This happens when both $Q_i$ and $K_j$ are near the boundary between two centroids, or when the clustering is suboptimal. The theoretical bound only guarantees high dot product given same-cluster assignment; it does not guarantee that all high-dot-product pairs will be co-clustered.

2. False positives: A query-key pair could have a low dot product but share a cluster, especially when clusters are balanced via top-$w$ selection: if a centroid doesn't have enough truly close tokens, it will be padded with tokens that are merely "closest among the remaining," which may not actually be close.

3. Attention weight distortion: Even when the correct pairs are co-clustered, the softmax normalization within the cluster may assign different relative weights than full-attention softmax would, because the denominator now includes only $w$ keys rather than all $n$ keys. Keys with moderate dot products that would have received small but non-zero attention in full attention are excluded entirely; the remaining weights are scaled up proportionally.

The JSD analysis (Table 6) addresses a different question — it shows that routing attention distributions differ from local attention distributions, which confirms they are attending to different tokens. But it does not address whether routing attention approximates full attention. The Random Transformer baseline (Table 1) shows that random key selection is worse than learned routing, which confirms that the clustering is non-random. But neither analysis tests the specific claim that routing attention preserves the high-dot-product pairs from full attention.

What evidence exists in the paper. No experiment in the paper directly measures MIPS approximation quality. The paper does not:

• Compare routing attention weights to full attention weights on CIFAR-10 (where full attention is computable).
• Report recall@$w$ — the fraction of the top-$w$ full-attention keys that routing attention includes in its cluster.
• Show examples of token pairs that are routed together and whether they correspond to semantically meaningful relationships.
• Measure the correlation between cluster assignment quality (intra-cluster dot product) and downstream task performance.

The closest the paper comes to mechanistic analysis is the JSD table (Table 6) and the qualitative hypothesis in Section 6.1:

"we conjecture that for every time step, the prediction depends on a small support of high value tokens: local attention facilitates local consistency and fluency, while a full dot product attention would facilitate global consistency."

This is presented as a conjecture, not a verified mechanism.

Mitigation status. The paper does not acknowledge this as a limitation. The MIPS framing is presented as a motivating analogy and theoretical justification, but the empirical validation only tests whether routing works (improves perplexity/bits-per-dim over baselines), not whether it works for the claimed reason (approximating MIPS). This is a common pattern in ML papers — a theoretically motivated design is validated on downstream metrics without verifying the intermediate mechanism — but it leaves open the possibility that routing helps for reasons orthogonal to the MIPS analogy (e.g., it acts as a regularizer that prevents overfitting to local statistics, or it provides a beneficial inductive bias toward global coherence that happens to correlate with but is not equivalent to MIPS). A practitioner who wants to improve upon routing attention or diagnose its failures would benefit from knowing the mechanism, not just the outcome.

---

Limitation 6: Single Model Family and Modality — No Evidence of Transfer Across Architectures or Tasks Beyond Text/Image Generation

The assumption or constraint. All experiments in the paper use a single base architecture: the standard Transformer (Vaswani et al., 2017) with multi-head self-attention, feedforward layers, and relative position encodings (Shaw et al., 2018) for language tasks. All experiments are on auto-regressive generative modeling tasks — language modeling and image generation — where the task is to predict the next token given previous tokens, and the evaluation metric is perplexity or bits-per-dimension (density estimation quality).

The consequence. The paper's claims about routing attention — that it captures complementary long-range dependencies, that it benefits from being placed in later layers, that it works best when combined with local attention — are established in this specific context. Several important deployment scenarios are untested:

1. Encoder-decoder architectures (e.g., machine translation, summarization). These use cross-attention between encoder and decoder in addition to self-attention. The routing mechanism would need to handle query-key pairs from different sequences (source and target), where the MIPS argument about dot product similarity has different semantics — a target-language query might need to attend to source-language keys with very different surface forms.

2. Bidirectional (non-causal) self-attention (e.g., BERT-style encoders). The Routing Transformer's causal masking workaround (sharing keys and queries, Algorithm 1 line 5) is specific to auto-regressive models. The paper notes that "for encoder self attention and encoder-decoder cross-attention, additional masking or sharing queries and keys is not necessary" (Section 4.1), but provides no experiments demonstrating that routing works in these settings.

3. Tasks beyond generative modeling (e.g., classification, question answering, retrieval). The paper only evaluates on density estimation (next-token prediction). The value of long-range attention for classification tasks — where the model must aggregate information from the entire sequence into a single representation — might be different from generation tasks, where the model predicts token-by-token.

4. Different model scales. The largest Routing Transformer has 22 layers (PG-19). Recent large language models operate at 70B+ parameters and 80+ layers. Whether the online k-means routing mechanism remains stable and beneficial at those scales — where the embedding space may have different geometric properties, and where the centroids must partition a much higher-dimensional space — is unknown.

5. Different hardware architectures. The paper trains on TPUv3 and notes that sparse operations lack TPU support (Section 6.3). GPU-specific implementations with sparse kernels (Gale et al., 2020) could change the efficiency tradeoff. Inference — where batching, latency constraints, and memory bandwidth differ from training — is not evaluated at all.

What evidence exists in the paper. The paper acknowledges the scope limitation implicitly by testing on two modalities (text and image) and noting in the conclusion that "our approach could prove useful in domains where the inputs are naturally sparse, such as 3D point clouds, social networks, or protein interactions" — but this is speculative. The paper provides only an existence proof: routing attention works on auto-regressive generation with Transformers on text and images. It does not establish whether the approach generalizes across architectures, task types, or scale.

Mitigation status. The paper does not frame this as a limitation, and does not conduct experiments on non-auto-regressive tasks, encoder-decoder architectures, or at larger scale. The generalization gap is inherent in the paper's scope — it demonstrates effectiveness on the benchmarks it targets — but a practitioner considering routing attention for a different architecture or task cannot extrapolate from the presented results with confidence. The paper suggests extensions (Section 7) but provides no empirical foothold for evaluating their likelihood of success.

7. Implications and Future Directions

How This Work Changes the Landscape

The Routing Transformer makes a methodological contribution that shifts the framing of efficient attention from "which fixed sparsity pattern works best?" to "can we learn which tokens should attend to each other before computing attention?" This is not a paradigm shift — it does not replace the Transformer or even the broader family of sparse attention methods — but it successfully reframes the problem in a way that opens a new design axis.

Prior to this work, the landscape of efficient attention was organized around a fundamental tradeoff: position-based sparsity (local windows, strided patterns) was computationally efficient but inflexible, while content-based sparsity (sparsemax, entmax) was flexible but required instantiating the full attention matrix — defeating the purpose of sparsity for efficiency. The Routing Transformer is the first method to demonstrate that this tradeoff is not inherent: you can have content-dependent sparsity patterns and natively sparse computation, provided you have a mechanism (clustering) to route queries to keys before computing dot products.

The significance of this reframing goes beyond the specific k-means clustering choice. It establishes that the MIPS framing — treating attention sparsity as an approximate maximum inner product search problem — is productive. This connects efficient attention to decades of research on similarity search, nearest neighbors, and vector quantization. Methods from those fields (hierarchical clustering, product quantization, graph-based ANN search) become natural candidates for future attention mechanisms in a way they were not before. The Reformer (Kitaev et al., 2020), published contemporaneously, independently made a similar move with LSH — the two papers together establish content-based routing as a viable paradigm, not a one-off trick.

The paper also shifts the conversation about model depth and attention span in a specific diagnostic way. The finding that routing heads are most effective in late layers — used to extreme effect on PG-19 (2 routing heads in the last 2 of 22 layers) — provides concrete evidence that long-range attention is not uniformly needed across the network. This aligns with Rae and Razavi (2020) and gives architectural guidance: efficient attention mechanisms can be targeted rather than applied everywhere. The earlier-generation approach (Child et al., 2019) of applying sparse patterns uniformly across all layers implicitly assumed that every layer benefits equally from global information. The Routing Transformer's architecture suggests otherwise, and the strong PG-19 result validates this design principle.

The paper resolves a latent tension in the literature: how can local attention — which completely discards long-range dependencies — be competitive with sophisticated recurrence and sparsity mechanisms on standard benchmarks? The answer, as the paper demonstrates, is that scaled-up local attention was consistently undertuned in prior work. By reporting strong local attention baselines across all benchmarks (19.8 on Wikitext-103, 1.10 on enwik-8, 3.48 on ImageNet-64, 39.3 on PG-19), the paper recalibrates expectations. A new efficient attention mechanism should be measured against a well-tuned local attention model at comparable scale, not against whatever local attention configuration was easiest to run. This raises the bar for the field in a productive way: methods that claim large gains over "local attention" without reporting their local baseline's scale and tuning should be viewed skeptically.

Finally, the paper provides the first direct comparison between learned routing (k-means) and fixed random routing (LSH-based, via the cross-paper comparison with Reformer) on a major benchmark. The 3.43 vs. 3.65 bits/dim gap on ImageNet-64 is large enough to motivate learning the routing function rather than fixing it, though the cross-paper nature of the comparison limits its conclusiveness. The paper establishes learned routing as the default choice pending controlled experiments.

Follow-Up Research This Work Enables

Controlled k-means vs. LSH routing ablation within a single framework. The paper's comparison between Routing Transformer (3.43 bits/dim on ImageNet-64) and Reformer (3.65 bits/dim) is cross-paper, confounded by differences in architecture, training recipe, and implementation. A direct ablation would implement LSH routing within the Routing Transformer's own codebase — replacing the learned centroids with random hyperplane projections while keeping everything else (layer count, head count, training hyperparameters, data pipeline) identical. The expected outcome is that k-means routing outperforms LSH routing by a meaningful but smaller margin than the cross-paper comparison suggests. A null result (LSH matches k-means) would indicate that the benefits of routing come from the cluster-and-attend structure itself, not from learning the partition. This experiment is straightforward to implement and would resolve the paper's central unanswered question about whether learned routing is worth the additional mechanism.

Direct measurement of MIPS approximation quality. The paper claims k-means routing approximates Maximum Inner Product Search — preserving high-dot-product query-key pairs while discarding low-dot-product ones — but never measures this directly. On a benchmark where full attention is tractable (CIFAR-10 or a moderate-length language modeling task), one could compute the recall@$w$: for each query, what fraction of the top-$w$ keys by full-attention dot product are included in the routing cluster? This measurement could be done per-layer and per-head, revealing whether routing's approximation quality degrades in deeper layers, whether some heads approximate MIPS well and others poorly, and whether recall@$w$ correlates with downstream performance. Low recall would indicate that routing works for reasons other than MIPS approximation (e.g., it implicitly regularizes attention or provides beneficial noise). This experiment would transform the MIPS framing from a theoretical analogy into a verified mechanism.

Systematic sweep over $k$ at fixed compute budget to find the practical optimum. The paper derives $k = \sqrt{n}$ as asymptotically optimal but then uses $k \ll \sqrt{n}$ in every experiment (e.g., $k = 8$ for $n = 12288$ on ImageNet-64, where $\sqrt{n} \approx 111$). A systematic experiment would sweep $k$ across a wide range (e.g., from 2 to $\sqrt{n}$) on a fixed benchmark and model configuration, reporting both accuracy and wall-clock time at each setting. This would reveal the practical Pareto frontier: how much accuracy is sacrificed by using small $k$ (which is faster due to smaller centroid assignment cost), and whether the theoretical optimum is ever practically relevant. A finding that accuracy is flat or even improves at $k \ll \sqrt{n}$ would suggest that the asymptotic analysis misses important effects — perhaps large $k$ over-fragments the embedding space, reducing cluster quality. This experiment would provide the practical guidance the paper currently lacks about how to set $k$ for a new task.

Attention pattern interpretability study — what do routing heads learn? The Jensen-Shannon divergence analysis (Table 6) shows routing attention differs quantitatively from local attention, but provides no semantic interpretation. A qualitative study could examine specific token pairs that routing heads attend to on a language modeling task like Wikitext-103: do routing heads consistently connect pronouns to antecedents? Named entity mentions across paragraph boundaries? Syntactic heads to distant dependents? Topic-related content words? This could be done by annotating a sample of high-attention-weight query-key pairs from routing heads and categorizing the linguistic relationship. If routing heads specialize in particular dependency types (e.g., one head primarily handles coreference, another handles topic coherence), this would explain the JSD between routing heads (Table 6: 0.1579–0.5820) and suggest that routing is capturing multiple distinct long-range phenomena. A null finding — that routing heads attend to seemingly random long-range tokens — would suggest routing helps via a different mechanism (perhaps providing a beneficial noise source or enabling the model to average over more diverse context).

Routing in encoder-decoder and bidirectional architectures. The paper's experiments are exclusively auto-regressive generative models with causal attention. A natural extension would apply routing attention to an encoder-decoder task like machine translation (WMT, IWSLT) and to a bidirectional encoder task like masked language modeling (BERT-style). The key question for encoder-decoder is whether routing works for cross-attention, where queries and keys come from different sequences with different statistical properties — a French query may need to attend to English keys with no surface similarity, challenging the MIPS assumption. The key question for bidirectional encoders is whether routing benefits from the additional context (attending to both past and future tokens) or whether the larger candidate set makes clustering harder. A positive result in both settings would establish routing as a general-purpose attention component; a negative result would bound its applicability to auto-regressive generation. Both findings are valuable.

Scaling to larger models — does routing remain stable at 1B+ parameters? The largest Routing Transformer in the paper has 22 layers and 8 heads (PG-19), which is small by modern standards. As model scale increases, the embedding dimension grows, the number of layers grows, and the geometry of the representation space changes in ways that could affect clustering: higher-dimensional spaces are "emptier" (most points are far from most other points), which could degrade k-means cluster quality. A scaling study would train Routing Transformers at increasing model sizes (e.g., 100M, 1B, 10B parameters) on a fixed benchmark, measuring whether the relative gain over local attention grows, shrinks, or stays constant with scale. The hypothesis from the paper — that routing captures long-range dependencies local attention misses — would predict that the gain grows with scale, because larger models have more capacity to exploit long-range information. A finding that the gain shrinks would suggest routing provides a regularization benefit that matters more at small scale. This experiment is expensive but would determine whether routing attention is relevant for the largest models.

Combining routing with segment-level recurrence. The paper explicitly frames sparse attention and recurrence-based methods (Transformer-XL, Compressive Transformer) as orthogonal approaches with complementary tradeoffs: recurrence allows deeper models on shorter segments, while sparse attention trains directly on long sequences (Section 6.2). A natural combination would augment a Routing Transformer with a Transformer-XL-style memory cache: train on long segments with routing attention within each segment, and maintain cached key-value representations from previous segments for cross-segment attention. This could combine the depth advantage of recurrence (train on manageable segment lengths) with the direct long-range modeling of routing (the model actually sees and clusters long-range dependencies during training). The hypothesis is that this hybrid would outperform either approach alone, particularly on very long documents like PG-19 where even 8,192-token segments cover only a fraction of the full document.

Practical Applications and Downstream Use Cases

Long-document language modeling and generation (PG-19-scale tasks). The Routing Transformer's strongest relative gain over local attention is on PG-19 (33.2 vs. 39.3 perplexity, a 15.5% reduction; Table 5), where documents average ~69,000 words. For applications involving book-length text modeling — long-form story generation, document summarization, legal document analysis, academic literature review — the Routing Transformer provides a deployment-ready architecture that trains on 8,192-token sequences directly without requiring segment-level recurrence infrastructure. A practitioner working with PG-19-like data can use the paper's specific configuration (22 layers, 2 routing heads in the last 2 layers, $k$ such that $w \approx 512$, Adafactor optimizer) as a starting point. The 15.5% perplexity reduction over local attention translates directly to better generation quality and more accurate perplexity-based evaluations for downstream tasks.

High-resolution image generation and pixel-level modeling. On ImageNet-64, the Routing Transformer achieves 3.43 bits/dim (Table 4), improving on the previous state-of-the-art Sparse Transformer (3.44) while using half the layers (24 vs. 48). For applications involving auto-regressive image generation — medical imaging, satellite imagery, texture synthesis, or any domain where images are large enough that the 12,288-pixel flattened representation stresses memory — the Routing Transformer provides a drop-in replacement for dense or strided attention that preserves global coherence (e.g., consistent textures across the image, symmetric features like eyes in faces) without the $O(n^2)$ cost. The finding that routing generalizes across modalities (text and images, tested in the paper) suggests it may apply to other sequential representations of spatial data — video frames, 3D point clouds, audio spectrograms — though this remains to be tested.

Efficient training pipelines for sparse long-sequence models. The paper provides a practical recipe for training sparse attention models on long sequences: use half local heads for fluency, add a small number of routing heads in late layers for global consistency, set $k$ such that the cluster size $w = n/k$ matches the desired attention budget (typically 256–2048 tokens per query), and update centroids via an EMA during training. This recipe is immediately applicable to practitioners who need to train Transformers on sequences longer than a few thousand tokens but do not want to implement segment-level recurrence or complex sparsity patterns. The CIFAR-10 ablation table (Table 1) provides concrete guidance on the tradeoff between routing heads and speed: 2 routing layers × 4 routing heads trains at 7.409 steps/sec while nearly matching full attention quality (2.986 vs. 2.983 bits/dim), making it suitable for applications where training throughput matters.