GRITLM: Unified Generation & Embedding Model

• GRITLM is a unified model that combines generative text and embeddings through instruction-based objective switching.
• It uses distinct training losses and special tokens to toggle between causal and bidirectional attention, enhancing dual-mode performance.
• The architecture streamlines retrieval-augmented generation by reducing inference passes while achieving state-of-the-art results on embedding and generative tasks.

The term GRITLM refers to "Generative Representational Instruction Tuning" (GRIT), a methodology and family of models for unifying text generation and embedding within a single LLM via instruction-based objective switching. The GRITLM architecture, training paradigm, and empirical validation address limitations of prior approaches that silo generative and representational (embedding) capabilities, establishing new performance benchmarks in both modalities and directly impacting downstream scenarios such as Retrieval-Augmented Generation (RAG) by enabling highly efficient unified inference (Muennighoff et al., 2024).

1. Model Variants and Architectural Framework

GRITLM consists of two primary open-model variants:

• GritLM 7B: Built upon the Mistral 7B model, this version is a decoder-only, causal-attention transformer supporting sliding-window attention to accommodate arbitrarily long contexts.
• GritLM 8×7B: Mixture-of-experts backbone based on Mixtral 8×7B, comprising eight 7B expert models (total 47B parameters), with roughly 13B parameters active per forward pass.

Both variants share core transformer parameters and are extended with mechanisms for dual-purpose instruction following.

Model	Backbone	Parameters	Active Parameters
GritLM 7B	Mistral 7B	7B	7B
GritLM 8×7B	Mixtral 8×7B (MoE)	47B	~13B

The design leverages standard causal self-attention for generation and bidirectional attention with mean-pooling for embeddings. Switching between these two regimes is achieved by special instruction tokens, notably <|user|>, <|assistant|>, and <|embed|>, with the latter explicitly signaling embedding-mode and triggering the corresponding loss and attention mask.

2. Training Objectives and Instruction-Based Modality Selection

GRITLM employs a dual-objective training schema:

• Contrastive Embedding Loss ($\mathcal{L}_{\mathrm{Rep}}$): $$\mathcal{L}_{\mathrm{Rep}} = -\frac{1}{M}\sum_{i=1}^M \log \frac{\exp(\tau \cdot \cos(f_\theta(q^{(i)}), f_\theta(d^{(i)})))}{\sum_{j=1}^M \exp(\tau \cdot \cos(f_\theta(q^{(i)}), f_\theta(d^{(j)})))}$$ Mean-pooled representations are computed from bidirectional attention for samples prefixed with <|embed|>.
• Generative Language Modeling Loss ($\mathcal{L}_{\mathrm{Gen}}$): $$\mathcal{L}_{\mathrm{Gen}} = -\frac{1}{N}\sum_{t=1}^{N} \log P_\eta(x^{(t)}|x^{(<t)};\theta)$$ Applied to instruction-response pairs using the causal attention mechanism.

Instruction tokens determine the format—embedding tasks use <|embed|>, generative tasks use <|assistant|>. At finetuning, both sample types are present in each batch (e.g., 2048 Rep/256 Gen per step for 7B), with token-level weighting maintaining balance between objectives. The combined objective is $$\mathcal{L}_{\mathrm{GRIT}} = \lambda_{\mathrm{Rep}} \mathcal{L}_{\mathrm{Rep}} + \lambda_{\mathrm{Gen}} \mathcal{L}_{\mathrm{Gen}}$$ Typical weighting ratios are $\lambda_{\mathrm{Rep}} \approx 4$, $\lambda_{\mathrm{Gen}} \approx 1$.

3. Pretraining and Finetuning Data Composition

GRITLM's mixed-batch finetuning draws from distinct data sources:

• Embedding (Rep) Data: ~11M examples (MEDI2 plus curated scientific S2ORC passages "E5S"), covering classification, retrieval, clustering, and similarity with diverse, explicit domain/intention/unit instructions. Batch sizes reach 2048 per step for the 7B model.
• Generative (Gen) Data: ~1M filtered instruction-response pairs from OpenAssistant, UltraChat, and cleaned OASST. Batch size 256 samples per step regardless of model.
• Batch Proportion: ~88% embedding vs. 12% generative at the sample level for 7B; token-level adjustment further normalizes gradient contribution.

Instruction formats and the <|embed|> switch ensure gradient routing exclusively follows the respective loss and attention regime per sample. All model parameters are shared; only the head and attention mask are task-dependent.

