Michał Chromiak's blog - Optimization

QLoRA: How to fine-tune large language models with less memory

2023-05-26T09:33:00+02:00

Adapting a language model to a new task usually costs much more GPU memory than simply running it. Training needs room for the model, intermediate calculations, and the information used to update its weights. For a model with billions of parameters, those extra costs can put fine-tuning out of reach even when inference fits.

QLoRA combines two ways to save memory: compress the pretrained model and train a small set of additional weights. In their 2023 paper, Tim Dettmers and colleagues use this approach to fine-tune a 65-billion-parameter LLaMA model on a single 48 GB GPU. They also introduce Guanaco, a family of chatbots trained with QLoRA. Paper, introduction.

The method is easiest to understand in two steps: first decide which weights need to change, then reduce the memory used by the weights that stay fixed.

Why fine-tuning needs so much memory🔗

A pretrained model's weights, also called parameters, are the numbers it learned during its original training. Fine-tuning adapts those numbers, or adds trainable ones, using examples of the behavior we want. For instruction tuning, an example might pair a question with a useful answer.

In full fine-tuning, all the model's weights can change. Training also stores gradients, which indicate how to adjust trainable weights, and optimizer state, the running statistics an optimizer uses to choose those adjustments. Activations, the intermediate results produced as inputs pass through the network, require memory too.

The paper estimates that conventional 16-bit fine-tuning of LLaMA 65B requires more than 780 GB of GPU memory. That is a training-memory estimate, not just the size of the model's weights. QLoRA reduces several parts of this memory budget together. Paper, sections 1 and 2.

LoRA: Learn a small update instead of rewriting the model🔗

LoRA, short for Low-Rank Adaptation, keeps the pretrained weights frozen: training does not update them. It adds small trainable components called adapters alongside selected layers. Each layer combines its original computation with an adjustment learned by its adapter.

For example, when adapting a model to answer customer-support questions, the base model continues to supply its learned language capabilities. Training adjusts the adapters using the new examples. The resulting behavior depends on both parts, even though only the adapters change.

An adapter represents an update using two smaller matrices, or grids of numbers, instead of one full-size matrix. Their narrow intermediate dimension is the adapter rank. This restriction greatly reduces the number of trainable parameters and the gradients and optimizer state associated with them. Ordinary LoRA still needs to keep the large base model in memory; QLoRA addresses that remaining cost.

For a fuller explanation of rank, adapter training, and merging updates for inference, see LoRA: Fine-tuning a model by learning a small update.

The algebra: How two small matrices change a layer

The QLoRA paper writes the LoRA computation as follows in equation (3):

$$ \mathbf{Y} = \mathbf{X}\mathbf{W} + s\mathbf{X}\mathbf{L}_1\mathbf{L}_2. $$

Here, $\mathbf{X}$ is the layer's input and $\mathbf{Y}$ its output. The frozen matrix $\mathbf{W}$ has shape $h \times o$, where $h$ and $o$ are the input and output widths. The trainable matrices $\mathbf{L}_1$ and $\mathbf{L}_2$ have shapes $h \times r$ and $r \times o$. Their product has the same shape as $\mathbf{W}$, but rank at most $r$. The scalar $s$ controls the size of the adapter's contribution.

For an illustrative layer with input and output widths of 4,096 and rank 8, the two adapter matrices contain 65,536 parameters in total. A full 4,096-by-4,096 weight matrix contains 16,777,216. The adapter therefore uses 256 times fewer trainable parameters for this layer. This is a count of parameters, not a prediction of the total GPU-memory saving.

In QLoRA, the base matrix is stored in quantized form and reconstructed approximately for computation. The adapters remain higher precision. The paper's equations (5) and (6) give the full expression, including the reconstruction of quantization scales. Paper, sections 2 and 3.

Figure 1. The paper's comparison of the three methods. Blue arrows show parameter updates, green arrows show gradient flow, and pink arrows show movement of optimizer state between GPU and CPU memory. In LoRA and QLoRA, gradients pass through the base model's computation to train the adapters, while the base weights stay fixed. Source: Dettmers et al., figure 1.

What QLoRA adds🔗

QLoRA builds on LoRA with three memory-saving components. Each addresses a different cost: storing the weights, storing the information needed to reconstruct those weights, and handling temporary peaks during training.

