CoLT5: Reading long documents with selective computation

Suppose a model must answer a question about a lengthy report. It needs enough context to find the relevant evidence, but not every passage deserves the same amount of processing. Repeated headings, background material, and the sentence containing the answer all consume space in the input; their usefulness to the task can be very different.

CoLT5 gives every token a lightweight path through the encoder, then spends extra computation on selected tokens. It learns which tokens to select separately at each layer and for each kind of operation. The aim is to improve the balance between answer quality and processing time, rather than simply making the model smaller.

Joshua Ainslie and colleagues introduced CoLT5, short for Conditional LongT5, in March 2023. It builds on LongT5 with conditional computation in the encoder, multi-query cross-attention for faster decoding, and a UL2-style pretraining objective. Their experiments cover long-document question answering, summarization, and inference over contracts, with a separate study extending input length to 64k tokens. Original paper.

Why efficient attention is only part of the problem🔗

A token is a chunk of text processed by the model. A token's representation is the vector of numbers the network uses to describe it at a particular layer. In an encoder-decoder model such as T5, the encoder builds representations of the input, and the decoder uses them to generate an answer or summary.

Two operations do much of the encoder's work. Attention lets token representations draw information from other positions. A feed-forward network, also called an MLP, transforms each position's representation independently. Attention also needs learned projections that turn representations into queries, keys, and values.

Ordinary full attention compares every position with every other position, so the number of comparisons grows quadratically with input length. LongT5 reduces that cost using local attention and, in its transient-global variant, summaries of token blocks. But the feed-forward networks and attention projections still process every token. Once attention becomes cheaper, these operations can dominate the computational budget.

CoLT5 therefore asks a second question: can the expensive feed-forward and projection work also be concentrated on fewer tokens? Paper, section 2.

One light path for everyone, extra work for selected tokens🔗

Each CoLT5 encoder layer has light and heavy branches for both attention and feed-forward processing:

Light branches process every token. They maintain a representation for the full input using local attention and a narrower feed-forward network.
Heavy branches process selected tokens. They add a wider feed-forward transformation or attention across selected positions from throughout the input.
Routers choose the tokens. The choices depend on the current input representations and can change from one layer to the next.

Figure 1. The CoLT5 encoder layer, reproduced from the paper's figure 1. The symbols $q$, $v$, and $m$ count routed attention queries, key/value tokens, and feed-forward tokens. The diagram reads from bottom to top.

Think of reading a report once for context and revisiting selected passages more carefully. The analogy has a limit: CoLT5 does not permanently discard the other passages or perform a single document-level selection. All tokens keep their light-path representations, and selection is repeated throughout the encoder.

Routing: What does "important" mean?🔗

A router gives each token a score by comparing its current representation with a learned vector. It selects the highest-scoring positions within a chosen budget. Training teaches the router which selections help the overall task; it does not receive human labels declaring particular words important.

Each layer has three independent routers: one for heavy feed-forward processing, one for heavy attention queries, and one for heavy attention keys and values. A token can therefore be useful for one operation without being selected for another.

This distinction matters when answering a question. Some positions may need information from elsewhere in the document; other positions may contain the evidence those positions need. CoLT5 can select these two groups separately. In the paper's TriviaQA analysis, question tokens and tokens matching the answer were more likely to receive heavy processing than other tokens, though routing is not a human-readable explanation of the model's reasoning. Paper, sections 3.1 and 4.6.

The routing scores and feed-forward update

The paper denotes token $i$'s representation by $X_i$ and a learned routing vector by $u$. Its routing score is:

$$ s_i = X_i \cdot u. $$

The dot product produces one score per token. A generalized softmax normalizes these scores so that their sum equals the target routing budget. The resulting normalized score is written $\tilde{s}_i$. The router selects the top-scoring tokens, and the score weights their heavy-branch contribution. This weighting supplies a training signal for the routing vector; a hard selection alone would not do that.

Using the paper's notation, the feed-forward update is:

$$ \begin{aligned} X_i ={}& X_i + \mathrm{FFd}_{\mathrm{Light}}(X_i) \\ &+ \tilde{s}_i \cdot \mathrm{FFd}_{\mathrm{Heavy}}(X_i). \end{aligned} $$

This is an in-place update: the right-hand side uses the token's incoming representation. The first term preserves that representation, the second adds light processing, and the third adds the weighted heavy result. In this update, $\tilde{s}_i$ is set to zero for non-routed tokens. The implementation computes the heavy branch only for selected positions, rather than computing it everywhere and multiplying unwanted outputs by zero.

The paper uses an entropy-regularized optimization procedure from Qian et al. to normalize scores. Appendix C specifies 50 iterations and $\epsilon=1.0$, and permits the top $9k/8$ tokens to have nonzero weight during training to improve the learning signal. This is a training detail beyond the simpler top-$k$ explanation above. Paper, section 3.1 and appendix C.

