[Paper Review] Generative Representational Instruction Tuning

Niklas Muennighoff, Hongjin Su+6 more|arXiv (Cornell University)|Feb 15, 2024

Innovative Teaching and Learning Methods13 citations

TL;DR

GRIT unifies embedding and generation in a single large language model, achieving state-of-the-art embedding performance and strong generative results, while enabling faster retrieval-augmented generation. The approach trains a model to handle both tasks via instruction-based signals and two concurrent losses.

ABSTRACT

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.

Motivation & Objective

• Motivate a single model that performs well on both text embeddings and generative tasks rather than separate systems.
• Develop a training paradigm that blends embedding representational learning with generative instruction tuning.
• Demonstrate that a unified GRIT model can match or exceed specialized embedding and generation baselines across tasks.
• Show practical benefits of GRIT for retrieval-augmented generation and model serving infrastructure.

Proposed method

• Finetune a pretrained LLM with two data streams: embedding (representational) and generative instruction data.
• Use bidirectional attention with mean pooling for embeddings and causal attention with a language modeling head for generation.
• Optimize a joint loss L_GRIT = lambda_Rep L_Rep + lambda_Gen L_Gen with controllable loss weights and a mix of sample- and token-level generative loss.
• Experiment with data sources (embedding: MEDI, MEDI2, E5; generative: Tülu 2, UltraChat, OpenAssistant) and base models (Mistral 7B, Mixtral 8x7B, etc.).
• Evaluate embedding on MTEB across 56 datasets and generation on standards like MMLU, GSM8K, BBH, TyDi QA, and HumanEval variants.
• Investigate architectural choices (attention type, pooling, embedding head, batch size, precision) and training regimes to balance dual objectives.

Experimental results

Research questions

• RQ1 Can a single model achieve state-of-the-art performance on text embeddings while maintaining strong generative capabilities?
• RQ2 Is joint optimization essential to performance, or do embedding-only and generation-only models suffice?
• RQ3 What are the practical gains in retrieval-augmented generation when using a unified GRIT model?
• RQ4 How do data sources, attention schemes, and loss configurations impact dual-task performance?
• RQ5 What are the trade-offs in training efficiency and inference for a unified model?

Key findings

• GritLM 7B sets a new state-of-the-art on the Massive Text Embedding Benchmark (MTEB) among open models.
• GritLM 7B also outperforms larger open generative models on several generation tasks.
• GritLM 8x7B becomes the best open generative model on the task average while using 13B at inference.
• Unified GRIT training matches embedding-only and generative-only variants, with ablations showing the joint loss is critical for dual-task performance.
• RAG with GRIT yields >60% faster inference for long documents by enabling caching and using a single model for both retrieval and generation.
• Using larger embedding batch sizes (e.g., 2048) improves embedding performance without harming generation.