1. NormalFloat: Store the base weights in 4 bits🔗

Quantization replaces high-precision numbers with a smaller set of representable values. Four bits provide 16 possible codes. Each stored weight selects one of those codes, which maps to an approximate numerical value.

QLoRA introduces 4-bit NormalFloat (NF4), whose representable values are arranged for a zero-centered normal distribution: a bell-shaped distribution with most values near zero. It places more of its limited precision in that crowded region, where many pretrained weights lie, rather than spacing all values evenly.

The weights are divided into small blocks, each with its own scaling factor. This lets the same set of NF4 values represent blocks with different numerical ranges. The paper finds that NF4 performs better than the other 4-bit formats it tests.

Four-bit storage does not mean four-bit arithmetic throughout training. When the model uses a group of quantized weights, it dequantizes them into a computation format, usually 16-bit BFloat16 (BF16), before matrix multiplication. Dequantization reconstructs an approximation; it does not recover the exact original weights. The adapters learn while operating alongside that approximation. Paper, sections 3 and 4.

2. Double quantization: Compress the scales too🔗

The scaling factors also occupy memory. If every block of 64 weights has a 32-bit scale, that scale adds half a bit per weight. Across billions of weights, this overhead is substantial.

Double quantization quantizes these scaling factors as well. The paper uses 8-bit values for the first set of scales, with higher-precision scales shared across larger groups. This reduces the scale overhead from 0.5 to approximately 0.127 bits per weight, saving about 0.37 bits per parameter, or roughly 3 GB for a 65B model.

The second quantization therefore targets the scales, not a second round of four-bit compression applied directly to the model weights. Paper, section 3.

3. Paged optimizers: Make room for temporary memory peaks🔗

A training run can fit in memory most of the time and still fail on a batch containing long sequences. The extra intermediate calculations create a temporary peak.

QLoRA's paged optimizers use NVIDIA unified memory to move optimizer state between GPU memory and CPU RAM as needed. When GPU memory is under pressure, some of that state can reside in CPU RAM; it returns when the optimizer needs it. This helps accommodate peaks without reserving enough GPU memory for all optimizer state at all times.

Paging works alongside gradient checkpointing, which saves activation memory by recomputing some intermediate results during the backward pass. Neither technique removes the need to budget for sequence length and batch size. Paper, sections 3 and 4.

One training step, from input to update🔗

Putting the pieces together, a QLoRA training step works like this:

Run the examples through the model. Quantized base weights are dequantized as needed for computation, and the adapters contribute their learned adjustments.
Measure the prediction error. The model predicts tokens, the chunks of text it processes. A loss measures how well those predictions match the training targets.
Backpropagate through the computation. Gradients pass through the frozen model's operations to determine how the adapter weights should change.
Update the adapters. The base weights stay fixed. Paged optimizer state can move between CPU and GPU memory when needed.

This explains how a frozen model can participate in training: its computations influence the adapter gradients, even though training never updates its own weights.

What the experiments establish🔗

The authors compare QLoRA with both full 16-bit fine-tuning and 16-bit LoRA. On the evaluated tasks, QLoRA broadly matches these baselines while using much less memory. The comparison depends on scale: the largest LLaMA experiments, from 7B to 65B, compare against 16-bit LoRA, not full fine-tuning of every weight at every size.

Those LLaMA experiments use MMLU, a multiple-choice benchmark spanning 57 subjects. Table 4 reports mean accuracy of 53.1% for NF4 with double quantization and 53.0% for the 16-bit adapter baseline, averaged across the tested model sizes and tuning datasets. Individual settings vary, but the overall result is close.

Adapter placement matters too. The paper finds that applying LoRA across all linear layers in the transformer blocks is important for matching full fine-tuning in its LLaMA 7B experiment. Adding adapters only to the query and value projections in attention does not achieve the same result. Paper, section 4.

Guanaco and the single-GPU results🔗

QLoRA is the training method; Guanaco is a family of models trained with it. The authors build Guanaco by instruction-tuning LLaMA on OpenAssistant's OASST1 conversation data. They use supervised learning, without a reinforcement-learning stage.

The original release reports two particularly useful hardware reference points:

