PaperlayerView exact revision on arXiv

Source-linked AI summary

Focused Transformer: Contrastive Training for Context Scaling

Szymon Tworkowski, Konrad Staniszewski, Mikołaj Pacek, Yuhuai Wu, Henryk Michalewski, Piotr Miłoś

arXiv:2307.03170v2cs.CLcs.AIcs.LG

TL;DR

Long-context use is limited because additional documents increase irrelevant keys and overlapping key representations, creating a distraction issue. FOT addresses this with differentiable contrastive-style crossbatch training and memory attention, enabling fine-tuned LONGLLAMA models to extend context. The models improve long-context performance and handle 256k-token passkey retrieval.

  • Problem

    As the number of documents grows, relevant keys become rarer and overlapping keys associated with different values make long-context attention difficult to distinguish.

  • Method

    FOT combines memory attention with differentiable crossbatch training that uses same-document contexts as positives and other-document contexts as negatives to structure key-value representations.

  • Results

    LONGLLAMA models fine-tuned from OpenLLaMA checkpoints improve long-context tasks and successfully handle a 256k context length on passkey retrieval.

  • Takeaways & Limitations

    FOT provides a way to extend the effective context of existing large language models through fine-tuning without requiring long context during training.

  • Takeaways & Limitations

    Scaling beyond the experiments requires distributed storage and training, while exact kNN search is not scalable to large memory and approximate search needs further evaluation.

Abstract

from arXiv · show

Large language models have an exceptional capability to incorporate new information in a contextual manner. However, the full potential of such an approach is often restrained due to a limitation in the effective context length. One solution to this issue is to endow an attention layer with access to an external memory, which comprises of (key, value) pairs. Yet, as the number of documents increases, the proportion of relevant keys to irrelevant ones decreases, leading the model to focus more on the irrelevant keys. We identify a significant challenge, dubbed the distraction issue, where keys linked to different semantic values might overlap, making them hard to distinguish. To tackle this problem, we introduce the Focused Transformer (FoT), a technique that employs a training process inspired by contrastive learning. This novel approach enhances the structure of the (key, value) space, enabling an extension of the context length. Our method allows for fine-tuning pre-existing, large-scale models to lengthen their effective context. This is demonstrated by our fine-tuning of $3B$ and $7B$ OpenLLaMA checkpoints. The resulting models, which we name LongLLaMA, exhibit advancements in tasks requiring a long context. We further illustrate that our LongLLaMA models adeptly manage a $256 k$ context length for passkey retrieval.

1 Introduction

The paper identifies distraction as a central obstacle to extending Transformer context across multiple documents and introduces Focused Transformer training to address it. Fine-tuned LONGLLAMA models extend context without long-context training and support long-context tasks, including 256k-token passkey retrieval.

  • Results: 94.5% accuracy at 100k and 73% at 256k are reported for LONGLLAMA 3B on passkey retrieval.The original OpenLLaMA baseline cannot handle contexts longer than its 2k training length.
  • Motivation: The distraction issue arises when increasing documents reduce the proportion of relevant keys and overlapping keys make relevant information harder to distinguish.This is presented as a primary obstacle to scaling context in multi-document settings.
  • Method: FOT uses a contrastive-learning-inspired objective to improve the structure of key-value representations and enable k-nearest-neighbor context extension.The method provides additional context through memory attention and kNN lookup.
  • Method: FOT can extend context through cost-effective fine-tuning of existing 3B and 7B OpenLLaMA checkpoints without modifying the architecture.The resulting models are named LONGLLAMAs.
  • Results: 256k context length is successfully handled by LONGLLAMA models on passkey retrieval.The paper reports this as a demonstrated long-context capability.
  • Results: FOT shows improvements on tasks benefiting from more few-shot demonstrations and on long-context benchmarks including dictionary lookup, language modeling, and formal proofs.A model trained with 512-token context extrapolates to 16 million tokens in dictionary lookup.

2 Related work

