PaperlayerView exact revision on arXiv

Source-linked AI summary

Generative Representational Instruction Tuning

Niklas Muennighoff, Hongjin Su, Liang Wang, Nan Yang, Furu Wei, Tao Yu, Amanpreet Singh, Douwe Kiela

arXiv:2402.09906v3cs.CLcs.AIcs.LG

TL;DR

Existing language models generally perform well at either generation or embedding, leaving a gap for a single model that handles both. GRIT trains one language model to distinguish these task types through instructions and separate objectives. The resulting models achieve strong cross-task performance, including state-of-the-art open-model MTEB results and more than 60% faster long-document RAG inference.

  • Problem

    Current models tend to perform well at generation or embedding rather than both, despite embeddings supporting important applications such as search and chatbots.

  • Method

    GRIT combines generative and representational instruction tuning, using instructions and separate loss functions to distinguish the two task streams in one model.

  • Results

    GRITLM 7B achieves state-of-the-art MTEB performance among open models while outperforming generative models up to its size, and GRITLM 8X7B is reported as the best open generative model tried while retaining strong embeddings.

  • Takeaways & Limitations

    The unified model matches embedding-only and generative-only variants without performance loss and enables more than 60% faster RAG inference for long documents.

  • Takeaways & Limitations

    Embedding performance for input sequences longer than 512 tokens is unknown because suitable benchmarks are lacking.

Abstract

from arXiv · show

All text-based language problems can be reduced to either generation or embedding. Current models only perform well at one or the other. We introduce generative representational instruction tuning (GRIT) whereby a large language model is trained to handle both generative and embedding tasks by distinguishing between them through instructions. Compared to other open models, our resulting GritLM 7B sets a new state of the art on the Massive Text Embedding Benchmark (MTEB) and outperforms all models up to its size on a range of generative tasks. By scaling up further, GritLM 8x7B outperforms all open generative language models that we tried while still being among the best embedding models. Notably, we find that GRIT matches training on only generative or embedding data, thus we can unify both at no performance loss. Among other benefits, the unification via GRIT speeds up Retrieval-Augmented Generation (RAG) by > 60% for long documents, by no longer requiring separate retrieval and generation models. Models, code, etc. are freely available at https://github.com/ContextualAI/gritlm.

1 Introduction

GRIT unifies embedding and generative instruction tuning in one model, addressing the poor embedding performance of current generative models. The authors report strong performance across both task types, faster long-document RAG, and simpler serving, while requiring additional finetuning compute.

  • 1 Introduction: Embedding tasks are important in applications such as search and chatbots, but using hidden states from current generative models produces poor embeddings.Generating embedding tensors directly is impractical because embeddings require high dimensionality and precision.
  • 1 Introduction: GRIT unifies generative and representational instruction tuning so one model handles both generation and embedding tasks.Instructions and separate loss functions distinguish the two training streams.
  • 1 Introduction: GRITLM 7B reaches state-of-the-art performance among open models on MTEB while outperforming larger generative models on generative tasks.The paper presents this as a unified model matching or exceeding specialized alternatives.
  • 1 Introduction: > 60% faster RAG inference is achieved for long documents by using one equivalent embedding-and-generation model and halving forward passes through cached computation.The unified model replaces separate retrieval and generative models in the described pipeline.
  • 1 Introduction: GRIT requires more finetuning compute than generative instruction tuning because it optimizes two objectives.The authors argue this cost is outweighed by avoiding separate pretraining and finetuning when both models are needed.
  • 1 Introduction: Additional contributions include bidirectional-attention mean pooling and same-dataset in-batch negatives for embeddings, plus mixed sample- and token-level loss aggregation for generation.The paper also proposes methods to reduce memory requirements during embedding-model training.

2 GRIT