Guanaco 33B: fine-tuning on a single 24 GB GPU, with a reported run taking less than 12 hours.
Guanaco 65B: fine-tuning on a single 48 GB GPU, with a reported run taking about 24 hours.

These are results for the authors' training recipe, not fixed runtimes for arbitrary datasets and sequence lengths. They show how the memory reductions make much larger models accessible on a single GPU. Paper, introduction and section 4; authors' implementation.

What "99.3% of ChatGPT" means🔗

The widely repeated figure comes from Guanaco 65B's score relative to the 2023 ChatGPT baseline in the paper's Vicuna benchmark evaluation. GPT-4 scored responses to 80 prompts; the authors expressed Guanaco's score as a percentage of ChatGPT's and averaged over both response orders to reduce ordering bias. The resulting 99.3% is a relative score under that evaluation, not a measurement of 99.3% of ChatGPT's overall capabilities.

The paper also compares human and model-based judgments and examines failures that aggregate scores miss. Its broader lesson is that the evaluation should match the intended use: strong performance on a multiple-choice knowledge benchmark need not imply equally strong conversational answers. Training-data suitability matters as well as size. Paper, sections 5 and 6.

A closer look: Design choices and memory costs🔗

The following comparisons show where the gains come from. The first tests adapter placement; the second tests number formats; the worked example separates weight storage from the rest of the training budget.

Why does adapter placement matter?

The paper compares several ways to fine-tune LLaMA 7B on Alpaca. Attention layers combine information from different positions in the text. Feed-forward layers, abbreviated FFN, transform the representation at each position. Adapters can be added to either group or both.

Figure 2. Adapter-placement comparison from Dettmers et al., figure 2. The original figure is linked at full size.

How to read it. Each dot is a run with a different random seed. The vertical axis is ROUGE-L, a measure of overlap between generated and reference answers based on their longest common subsequence. Higher is better on this metric. Blue dots use four-bit base weights; orange dots show the 16-bit full-fine-tuning baselines.

QLoRA-All places adapters across all linear layers in the transformer blocks. It scores above the FFN-only and attention-only variants and is comparable to the authors' tuned full-fine-tuning baseline, labeled "Alpaca (ours)." The lower Stanford-Alpaca results also show why tuning the baseline matters: outperforming a weak training setup would be less convincing.

The useful distinction is where the model can learn an adjustment, not just how few parameters it updates. This experiment supports broad adapter coverage in this setup; ROUGE-L alone does not establish general chatbot quality. Paper, section 4.

What do NF4 and double quantization each contribute?

The paper also isolates the effect of the number format by evaluating quantized LLaMA models. Here, zero-shot means the evaluation supplies no worked examples in the prompt. This comparison tests the quantized models, rather than Guanaco's instruction-tuned conversational behavior.

Figure 3. Quantization comparison from Dettmers et al., figure 3. The original figure is linked at full size.

How to read it. The horizontal axis is total model bits on a logarithmic scale; farther left means less storage. The vertical axis is mean accuracy across five benchmarks: WinoGrande, HellaSwag, PIQA, ARC-Easy, and ARC-Challenge. Higher means better performance on that collection of tasks.

The orange NormalFloat curve sits above the blue ordinary-float curve. Both use four-bit weight codes, but choosing the representable values differently changes the error introduced by quantization. The green NormalFloat-plus-DQ curve stays close to the orange curve while shifting left: compressing the scales saves storage with little change in accuracy here.

The two components therefore do different jobs. NF4 improves the representation of the weights; double quantization reduces the overhead of their scales. These results explain the choice of quantization scheme. The fine-tuning comparisons in table 4 address the separate question of how well the models perform after adapter training. Paper, section 4.

A worked memory budget for 65 billion weights

Start with an idealized model containing exactly 65 billion weights. A byte contains eight bits, so storing every weight in 16 bits requires 130 GB. Four bits per weight reduces that to 32.5 GB. These calculations use decimal gigabytes and count only the stored weight values.

What is stored	Approximate size
Weight values at 16 bits each	130 GB
Weight values at 4 bits each	32.5 GB
Four-bit values plus the paper's double-quantized scale overhead	33.5 GB

