LoRA for Long-Context Adaptation

A 7B model trained on 4 096-token sequences cannot simply be asked to process 32 000 tokens: the rotary position embeddings (RoPE) were never seen at angles beyond the training range, and the attention matrices grow quadratically with sequence length. Fine-tuning the full model to fix this costs tens of thousands of GPU-hours. LongLoRA (Chen et al., 2023) demonstrated that a carefully structured LoRA approach can extend Llama-2 7B to 100 000 tokens on a single 8×A100 node - a gap of roughly 25× - but only after identifying two non-obvious failure modes that vanilla LoRA misses entirely.

Why vanilla LoRA under-delivers on long context

Standard LoRA freezes all pretrained weights and learns two small matrices, \(A \in \mathbb{R}^{d \times r}\) and \(B \in \mathbb{R}^{r \times d}\), added to each target projection:

The rank \(r \ll d\) keeps parameter count low. For task adaptation this works well because the base model already "knows" the task distribution; LoRA nudges the residual.

Long-context adaptation is different. The model has never encountered position indices beyond the training cutoff. RoPE encodes relative distance by rotating query and key vectors by angles \(\theta_i \cdot m\) where \(m\) is the position and \(\theta_i\) decreases geometrically with head dimension \(i\). At positions far outside the training range, those cosine/sine values enter regions the model was never optimised over. Low-rank updates to attention projections cannot compensate for this because the error is in the positional geometry, not the weight subspace.

Two independent but complementary fixes are therefore needed:

1. Position rescaling - bring out-of-range position indices back into a familiar angular neighbourhood.
2. Efficient attention during training - make the quadratic attention cost tractable so the model actually sees long sequences.

Position interpolation: the prerequisite

Before applying any LoRA, the position indices must be remapped. The simplest approach, proposed by Chen et al. (2023, "Extending Context Window via Positional Interpolation"), is linear interpolation: divide every position index by the ratio \(s = L_\text{new} / L_\text{orig}\), so the maximum index seen during fine-tuning maps to \(L_\text{orig}\) rather than \(L_\text{new}\).

This compresses previously unseen positions back into the angle range the model trained on. The cost is a small blurring of nearby-token distinctions - positions that were once distinct now share a tighter angular neighbourhood. A brief LoRA fine-tuning (roughly 1 000 steps) corrects for this blurring.

Without the interpolation step, LoRA fine-tuning on long sequences shows near-random perplexity on the extended portion of the context even after thousands of gradient steps - the model cannot learn to interpret angles it has never seen via low-rank residuals alone.

Shifted sparse attention: making training feasible

Full attention over sequences of 32 000+ tokens is quadratic in memory and compute. LongLoRA replaces full attention during training only with shifted sparse attention (S2-Attn):

• Split the sequence into fixed-size groups (e.g., 2 048 tokens each).
• Attend fully within each group (cheap, parallelisable).
• On alternate heads, shift all tokens by half a group size before grouping, so each shift boundary is covered by at least one head.

Group A: [tok 0 ... tok 2047]   (head 0, 2, 4, ...)
Group B: [tok 0 ... tok 2047]   (head 1, 3, 5, ...)
  shifted by 1024 → effectively [tok 1024 ... tok 3071]

This achieves approximate full-context coverage with linear memory per group. Crucially, inference reverts to standard dense attention - the sparse pattern is a training-only approximation, so no custom kernel is needed at serving time.

Which components to adapt

LongLoRA's key empirical finding is that adapting only attention projections (\(Q\), \(K\), \(V\), \(O\)) with LoRA is insufficient for long-context extension. Two additional components must remain trainable (full-rank, not LoRA):

Freezing normalisation layers during long-context LoRA reliably raises perplexity by 1-3 points on long benchmarks. The intuition: layer normalisation statistics shift when the mean and variance are computed over much longer sequences; the learned scale and bias need to adjust.

A training recipe

A practical pipeline for extending a 7B RoPE model from 4 096 to 32 768 tokens with LoRA:

1. Rescale RoPE base: multiply the base frequency \(\theta\) by the extension ratio (e.g., \(\theta' = \theta \times 8\) for 8× extension), or apply linear position interpolation by adjusting the position index divisor. Both are single-line config changes in most frameworks.

2. Enable S2-Attn during training: patch the attention implementation to use shifted groups. In practice this is two lines of index manipulation before the attention call.

3. Configure LoRA targets: attention projections with rank 64-128 (higher rank than typical task adaptation because the positional geometry shift demands more expressive residuals).

4. Unlock embeddings and norms: set requires_grad=True for the embedding table and all normalisation parameters. These add fewer than 0.1% of total parameters but are critical.

5. Train on long sequences: the dataset must contain actual long documents. Concatenating short documents with separator tokens does not reproduce genuine long-range dependencies.

A rough compute estimate for this recipe on a 7B model extending to 32 k tokens: around 1 000 gradient steps on 8×A100 80 GB cards, 1-2 days of wall time - versus weeks for full fine-tuning.

When it falls down

Inadequate training data length. If the fine-tuning corpus does not contain documents genuinely longer than the target context, the model learns to interpolate positions but has no signal for what long-range dependency patterns look like. Perplexity on short sequences stays good; retrieval benchmarks (e.g., "passkey" tasks) degrade.

Rank too low for large extensions. Extending from 4 096 to 8 192 tokens may work with rank 16. Extending to 100 000 tokens requires rank 64-128 or higher. The positional geometry shift is larger, and a low-rank subspace cannot span it. A common failure symptom is good perplexity on medium lengths but collapse on the very long tail.

Catastrophic forgetting on short context. Over-training (too many steps, too high a learning rate) on long documents causes the model to degrade on its original 4 096-token tasks. LoRA's parameter isolation protects less than expected here because the embeddings and norms are full-rank trainable. Monitoring short-context perplexity during training and early-stopping accordingly is essential.

RoPE base scaling vs. interpolation mismatch. There are two common position rescaling strategies (base frequency scaling and direct position interpolation), and they interact differently with LoRA rank and fine-tuning duration. Mixing them inconsistently, or applying one during fine-tuning and another at inference, produces erratic attention patterns at long ranges.

Attention sink concentration. In very long sequences, a small number of "sink" tokens (typically the BOS token) attract disproportionate attention. Standard LoRA fine-tuning can amplify this rather than correct it, particularly when rank is low. The symptom is high accuracy on retrieval tasks where the answer is near the start of the document, but poor accuracy when the answer is in the middle.

Flash Attention compatibility. S2-Attn requires specific sequence manipulations that are not always compatible with FlashAttention kernels out of the box. Without FlashAttention, the quadratic memory cost makes >32 k sequences impractical even with S2-Attn. Ensuring the custom attention patch integrates with the FlashAttention 2 API is a non-trivial engineering step.

Further reading

• Yukang Chen et al., "LongLoRA: Efficient Fine-tuning of Long-Context Large Language Models" (2023) - https://arxiv.org/abs/2309.12307
• Shouyuan Chen et al., "Extending Context Window of Large Language Models via Positional Interpolation" (2023) - https://arxiv.org/abs/2306.15595
• Edward Hu et al., "LoRA: Low-Rank Adaptation of Large Language Models" (2021) - https://arxiv.org/abs/2106.09685
• Jianlin Su et al., "RoFormer: Enhanced Transformer with Rotary Position Embedding" (2021) - https://arxiv.org/abs/2104.09864