Prior work extends Transformer context through architectural sparsity, recurrence, retrieval, or representation learning. FOT instead frames memory access as context extension and uses cross-document negatives to shape representations for that purpose.

  • Long-context architectures: Transformer-XL, Longformer, BigBird, and LongT5 primarily address long contexts through caching, sparse attention, or reduced attention complexity.These approaches differ in how they reduce the computational burden of long-context processing.
  • Fine-tuning and retrieval: RETRO and Memorizing Transformer fine-tune existing language models to add capabilities, whereas FOT is framed as context extension rather than retrieval.The paper contrasts FOT’s single-stage context extension with RETRO’s two-stage approach.
  • Zero-shot methods: KNN-LM combines model probabilities with retrieved-token similarity, while FOT extends context within selected attention layers to permit reasoning over additional context.The distinction is between augmenting output probabilities and expanding attention context.
  • Contrastive learning: Contrastive learning compares positive and negative examples, and prior work links larger negative-example batches with improved representation quality.FOT applies this general principle to memory-augmented language-model training.
  • FOT relation to prior work: FOT trains current-context tokens against same-document positives and other-document negatives through differentiable crossbatch attention.The negative contexts are intended to better shape the key-value space.

3 FOT: Focused Transformer

FOT combines memory attention with differentiable crossbatch training so selected layers can retrieve relevant key-value pairs from additional context. The training procedure targets distraction by contrasting same-document positives with unrelated-document negatives.

  • Architecture: FOT is a plug-and-play extension that uses memory attention and crossbatch training to support longer-context training or fine-tuning.The method is designed to improve key-value representations without changing the model architecture.
  • Memory attention: Memory attention retrieves the top k matching preceding keys from an additional memory using exact kNN search, extending inference context.Queries attend to local preceding keys and memory keys ranked by inner product.
  • Crossbatch training: Crossbatch training exposes each current context to same-document previous context as positive examples and d −1 other-document contexts as negatives.The procedure is differentiable and operates through the data pipeline and a minor self-attention change.
  • The distraction issue: The distraction issue is that standard training does not incentivize distinguishing keys from different documents, causing attention to spread across related and unrelated contexts.The paper describes this as a key problem for long multi-document contexts.
  • Evaluation: Figure 3 compares FOT variants and a standard Transformer by plotting the number of other-document contexts against average attention mass on the current document’s previous local context.The comparison tests how parameter d relates to attention allocation.
  • The distraction issue: The positive attention mass is observed near 1/d, indicating that attention is approximately equally distracted by positive and negative keys.The quantity measures attention assigned to the current document’s positive context.

: extending LLaMA’s context length with FOT

FOT fine-tuning extends existing OpenLLaMA models beyond their training context while retaining short-context performance. LONGLLAMA models show long-context gains, including strong passkey retrieval at 256k tokens.

  • Model extension: FOT fine-tunes existing OpenLLaMA-3B and OpenLLaMA-7B models into LONGLLAMAs that extrapolate beyond their training context length.The models retain compatibility with existing LLaMA inference code.
  • Long-context retrieval: 73% accuracy at 256k tokens demonstrates the 3B LONGLLAMA’s passkey retrieval capability beyond its 8K training context.The model achieves 94.5% accuracy at 100k-token prompts.
  • Few-shot learning: LONGLLAMA gains accuracy from longer few-shot contexts on TREC and shows some improvements on WebQS.The evaluation uses in-context demonstrations for both tasks.
  • Extrapolation beyond training: FOT achieves further accuracy gains at 6K and 8K evaluation contexts after fine-tuning with 4K context, unlike the standard fine-tuning baseline.Both models were fine-tuned on 1B tokens with a 4K context length.
  • Short-context compatibility: LONGLLAMA maintains performance on most original short-context tasks, supporting use as a drop-in LLaMA replacement.The evaluation addresses whether longer-context fine-tuning harms performance at the original 2K context length.

5 Analysis of FOT

Experiments analyze FOT across context scaling, model extension, distraction handling, and long-context language modeling. FOT extrapolates far beyond its training context, uses very large memories, and improves perplexity while handling multi-document distractions.

  • Evaluation settings: FOT experiments evaluate single-document and multi-document settings, with the latter retaining memory across documents to test distraction handling.Single-document evaluation clears memory for each document, whereas multi-document evaluation does not.
  • Context extrapolation: FOT fine-tuning enables steady perplexity gains up to 64K tokens despite only 2K total differentiable context during training.The comparison uses Transformer-XL and Memorizing Transformer baselines fine-tuned from the same base checkpoint.
  • Dictionary lookup: Above 92% accuracy after 5k training steps shows that FOT can use a 16M-token memory on the synthetic dictionary lookup task.The models were trained with documents of length 512.
  • Distraction handling: Higher crossbatch dimension d improves perplexity, consistent with benefits from mitigating the distraction issue.The analysis varies multi-document memory size on PG-19 language modeling.
  • Multi-document language modeling: FOT uses memory substantially better than a standard Transformer, while perplexity rises only 0.18 when memory scales beyond 500k tokens.Its perplexity approaches the single-document Memorizing Transformer comparison treated as a soft lower bound.
  • Single-document extrapolation: FOT improves perplexity as context grows far beyond training length, and larger crossbatch dimension d performs better in the tested settings.The analysis also varies the number of previous contexts used during crossbatch training.
  • Ablations: FOT’s ablations identify differentiability and inclusion of negatives as important design choices.The appendix also discusses the relation to Memorizing Transformer and memory-layer selection.