GRIT trains a pretrained language model on embedding and generative instruction data in a consistent format, using task-specific attention, pooling, and loss objectives. Its formulation supports separate embedding and generative batch sizes while combining both streams through weighted losses.

  • 2 GRIT: Embedding inputs use bidirectional attention and mean pooling to produce a final representation from the hidden states.The representation loss uses a contrastive objective with in-batch negatives.
  • 2 GRIT: Generative inputs use causal attention and a language-modeling head that predicts the next tokens.The generative loss is computed over predicted response tokens.
  • 2 GRIT: Mean pooling averages hidden states across the input sequence while excluding instruction and format tokens from the average.Those tokens can still influence the representation through self-attention.
  • 2 GRIT: Generative loss aggregation can operate at the sample or token level, trading equal sample weighting against equal token weighting.Sample-level aggregation can bias the model toward short generations.
  • 2 GRIT: GRIT combines representation and generation losses as weighted objectives while allowing different numbers of embedding samples and generative samples or tokens.This permits increasing the embedding batch size while keeping the generative batch size fixed.

3 Experiments

The experiments evaluate GRITLM across embedding, generative, retrieval, alignment, and ablation settings. GRITLM 7B achieves leading performance on both embedding and generation, while unifying the objectives without performance loss, with trade-offs from compute and configuration choices.

  • Main results: GRITLM 7B outperforms prior open models on MTEB while outperforming generative models up to 7B parameters.GRITLM nearly doubles Llama 70B’s MTEB score of 35.6 and outperforms Llama 70B on generative tasks by more than 20%.
  • Main results: GRITLM matches embedding-only and generative-only variants, unifying both objectives without performance loss.The unified model uses more compute per training step because it processes both embedding and generative data.
  • Main results: Few-shot examples overall worsen embedding performance, with only small, marginal, and inconsistent gains on some PairClassification tasks.Including few-shot embedding samples in around 5% of MEDI2 training data did not make the model use such examples effectively.
  • Main results: KTO alignment improves average generative performance but decreases the MTEB average from 66.8 to 66.7 and from 65.7 to 65.2.AlpacaEval increases by more than 10 points for both KTO-tuned models, while no further embedding training is performed during alignment.
  • Ablations: Increasing the embedding batch size to 4096 produces a 1.0 gain on the embedding average while generative performance remains stable.The increase especially benefits the 15 retrieval datasets included in the embedding average.

4 RAG with GRIT

GRITLM simplifies RAG by combining embedding and generation, enabling cached representations and key-value states to reduce inference cost. Caching is most beneficial for long queries or documents, but introduces performance and storage trade-offs.

  • Caching variants: GRITLM unifies the embedding and generative models, reducing traditional RAG’s four forward passes through query and context processing.Query caching reuses query key-value states, while query-document or document-query caching also reuses cached document states.
  • Performance: RAG outperforms the “No RAG” baseline, while Query Caching reaches a match score of 25.46 despite a performance drop from traditional RAG.The drop is attributed to attention and formatting mismatches between embedding and generation.
  • Limitations: Query-document and document-query caching can approach the “No RAG” baseline, limiting their practical usefulness in the reported experiments.The authors attribute this behavior to an additional attention mismatch and leave RAG-specific finetuning for future work.
  • Latency: 54% and 33% faster CPU and GPU latency, respectively, make Query Caching advantageous for 4000-token queries.For 4000-token documents, Doc Caching is 63% and 31% faster on CPUs and GPUs, respectively; caching only 250 tokens provides negligible speed-up.
  • Latency: Caching mode should follow input length: Doc Caching for very long documents and Query Caching for very long queries.A production system could switch modes using a simple input-length check.
  • Storage: Doc Caching removes the need to store document text but requires around 30TB of key-value states for 2,681,468 documents with GRITLM 7B.These states can be offloaded to disk; loading one sample requires around 12.5MB in memory.

5 Discussion

The discussion presents GRITLM as a unified text embedding and generation model that maintains strong performance across both tasks. It also identifies remaining opportunities in multilinguality, multimodality, and RAG-specific optimization.

  • Open directions: Multilingual performance may improve through language-targeted data and architecture changes, while unifying multimodal embedding and generation remains unexplored.The discussion specifically identifies image-text and speech tasks as outside the current text-only unification scope.
  • Unification: GRITLM unifies representational instruction tuning and generative instruction tuning in one model at no performance loss on either task.On MEDI2, adding the generative objective even improves embedding performance over an otherwise equivalent embedding-only model.
  • Unification: The authors interpret embedding and generative language modeling as related expressions of deep natural-language understanding.They hypothesize that a small number of parameters may switch the model between embedding-oriented representations and language-model generation.
  • RAG optimization: Jointly optimizing retrieval and generation remains a direction for improving RAG beyond the frozen-model caching setup presented here.Prior work cited by the authors finds that jointly optimizing both models gives the best performance.