The last row adds approximately 0.127 bits per weight for scales, using the paper's block sizes: 64 weights per first-level block and 256 scales per second-level block. That contributes about 1 GB across 65 billion weights. Without double quantization, the 0.5-bit scale overhead would contribute about 4.1 GB instead. This is where the roughly 3 GB saving comes from. Paper, section 3.

The 33.5 GB estimate is not the total training requirement. A real run also needs adapters, their gradients and optimizer state, activations, temporary computation buffers, and any parameters kept at higher precision. Sequence length and batch size affect several of these costs. Gradient checkpointing reduces saved activations through recomputation; paging lets some optimizer state reside in CPU RAM when GPU memory is tight.

This calculation explains how the frozen weights can fit within the paper's 48 GB training setup while leaving some room for the rest. It is an illustrative storage budget, not a reproduction of the authors' measured peak memory. Their appendix G examines the memory footprint in more detail.

Figures 2 and 3 are reproduced unmodified from the original paper by Dettmers et al., distributed under CC BY 4.0.

The takeaway🔗

QLoRA makes fine-tuning more accessible by separating the large body of pretrained weights from the much smaller set of weights that must change. LoRA reduces the trainable part; NF4 and double quantization compress the frozen part; paged optimizers help handle temporary memory pressure.

The result is a practical route to adapting larger models within a limited GPU-memory budget. The paper's strongest message is the combination: carefully chosen quantization and well-placed adapters can preserve strong task performance without the memory cost of updating the whole model.

For a different approach to training efficiency, see I-JEPA: Learning from images by predicting missing features. I-JEPA learns visual representations from unlabeled images; QLoRA adapts pretrained language models. The articles address different stages of learning, rather than two versions of the same method.

For an architectural example built around dense matrix operations, see MLP-Mixer: How image patches communicate without attention. MLP-Mixer is a vision backbone, whereas QLoRA is a memory-efficient way to adapt a pretrained language model; the connection is computational perspective, not task or method.

Original sources🔗

Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, and Luke Zettlemoyer. QLoRA: Efficient Finetuning of Quantized LLMs, 2023. This explanation follows the original May 2023 version.
QLoRA code and Guanaco model information, maintained by the authors.

CoLT5: Reading long documents with selective computation

2023-03-30T14:33:00+02:00

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.

LoRA: Fine-tuning a model by learning a small update

2023-03-18T13:01:00+01:00

Suppose you want a language model to turn questions into database queries. It already knows a great deal about language, but it needs examples of the task and the kind of answer you expect. Fine-tuning can teach that behavior. The expensive part is allowing billions of existing weights to change and storing everything needed to update them.

LoRA keeps the pretrained weights fixed and learns a small adjustment alongside them. Instead of training a full-size update for each chosen weight matrix, it represents that update with two much smaller matrices. The model still uses all its pretrained weights; only the additional matrices are trained.

Edward J. Hu and colleagues introduced Low-Rank Adaptation in June 2021, and the paper appeared at ICLR 2022. Its central question is practical: how much of full fine-tuning's performance can we retain while learning and storing far fewer task-specific parameters? Original paper.

Why train an update instead of the whole model?🔗

A model's parameters, often called weights, are the numbers it learns during training. Full fine-tuning starts from a pretrained model and updates all of those parameters on a new dataset. It is not training from scratch, but it still requires gradients and optimizer state for the weights being updated.

Gradients indicate how the training loss changes with each parameter. An optimizer such as Adam also keeps running statistics to help choose the updates. These extra tensors can take much more memory than the final model checkpoint alone. Intermediate results, called activations, add another substantial cost.

There is also a storage problem. If each task gets a complete fine-tuned copy of a large model, every new task adds another large checkpoint. LoRA separates what can be shared, the pretrained model, from what changes between tasks, the learned update.

This makes LoRA a parameter-efficient fine-tuning (PEFT) method. Parameter-efficient means that relatively few parameters are trained; it does not mean the large base model disappears from memory.

Two paths through a layer🔗

A dense layer multiplies an input vector by a matrix of learned weights. LoRA adds a second path alongside that operation:

The original matrix processes the input using its frozen weights.
A small trainable matrix, called $A$, maps the same input into a narrow intermediate space.
Another trainable matrix, $B$, maps that intermediate result back to the layer's output width.
The layer adds the original result and the learned adjustment.