Feed-forward branches: Wider is not deeper🔗

The light feed-forward branch uses half the standard T5 hidden width, while the heavy branch uses roughly four times that width. The heavy branch is more capable per selected token, but only a small fraction of the input uses it.

The branches do not achieve this by using different numbers of encoder layers. Corresponding CoLT5 and LongT5 models have the same layer count and model width. CoLT5 has more total parameters because it includes the additional heavy branch, yet accesses those parameters sparsely.

Model dimensions and a worked compute example

Figure 2. Model configurations from the paper's table 7. The original image is linked at full size. Parameter count measures stored capacity, not how much computation is performed on each token.

For example, the Large models both have 24 layers and a model dimension of 1,024. LongT5-L uses feed-forward hidden width 2,816; CoLT5-L uses 1,408 in the light branch and 11,264 in the heavy branch.

For its simplified FLOP accounting, the paper writes the feed-forward cost as:

$$ \mathrm{FLOPs}_{\mathrm{FFd}} = 8nr_Ld^2 + 8mr_Hd^2. $$

Here, $n$ is input length, $m$ is the number of routed feed-forward tokens, $d$ is model dimension, and $r_L$ and $r_H$ are the light and heavy hidden-width ratios. Substituting the main experiment's ratios, $r_L=1/2$, $r_H=4$, and $m=n/16$, gives:

$$ \mathrm{FLOPs}_{\mathrm{FFd}} = 4nd^2 + 2nd^2 = 6nd^2. $$

That is 75% of the standard $8nd^2$ feed-forward cost in this accounting. The heavy network is wider, but it runs on sufficiently few tokens to reduce the combined work. This is a component-level FLOP estimate, not a prediction of a 25% reduction in total running time. The paper counts each multiply-add as one FLOP. Paper, sections 2 and 3.1.

Attention branches: Separate who asks from who supplies information🔗

An attention query represents a position seeking information. Keys are used to score possible matches, and values supply the information combined into the output.

The light attention branch gives every token access to a local window. The heavy branch lets selected queries attend to a separately selected set of keys and values from across the input. Here, "global" means that the selected positions can be far apart; it does not mean that every token attends to every other token.

In the main 16k-input setup, the routing budgets are:

1,024 tokens for the heavy feed-forward branch;
1,024 query tokens for heavy attention; and
2,048 key/value tokens for heavy attention.

The query and feed-forward budgets have the same size, but their routers need not choose the same positions. The key/value budget is twice as large, allowing a smaller group of queries to gather evidence from a larger group of source positions. The heavy attention results are added to the selected queries' light-path updates. Paper, section 3.1.

Two other changes complete the recipe🔗

Multi-query cross-attention speeds up the decoder🔗

Making the encoder cheaper does not automatically make answer generation fast. The decoder repeatedly reads the encoded document as it generates output tokens, and moving the keys and values through memory can become a bottleneck for long inputs.

CoLT5 uses multi-query attention (MQA) in the decoder's cross-attention layers. Query heads remain separate, but share keys and values. This reduces the amount of key/value data that must be repeatedly read. MQA addresses decoder memory bandwidth, while conditional computation addresses encoder work; they solve different parts of the runtime problem. Paper, section 3.2; original MQA paper.

UL2 supports learning from examples in the input🔗

CoLT5 replaces LongT5's PEGASUS pretraining objective with a variant of UL2, a mixture of text-reconstruction tasks. Its recipe combines prefix language modeling, which predicts a continuation, with span corruption, which reconstructs removed spans of different lengths.

This supports in-context learning: the model receives examples of the desired task in its input, then responds to a new example without updating its weights. A longer input can accommodate more demonstrations as well as the document to be processed.

The paper tests this on Natural Questions and TriviaQA. For the 16k-context experiment, it first continues pretraining a CoLT5-Large model at that length for another 100,000 steps. The authors found that in-context learning did not reliably extend beyond the length used in training. Thus, the demonstration is not simply a 4k-trained model being handed a much longer prompt. Paper, sections 3.3 and 4.4; UL2 paper.

What the experiments show🔗

Quality and speed at roughly 16k input length🔗

The main comparison uses 16,384 input tokens, except ContractNLI, which uses 8,192. It covers nine datasets: TriviaQA and arXiv, plus the seven tasks in SCROLLS. Those tasks span document question answering, summarization, and identifying whether a contract supports a given statement.

Figure 3. The paper's figure 2, showing the main quality-speed comparison, not the separate 64k experiment. Higher and farther left is better. Batched timings are normalized per sample per TPUv4 chip; they are not single-request latency measurements. The LongT5 inference comparison includes MQA for a conservative baseline.