6 Related Work

Prior work unified tasks within either embedding or generation, but not both. GRITLM addresses this gap by using instructions to support the full breadth of embedding and generative tasks in one model.

  • Embedding Models: Embedding models evolved from word representations to contextual sentence and passage embeddings, but strong performance still required separate symmetric and asymmetric models.
  • Generative Models: Generative models progressed from task-specific systems to large language models capable of handling practically any generative task.
  • Earlier work unified embedding and generative models only within their respective streams, leaving their combination as the next step toward general multi-task models.
  • GRITLM uses instructions rather than parameter switching, supports symmetric and asymmetric use cases, and covers both embedding and generative tasks.

7 Conclusion

The paper presents GRITLM as one model for text embedding and generation. Its 7B and 8x7B variants achieve strong results across both streams, while unified operation supports retrieval and reranking applications.

  • GRITLM unifies text embedding and generation in a single model, matching equivalent embedding-only and generative-only variants at no performance loss.
  • GRITLM 7B achieves state-of-the-art performance on MTEB among open models and beats generative models up to its size.
  • With only 5B additional inference parameters, GRITLM 8x7B is the best open generative language model among the models tried while retaining strong embedding performance.
  • Using GRITLM as both bi-encoder and cross-encoder improves performance on 15 of 16 retrieval datasets.

D Additional RAG results

Additional experiments examine document caching and BGE embeddings for retrieval-augmented generation. Document caching preserves RAG performance on new datasets, whereas BGE performs significantly worse with GRITLM as generator.

  • Document caching maintains or slightly improves regular RAG performance on TriviaQA and MMLU despite attention mismatch.
  • Document caching is identified as the most promising proposed caching variant because documents are longer than queries and can be cached without reducing regular RAG performance.
  • BGE embedding models produce significantly worse RAG performance than GRITLM in the reported comparison.
  • Manual inspection suggests BGE commonly retrieves irrelevant passages that confuse the generative model.

E Evaluation

GRITLM is evaluated using established benchmarks for both embedding and generation. The evaluation covers seven main MTEB tasks and follows the Tülu setup for generative performance.

  • The evaluation selects commonly used embedding and generative benchmarks for assessing GRITLM.
  • Embedding: Embedding performance is evaluated on seven main MTEB tasks, including classification, clustering, and pair classification.
  • Embedding: Classification uses logistic regression over text embeddings and reports F1, while clustering uses k-means and reports V-measure.
  • Generation: Generative evaluation largely follows Tülu and includes multiple-choice question answering on MMLU and problem solving on GSM.

F Ablations Detailed Results

The appendix breaks down ablations for unified, embedding-only, and generative-only models across attention, pooling, datasets, precision, batching, losses, and formats.

  • Attention, pooling, and model setup: Tables 12–15 cover attention, pooling, base-model, and generative-versus-embedding ablations across unified and single-task models.The attention tables specify causal or bidirectional configurations, while the base-model table uses 100-step runs and suboptimal settings.
  • Training data: Tables 16–18 vary embedding datasets, natural instructions, negative sources, and generative training epochs.MEDI2 variants distinguish the inclusion of natural instructions and evaluation instruction sets; the generative ablation reports epoch counts.
  • Optimization and representation choices: Tables 19–23 examine embedding heads, batch sizes, numerical precision, in-batch negatives, and generative formats.The embedding-head ablation includes a 4096-to-1024 down-projection, while the precision table contrasts BF16 and FP32.

G GRITLM MTEB Full Results

This section provides the complete MTEB results underlying the paper’s embedding evaluation. It includes per-dataset scores and aggregate results, including an average row.

  • Full MTEB results: Table 26 expands the full MTEB results from Table 1 into detailed benchmark results.The appendix directs readers to per-dataset results in the MTEB leaderboard and released result files.
  • Per-dataset results: SummEval reports 21.14, 30.18, 30.37, and 29.82 across the compared systems.The passage provides the row values without naming the corresponding system columns.
  • Aggregate results: Average scores are 41.21, 66.82, 66.76, and 65.66 across the compared systems.The aggregate row reports four values for the systems included in the table.

H Reducing Embedding Training Memory