6 Limitations and future work

The paper identifies engineering, training-cost, and methodological limitations that constrain scaling FOT, while proposing broader contrastive-learning and hybrid-method directions.

  • Scaling up context: Storing more than 16M key-value pairs requires a distributed multi-node system, and exact kNN search is not scalable to large memory.Approximate kNN would require substantial engineering and careful evaluation of its effect on model performance.
  • Scaling up crossbatch: Increasing crossbatch dimensionality is beneficial, but experiments are limited to d = 64 or d = 128 by single-TPU memory.The authors also plan larger-memory or multi-node training and note that crossbatch increases training cost in only a subset of layers.
  • Exploring contrastive learning: FOT uses basic contrastive-learning techniques, leaving hard-negative mining and other contrastive methods for future work.The authors report that the current approach improves key structure and mitigates distraction.
  • Combining with other methods: The authors suggest combining FOT with other long-context methods because their interactions may be mutually beneficial.This is presented as a future direction rather than an evaluated result.

Broader Impact

The broader-impact discussion states that FOT is a generic capability-extending technique and could amplify existing language-model threats, while asserting that it introduces no new threats.

  • Broader Impact: FOT is described as a generic technique applicable in principle to virtually any language model.The paper frames this generality as a capability-extension property.
  • Broader Impact: Extending model capabilities could exacerbate existing threats, but the authors state that FOT does not create new threats.The paper refers readers to existing broader-impact work on language models.
  • Model background: OpenLLaMA uses a decoder-only architecture with rotary positional embeddings, RMSNorm pre-normalization, SiLU activation, and a 32k SentencePiece vocabulary.These implementation details describe the model family used for the paper’s LongLLaMA experiments.
  • Implementation: For LLaMA-based models, local positional encodings are retained while tokens outside the local context receive the first local token’s position.This design supports backward compatibility with the original LLaMA.
  • Fine-tuning dataset: The fine-tuning mixture combines RedPajama and The Stack, with Python files concatenated by subdirectory to create longer training documents.Short documents are filtered for some subsets, and documents too short for crossbatch are concatenated with the next document.

A.4 Language Model Evaluation Harness

The evaluation section checks whether LongLLaMA preserves short-context performance and benefits from added context, while describing memory-attention implementation choices and alternatives.

  • Language Model Evaluation Harness: LongLLaMAs maintain good performance on short-context tasks when evaluated as standard OpenLLaMA models without context extension.The Language Model Evaluation Harness comparison omits CB and WSC, matching the OpenLLaMA evaluation setup.
  • Context utilization: LongLLaMA 3B benefits from extended context on the Qasper validation set.The evaluation uses zero-shot performance at different context lengths and separates yes/no from open questions.
  • Memory attention: Memory attention retrieves the k highest-inner-product key-value pairs from additional memory and merges them with the preceding local context.The merged entries are used to compute the attention value with the standard Transformer formula.
  • Memory attention: The softmax score can use a temperature parameter, while the described approach does not distinguish between local context and memory.The supplied passages define the score and temperature but do not reproduce the complete formula.
  • Memory integration: A gating alternative computes separate memory and local-context attention values before combining them with a sigmoid-controlled gate.The authors report that their simpler direct-attention approach is equally effective and requires no additional parameters, so they use it in experiments.
  • LongLLaMA implementation: For LongLLaMA, the implementation replaces kNN search with dense attention, retains positional encodings, and omits the temperature parameter.The temperature is omitted because the architecture does not normalize keys and queries.

B.3 Crossbatch training procedure