At Large size, both models have an average score of 45.3 across the reported tasks. CoLT5-XL reaches 47.4 versus 46.6 for LongT5-XL. CoLT5-Base is faster but has a lower average score, 42.4 versus 43.1. The improvement is therefore a better overall quality-speed trade-off, not a win on every task at every model size.

The paper reports 35-75% training speedups and 50-100% inference speedups for Large and XL in this comparison, beyond MQA's contribution. Its CoLT5-XL result was state of the art on the SCROLLS leaderboard at the time. The average combines different task metrics: answer F1, exact match, and ROUGE-based summarization scores. It is not a single accuracy percentage. Paper, table 3 and section 4.2.

The 64k result: More context without more routed tokens🔗

The long-input scaling study uses NarrativeQA, where answering questions can require information from lengthy stories. The authors compare LongT5-Large and CoLT5-Large at increasing input lengths. CoLT5 achieves a better quality-speed trade-off and continues to gain answer F1 when its input grows from 32k to 64k.

For that last increase, the heavy routing budget stays fixed: 2,048 feed-forward tokens and queries, with 4,096 key/value tokens. The model sees twice as much input without doubling the number of heavy-branch tokens. Their fraction of the input therefore falls.

This supports the intuition that a longer document may contain proportionally fewer positions that need expensive processing. It does not show that the absolute number of useful tokens must decrease, or that the total model runs in sublinear time. Paper, section 4.3.

Does conditional attention remove quadratic scaling?

Heavy attention compares $q$ selected queries with $v$ selected key/value positions, so its comparison work depends on $qv$. With the main fixed fractions, $q=n/16$ and $v=n/8$, their product is $n^2/128$. This is a calculation of query-key pairs, not the paper's complete attention FLOP formula: projections, heads, and local attention also contribute.

The smaller fraction greatly reduces the quadratic term's coefficient, but does not eliminate it. With a fixed heavy routing budget, that part stops growing with input length. The light branches still process every token, and the routers still need to score the input. The 32k-to-64k result demonstrates useful scaling of a selected-token budget, not constant-time or sublinear processing of the entire document.

What the ablations teach us🔗

The authors also change individual parts of the recipe to test their contribution:

Learned routing matters. Selecting evenly spaced positions instead lowers the average score in the Base ablation from 42.5 to 40.5.
The routing budget trades speed for quality. Raising the feed-forward/query budget from 512 to 1,024 helps; further increases show diminishing returns in the tested setup.
Not every query needs every key/value. Letting selected queries attend to the entire input gives little average quality improvement at greater cost.
MQA is a deliberate trade-off. Ordinary multi-head cross-attention gives a slightly higher average score in the ablation, but much slower inference.
UL2 is not the source of every fine-tuning gain. PEGASUS is slightly stronger on the fine-tuned task average in this ablation; UL2 is retained for its in-context-learning capability.

These results are useful when interpreting the architecture: routing, context length, decoder bandwidth, and the training objective affect different aspects of performance. Paper, section 4.5.

Where the approach is useful🔗

CoLT5 is aimed at encoder-decoder tasks with long inputs: finding evidence in documents, summarizing reports or meetings, and answering questions about extended narratives. It is especially relevant when only part of the input needs costly global interactions, but retaining local representations of the rest remains useful.

The experiments suggest three practical questions for evaluating such a design: does extra context improve the task, how much quality is lost when the heavy-token budget shrinks, and is the runtime bottleneck in the encoder or decoder? Parameter count alone cannot answer them. CoLT5's heavy branches increase stored capacity, while conditional execution reduces the work applied to each example.

Sources and implementation notes🔗

This explanation follows the March 2023 CoLT5 paper. The authors describe an implementation using JAX, Flax, Flaxformer, and T5X. That implementation description should not be confused with a downloadable pretrained CoLT5 checkpoint.

Phil Wang's CoLT5-attention implements conditionally routed components in PyTorch. It is not an official release of the paper's complete trained models, and its README identifies implementation choices that differ from, or interpret gaps in, the paper.

For the routing background, the original CoLT5 paper cites Qian et al., Multi-Vector Retrieval as Sparse Alignment. The community implementation also points to Conditional Adapters (CoDA) and Wright's coordinate-descent review. CoDA's arXiv submission came after CoLT5's original March release, so it is better treated here as related reading than as the original paper's stated source.

Other foundations are LongT5, T5.1.1, multi-query attention, and UL2. The figures reproduced from the CoLT5 paper are attributed to Ainslie et al.; the original preprint is distributed under CC BY 4.0.

The takeaway🔗

CoLT5 keeps a lightweight representation of the whole input while concentrating additional capacity on learned selections of tokens. Its separate routers distinguish positions that need information from positions that can supply it. Together with faster cross-attention and a suitable pretraining objective, this lets the model make better use of long documents within a given computation budget.