The appendix describes memory-saving methods for embedding training and documents evaluation and benchmarking choices, including instruction handling, formats, precision, and evaluator stability.

  • Memory requirements: 3×bs samples are required for naive in-batch-negative training because each batch contains query, positive-document, and negative-document samples.The appendix motivates memory reduction because all samples must be forwarded together for in-batch negatives.
  • Memory reduction strategies: Triplet splitting reduces memory by separating query and passage computation across forward and backward passes, with extra gradient-free forwards.The implementation reuses detached representations while backpropagating through different triplet components.
  • Memory reduction strategies: Distributed training and GradCache preserve nearly unlimited in-batch negatives while reducing the per-device batch-memory requirement.Distributed training gathers representations across GPUs; GradCache is used when enough GPUs are unavailable.
  • Evaluation settings: MTEB evaluation adds instructions for instruction-tuned models, while HumanEval evaluation uses the intended chat format.The reported comparisons found similar performance with and without these choices, with instructions slightly better for Mistral embeddings.
  • Precision: BF16 mixed-precision evaluation was verified as comparable to FP32 on MTEB, while pooling and similarity computations remained in FP32.The appendix also evaluates BF16 embedding caching and documents the compared precision settings.
  • Generative evaluation: MT-Bench absolute scores changed significantly between GPT-4 and GPT-4 Turbo, motivating the use of comparison-based AlpacaEval.The authors describe absolute ratings as evaluator-dependent and comparison-based judgments as more stable.

O Dataset Samples

The appendix illustrates dataset samples and the instruction formats used for retrieval, clustering, similarity, and other embedding tasks. Examples cover E5, MEDI, MEDI2, and Tülu 2 data.

  • Dataset samples: The examples include MEDI, MEDI2, and Tülu 2 samples alongside retrieval questions and Wikipedia passages.The displayed retrieval example asks which two plates meet along the west coast of the USA and includes relevant tectonic text.
  • Dataset samples: E5 samples omit instructions for documents, so both positive and negative document samples lack instructions.The figure caption contrasts document formatting with instruction-based examples elsewhere in the appendix.
  • MTEB instructions: MTEB retrieval instructions pair query prompts with tasks such as retrieving answering passages, similar texts, duplicates, or supporting documents.The instruction list spans datasets including NQ, MSMARCO, STS*, FEVER, and SciFact.
  • MTEB instructions: Clustering instructions ask representations to identify topics or categories from titles, abstracts, posts, articles, or community content.Examples cover arXiv, bioRxiv, medRxiv, Reddit, StackExchange, StackOverflow, and Twenty Newsgroups.
  • Instruction conventions: MEDI2 evaluation uses separate query and document instructions for asymmetric datasets, while E5 evaluation omits document instructions for retrieval.The appendix states that Q and D designate query and document instructions in MEDI2.
  • MTEB instructions: Representation instructions frame inputs for finding semantically related, duplicate, paraphrased, or same-source texts.The listed tasks include scientific-paper, social-media, forum, news, and community matching variants.

R Limitations and Future Work

The paper identifies several directions for improving GRITLM, including self-retrieval, more efficient training formats, and broader training strategies. It also records implementation updates and experiments concerning training efficiency and model improvements.

  • GRITLM Agents: Self-retrieval could let GRITLM search an index when it is uncertain, reducing reliance on external retrieval plugins.Teaching the model to invoke its embedding capability for this purpose likely requires additional finetuning.
  • Pretraining: Pretraining GRIT from scratch may require unsupervised embedding objectives or data strategies because labeled embedding data is likely scarce.Suggested approaches include RetroMAE, data augmentation, pruning, and multi-epoch training.
  • Format Efficiency: 13-token and 15-token task-format encodings make the current format inefficient, while special tokens could make training and inference slightly cheaper.The embedding format requires 13 tokens and the generative format requires 15 tokens.
  • Training efficiency: Packing and Reusing: Embedding and generative samples could be packed together and reused within one training sample to improve training efficiency.This extends existing sample-packing practices for generative instruction tuning and embedding training.
  • Revision updates: The revision fixed a figure reference, expanded training discussion, added Appendix D, and added KTO experiments.It also discussed compute considerations and potential speed-performance trade-offs involving smaller embedding models or intermediate-layer embeddings.
Loading 2402.09906v3…