Figure 1. The LoRA construction from Hu et al., figure 1. Read from bottom to top. The narrow width $r$ is the adapter rank. The figure uses a square weight matrix; the same idea works for rectangular matrices.

The small matrices are often called a LoRA adapter. They learn together with the task's training objective. LoRA does not first compute a full fine-tuned model and then compress the difference: the small matrices are the parameters it trains from the start.

What does "low rank" mean?🔗

The rank of a matrix measures how many independent directions it can express. In LoRA, the narrow intermediate width limits the number of directions available to the update. If that width is 8, the update matrix has rank at most 8, even when the original layer has thousands of inputs and outputs.

Crucially, the pretrained matrix is not forced to have low rank. It remains intact. The assumption is that adapting an already capable model to a particular task may need a much simpler change than learning the model in the first place.

The authors draw motivation from work on the low intrinsic dimensionality of fine-tuning: some tasks can be learned by optimizing within a much smaller parameter space. LoRA turns that motivation into a specific constraint on individual weight updates. It is an empirical hypothesis about adaptation, not a proof that every useful update must be low-rank. Paper, section 4.1; Aghajanyan et al..

The algebra, with the paper's notation

Let the frozen weight matrix be $W_0 \in \mathbb{R}^{d \times k}$. It maps an input $x$ with $k$ components to an output $h$ with $d$ components. LoRA represents the update as:

$$ \Delta W = BA, $$

where $A \in \mathbb{R}^{r \times k}$ and $B \in \mathbb{R}^{d \times r}$, with $r \ll \min(d,k)$. The product $BA$ has the same dimensions as $W_0$, but rank at most $r$.

Equation (3) of the paper writes the modified forward pass as:

$$ h = W_0x + \Delta Wx = W_0x + BAx. $$

Immediately afterward, the paper specifies that the update contribution is scaled by $\alpha/r$. Writing that scaling explicitly gives:

$$ h = W_0x + \frac{\alpha}{r}BAx. $$

Here, $\alpha$ controls the strength of the update relative to its rank. The paper holds $\alpha$ fixed when varying $r$ in this construction. Rank and scaling are related settings, but changing rank changes the space of possible updates, not just their size.

The paper initializes $A$ with random Gaussian values and $B$ with zeros. Consequently, $BA=0$ initially: the adapter starts by adding nothing to the original layer. Training then learns a nonzero adjustment. Paper, section 4.1.

How much smaller is the trainable part?🔗

For a layer with $k$ inputs and $d$ outputs, full fine-tuning updates $dk$ matrix entries. LoRA trains $rk+dr=r(k+d)$ entries instead, excluding biases or any separately trained task head.

Consider a 4,096-by-4,096 matrix and rank 8:

Trainable component	Parameter count
Full weight matrix	16,777,216
LoRA matrix A, shape 8 by 4,096	32,768
LoRA matrix B, shape 4,096 by 8	32,768
Both LoRA matrices	65,536

That is 256 times fewer trainable parameters for this matrix. The frozen 16.8 million weights are still present, and the model still computes with them. Total training memory also includes activations and temporary buffers, so it will not fall by the same factor.

A training step with frozen weights🔗

For our question-to-SQL example, training proceeds as follows:

Run a question through the model. Each adapted layer adds its LoRA adjustment to the original computation.
Compare the predicted output with the target SQL query and calculate a loss.
Backpropagate through the computation to determine how the adapter matrices should change.
Update the adapter parameters, leaving the pretrained weights fixed.

Frozen weights still participate in backpropagation. Their operations help determine the gradients needed by adapters earlier in the network. LoRA saves the work and storage associated with updating the frozen weights; it does not eliminate the backward pass through the model.

Where should the adapters go?🔗

Transformer attention uses projections called queries, keys, values, and an output projection. Queries and keys determine which positions attend to one another; values provide the information that attention combines.

The original LoRA study concentrates on attention weights, usually the query and value projections, while leaving the feed-forward networks frozen. LoRA itself is not restricted to those locations. The choice determines where the model can learn adjustments, and it matters alongside rank.