4. Empirical Performance on Embedding and Generation Tasks

4.1 Text Embedding (MTEB Benchmark)

GritLM 7B sets a new open model state of the art on the Massive Text Embedding Benchmark (MTEB, 56 datasets), outperforming embedding-only and concurrent 7B models, and demonstrating robust retrieval nDCG and classification accuracy:

Model	Params	MTEB Avg.
E5 Mistral 7B	7B	66.6
Emb-only 7B	7B	66.8
GritLM 7B	7B	66.8
GritLM 8×7B	47B	65.7

4.2 Generative Task Performance

Tested across MMLU, GSM8K, BBH, TyDi QA, HumanEvalSynthesize, and AlpacaEval, GritLM 7B leads all open 7B models; GritLM 8×7B outperforms all open generative models (including Llama 2 70B):

Task	GritLM 7B	GritLM 8×7B
MMLU (EM)	57.6	66.7
GSM8K (EM)	57.5	61.5
BBH (EM)	54.8	70.2
TyDiQA (F1)	55.4	58.2
HumanEvalSynthesize (pass@1)	32.8	53.4
AlpacaEval (% win)	74.8	84.0
Avg.	55.5	65.7

4.3 Objective Ablation

Ablation experiments confirm that joint GRITLM training achieves full generative and embedding performance, whereas "only-generative" or "only-embedding" training yields random or collapsed outcomes on the other modality (e.g., generative-only 7B: MTEB ≈ 41.2, embedding-only: generative avg ≈ 7.6).

5. Unified Model Impacts on Retrieval-Augmented Generation

Traditional RAG pipelines require separate inference with a retriever and a generator, necessitating up to four forward passes per query-context pair. GRITLM's unification enables parameter sharing (θ) between retrieval (embedding) and generation, supporting direct reuse of key/value Transformer cache states for substantial speed-ups:

• Query Caching: Embedding the query and caching (k, v) for reuse in generation saves 1 pass/query.
• Doc Caching: Storing precomputed (k, v) per document allows further elimination of duplicate model passes.
• Joint Caching: Full reuse of cached states but with a minor trade-off in answer accuracy due to attention regime mismatches.

On NaturalQuestions with 4000-token documents:

Method	GPU Latency (s)	CPU Latency (s)	RAG Match (0-shot)
RAG baseline	0.40 / 11.6	—	30.5%
Query Caching	0.27 / 6.9	—	25.5%
Doc Caching	0.27 / 5.3	—	33.4%

Query caching achieves >30% GPU and >50% CPU speed-up for long inputs; doc caching yields >60%, typically at equal or improved accuracy.

6. Scaling Behavior, Trade-Offs, and Implementation

• Scaling:

Generative performance increases with model size: 7B→8×7B yields generative average 55.5→65.7 (+10.2) with only a marginal drop in embedding average (66.8→65.7). Inference cost grows sublinearly due to the mixture-of-experts setup (7B→~13B active parameters).

• Resource and Deployment Considerations:
  • GritLM 7B inference requires ~13GB GPU VRAM (BF16 precision); 8×7B requires ~25GB (for 13B active).
  • Mixed-precision inference matches FP32 in quality for both modalities, enabling memory and performance gains.
  • The model exposes a unified endpoint ("/generate", "/embed"), handling both modalities via input tokens.
  • Sliding-window sparse attention supports arbitrarily long context inputs.
• Limitations and Trade-Offs:
  • Training cost is elevated due to dual-objective passes (~3,072 GPU-h for 7B, ~20,480 for 8×7B).
  • Unified serving doubles latency over a standalone embedder but eliminates cross-system complexity.
  • Efficient training of large Rep batches necessitates distributed/GradCache techniques.
  • Few-shot generative prompting has minor negative effects on embedding performance, but this scenario is rare in retrieval settings.

7. Code and Model Release

GRITLM models, training scripts, benchmarks, and downstream indices are released under open-access terms:

• Source: https://github.com/ContextualAI/gritlm
• GritLM 7B weights: https://huggingface.co/GritLM/GritLM-7B
• GritLM 8×7B weights: https://huggingface.co/GritLM/GritLM-8x7B
• Prebuilt RAG index: https://huggingface.co/datasets/GritLM/index

The unified GRITLM paradigm demonstrates that dual-mode training (generation+embedding with instruction-driven objective switching) is scalable, efficient, and achieves state-of-the-art results in both text representation and generative tasks, with wide practical benefits in retrieval pipelines and multimodal language technology (Muennighoff et al., 2024).