Crossbatch training exposes selected attention layers to previous-context key-value pairs from the same and other documents, using differentiable representations and negatives to reduce distraction. These choices improve long-context perplexity and retrieval behavior relative to the Memorizing Transformer in the reported experiments.

  • Training procedure: Selected attention layers mix key-value pairs from the current document’s previous context with contexts from other documents during training.For each batch document, the additional contexts are fixed, and positional encoding is disabled in the augmented layer.
  • Training procedure: FOT backpropagates through previous-context key-value pairs, unlike the Memorizing Transformer’s non-differentiable retrieved memory.The authors identify differentiability as important for jointly structuring key, value, and query embeddings.
  • Comparison with Memorizing Transformer: FOT’s crossbatch training differs from Memorizing Transformer training because it attends to previous contexts from the same and different documents rather than only same-document retrieved tokens.The section also notes that FOT does not use memory during training and uses up-to-date values.
  • Comparison with Memorizing Transformer: The authors leave blending FOT and Memorizing Transformer for future work because Memorizing Transformer may provide harder negatives.They report only a proof-of-concept experiment combining the approaches.
  • Empirical findings: FOT is significantly better than Memorizing Transformer when context is expanded during inference in the multi-document experiment.Both models were trained with a local context of 512; the comparison attributes the advantage to differentiable keys and values.
  • Empirical findings: Increasing the number of negatives increases attention mass on relevant tokens, while training without negatives causes significant perplexity deterioration as context length grows.The comparison uses d = 1 versus a schedule that reaches d = 64, alongside a Memorizing Transformer baseline with 16K memory.

F Dictionary lookup task

The dictionary lookup task tests whether FOT can retrieve values from a growing key-value memory after training on short documents. Evaluation increases the number of definitions while keeping the final query block fixed, and FOT maintains high accuracy at large memory sizes.

  • Task design: The task presents key-value definitions followed by queries for randomly selected keys, with accuracy measured on the predicted values after the query marker.The dictionary uses special key, value, and query tokens, and the loss is masked to the queried values.
  • Results: FOT demonstrates high accuracy even when the dictionary memory is large.Figure 10 plots accuracy against the number of dictionary tokens after training on 512-token examples with an extended evaluation context.
  • Training and evaluation: Training documents contain 512 tokens, with definitions in the first half and questions in the second half; FOT uses a local context of 256.FOT increases d from 1 to 128 after reaching 98% training accuracy and retrieves k = 32 keys during inference.
  • Training and evaluation: Evaluation lengthens documents by adding definitions while keeping only the last 256 tokens as questions.Thus, the number of definitions grows with the token axis while the number of queries remains unchanged.

H Datasets

The study evaluates FOT on text, code, and formal-proof corpora, including PG-19, arXiv, GitHub, and Isabelle. The datasets vary from books and scientific papers to repositories and mathematical theories, with repeated runs used for selected experiments.

  • Datasets: PG-19 and arXiv are among the principal datasets used in the study, while additional experiments use GitHub repositories and Isabelle formal theories.The GitHub corpus covers several programming languages, and Isabelle combines the Archive of Formal Proofs with the Isabelle standard library.
  • Datasets: GitHub data consists of permissively licensed source-code files concatenated within repositories while preserving subdirectory structure before shuffling files and subdirectories.The selected extensions include C, C++, Java, Python, Go, and TypeScript.
  • Datasets: Isabelle data consists of open-source formal mathematical theories covering topics such as logic, analysis, algebra, and cryptography.Each theory contains multiple files of formal proofs.
  • Evaluation protocol: Figure 10 reports minimum and maximum values over 10 runs, while some other experiments use fewer repeated runs because of resource constraints.The arXiv baseline experiment in Appendix K uses three seeds and reports a standard deviation of 0.002 perplexity.
  • Evaluation protocol: Test perplexities are calculated using 1M tokens.

K Additional experimental results

Additional experiments compare FOT with Memorizing Transformer across training and evaluation context lengths on arXiv. FOT benefits from trainable keys and values and from using negatives, while the baseline shows limited gains when its training memory becomes very large.

  • Experimental setup: The additional experiments train both models for 500k steps with a 2K local context and evaluate them at different context lengths on arXiv.The setup follows the single-document protocol used in prior Memorizing Transformer work.
  • Memorizing Transformer results: A Memorizing Transformer trained with 2K memory struggles beyond 32K evaluation tokens, whereas 8K memory performs better at longer contexts and 16K yields diminishing returns.
  • FOT results: FOT outperforms the Memorizing Transformer in an apples-to-apples comparison with 2K additional training context, supporting the importance of trainable keys and values.
  • FOT results: 2.148 perplexity is achieved by FOT with 4K additional training context, compared with 2.164 for Memorizing Transformer with 16K additional context.
  • FOT results: d = 2 works significantly better than d = 1 across the reported FOT settings.The comparison varies d and w while evaluating at different context lengths.
  • Implementation: The memory-attention implementation combines local and additional attention scores before softmax, using previous-context pairs during training and nearest-neighbor memory pairs during inference.The same layer returns a context vector after weighting the corresponding values.
Loading 2307.03170v2…