In the paper's GPT-3 experiment with an approximately 18-million-parameter budget, adapting both query and value matrices at rank 4 works better on WikiSQL than spending the same budget on query matrices alone at rank 8. This illustrates why distributing capacity across useful locations can matter more than increasing the rank at one location. Paper, sections 4.2 and 7.1.

Is a larger rank always better?🔗

No. A larger rank allows a more expressive update and increases the number of trainable parameters, but does not guarantee better validation performance. On the paper's GPT-3 WikiSQL and MultiNLI experiments, query-and-value adapters already work well at very low ranks, including rank 1.

That result is specific to those models and tasks. The authors explicitly caution that small ranks need not suffice for every dataset. Rank, placement, learning rate, and the training data should be evaluated together. A low-rank constraint can limit fitting capacity, but it does not by itself guarantee less overfitting. Paper, section 7.2.

What the paper's results establish🔗

The experiments cover language understanding and generation using RoBERTa, DeBERTa, GPT-2, and GPT-3. They show that LoRA can match or exceed full fine-tuning on a range of tested tasks while training far fewer parameters. They do not establish that LoRA is always the best method for every adaptation problem.

The GPT-3 175B results make the scale of the savings concrete:

Training memory: the paper reports a reduction from 1.2 TB to 350 GB in its setup.
Task-specific storage: with rank-4 query-and-value adapters, the stored update is about 35 MB instead of a roughly 350 GB full-model checkpoint. The base model is still required.
Training speed: the reported throughput rises from 32.5 to 43.1 tokens per second per V100 GPU, using the same number of model-parallel weight shards. That corresponds to about 25% less time for the same token workload, not a universal runtime guarantee.

The approximately 10,000-fold reduction, or about 0.01% of the full parameter count, refers to the rank-4 storage example. It is not a fixed fraction for every LoRA configuration. Paper, section 4.2.

A closer look at the GPT-3 task scores

The following selection comes from the paper's table 4:

Method	Trainable parameters	WikiSQL accuracy	MNLI-m accuracy	SAMSum ROUGE-L
Full fine-tuning	175,255.8 million	73.8%	89.5%	44.5
LoRA	4.7 million	73.4%	91.7%	45.9
LoRA	37.7 million	74.0%	91.6%	45.1

WikiSQL reports logical-form validation accuracy for question-to-SQL prediction, rather than execution accuracy. MNLI-m measures whether one sentence entails, contradicts, or is neutral with respect to another, on the matched validation set. SAMSum evaluates dialogue summaries; ROUGE-L measures overlap with reference summaries using the longest common subsequence.

These are different metrics, not interchangeable measurements of general intelligence. The smaller LoRA configuration is slightly below full fine-tuning on WikiSQL, within the paper's stated typical fluctuation of about 0.5 percentage points, while exceeding it on the other two reported metrics. The point is competitive task performance with much less trainable state, not a claim of a win in every table cell. Paper, table 4.

Serving models and switching tasks🔗

During training, LoRA has two computational paths. At deployment, their weights can be merged: add the learned update to the pretrained matrix once, then use the resulting matrix as an ordinary dense layer.

With the paper's scaling made explicit, the merged matrix is:

$$ W_{\mathrm{merged}} = W_0 + \frac{\alpha}{r}BA. $$

This is the basis of the paper's no additional inference latency claim. After merging, there is no separate adapter computation at that layer. Keeping adapters separate can make task switching convenient, but retains the extra operations.

Different tasks can share one base model and store separate small adapters. To switch a merged model between tasks, the paper describes subtracting the first update and adding the second. In practice, keeping the original base checkpoint also provides a clean starting point for each merge. Once weights are merged for one task, serving different adapters together in a batch requires additional handling; the original paper identifies this as a limitation. Paper, sections 4.1 and 4.2.

Does freezing prevent catastrophic forgetting?🔗

Freezing preserves the original parameter values, so disabling the adapters recovers the base model's computation when no other parameters have been changed. That is useful when maintaining several task-specific versions.

It does not guarantee that the adapted model retains every original capability. Its outputs depend on the adapters as well as the base weights, and the learned adjustment can harm performance on other tasks. Retention therefore needs evaluation alongside the new task's score. The original paper does not establish a general guarantee against catastrophic forgetting.

How LoRA fits among PEFT methods🔗

The taxonomy below places LoRA in a broader family of approaches:

Additive methods introduce trainable components, such as nonlinear adapter modules or learned soft prompts.
Selective methods update a subset of existing parameters, such as biases.
Reparameterization-based methods, including LoRA, express a trainable change through a more economical parameterization.

Figure 2. Taxonomy reproduced from the 2023 survey Scaling Down to Scale Up. The overlapping regions represent methods that combine ideas. This is a survey figure, not a figure from the original LoRA paper.

The word adapter can cause confusion here. A LoRA adapter is a pair of matrices representing a weight update. The classic bottleneck adapters compared in the original paper insert additional layers with a nonlinearity. LoRA's linear update can be merged into the original matrix; those nonlinear modules generally cannot be absorbed in the same way.

From LoRA to QLoRA🔗

LoRA reduces the part of the model that must be trained. It leaves another large expense: storing the frozen base weights. QLoRA addresses that expense by storing the base model in four-bit form while training higher-precision LoRA adapters.

The distinction is simple: LoRA is a method for learning a compact update; QLoRA combines that method with quantization and additional memory-saving techniques. QLoRA is not merely LoRA with a smaller rank, and its placement experiments extend beyond the original LoRA paper's attention-focused study.

Continue with QLoRA: How to fine-tune large language models with less memory for NormalFloat, double quantization, paged optimizers, and the single-GPU results.

The takeaway🔗

LoRA does not make a large pretrained model small. It makes the task-specific change small. Two trainable matrices can provide a useful adjustment while the original model remains shared and fixed. The practical gains are fewer gradients and optimizer states, smaller task checkpoints, and ordinary inference computation when the update is merged.

Its success depends on whether the chosen rank and adapter locations give the model enough freedom to learn the task. That is the trade-off to test: task quality against the memory, storage, and training cost of the update.

Sources and implementation🔗

Hu et al., LoRA: Low-Rank Adaptation of Large Language Models, first submitted in 2021 and published at ICLR 2022. This article follows arXiv version 2.
Microsoft's official LoRA repository and loralib, including examples for adding LoRA layers, saving adapter weights, and merging for inference. Its default setup trains only LoRA parameters; training biases is an explicit option.
Aghajanyan et al., Intrinsic Dimensionality Explains the Effectiveness of Language Model Fine-Tuning, the motivation cited by the LoRA authors.
Lialin et al., Scaling Down to Scale Up: A Guide to Parameter-Efficient Fine-Tuning, the source of the broader PEFT taxonomy.

Michał Chromiak's blog - Optimization

QLoRA: How to fine-tune large language models with less memory

Why fine-tuning needs so much memory🔗

LoRA: Learn a small update instead of rewriting the model🔗

What QLoRA adds🔗

1. NormalFloat: Store the base weights in 4 bits🔗

2. Double quantization: Compress the scales too🔗

3. Paged optimizers: Make room for temporary memory peaks🔗

One training step, from input to update🔗

What the experiments establish🔗

Guanaco and the single-GPU results🔗

What "99.3% of ChatGPT" means🔗

A closer look: Design choices and memory costs🔗

The takeaway🔗

Related reading🔗

Original sources🔗

CoLT5: Reading long documents with selective computation

Why efficient attention is only part of the problem🔗

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

Routing: What does "important" mean?🔗

Feed-forward branches: Wider is not deeper🔗

Attention branches: Separate who asks from who supplies information🔗

Two other changes complete the recipe🔗

Multi-query cross-attention speeds up the decoder🔗

UL2 supports learning from examples in the input🔗

What the experiments show🔗

Quality and speed at roughly 16k input length🔗

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

What the ablations teach us🔗

Where the approach is useful🔗

Sources and implementation notes🔗

The takeaway🔗

LoRA: Fine-tuning a model by learning a small update

Why train an update instead of the whole model?🔗

Two paths through a layer🔗

What does "low rank" mean?🔗

How much smaller is the trainable part?🔗

A training step with frozen weights🔗

Where should the adapters go?🔗

Is a larger rank always better?🔗

What the paper's results establish🔗

Serving models and switching tasks🔗

Does freezing prevent catastrophic forgetting?🔗

How LoRA fits among PEFT methods🔗

From LoRA to QLoRA🔗

The takeaway🔗

Sources and implementation🔗