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

Generative language models produce poor hidden-state embeddings despite embeddings’ importance for search and chatbots. GRIT trains one language model for generation and embedding, distinguishing tasks through instructions and separate losses. GRITLM 7B achieves state-of-the-art MTEB performance among open models and outperforms generative models up to its size while matching equivalent task-specific variants.

  • Problem

    Existing generative models yield poor embeddings from their hidden states, although embeddings support critical applications such as search engines and chatbots.

    Source: §1 · p. 2

  • Method

    GRIT jointly trains one language model for generation and embedding, using instructions and separate loss functions to distinguish the tasks.

    Source: §1 · p. 2

  • Results

    GRITLM 7B sets the open-model MTEB state of the art and outperforms generative models up to its size, while matching equivalent embedding-only and generative-only variants.

    Source: §7 · p. 16

  • Takeaways & Limitations

    > 60% faster RAG inference with long documents is possible with GRITLM by caching computations and halving the required forward passes.

    Source: §1 · p. 3

    Embedding quality for inputs longer than 512 tokens remains unknown because suitable benchmarks are lacking.

    Source: §3.3 · p. 11

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

Generative representational instruction tuning (GRIT) unifies generation and embedding in one model, matching separately trained variants while enabling faster retrieval-augmented generation.

1 Introduction

  • Existing generative models yield poor hidden-state embeddings, while generating embedding tensors directly is impractical because of their dimensionality and precision requirements.T5 requires finetuning to make its hidden states useful for embeddings, during which it loses its generative capabilities.Source: §1 · p. 2
  • 7B-parameter GRITLM sets the open-model state of the art on MTEB while outperforming larger models, including Llama 2 70B, on generative tasks.GRITLM 8X7B achieves the best open-model generative task average while using 13B parameters at inference. Unified models match embedding-only and generative-only variants, sometimes outperforming them.Source: §1 · p. 2
  • GRIT combines generative and representational instruction tuning, using instructions and separate loss functions to distinguish generation from embedding tasks.The approach is tested on models with up to 47B parameters. Sliding window attention [20] [9] allows generative and embedding inputs of arbitrary length.Source: §1 · p. 2
  • > 60% faster RAG inference with long documents is enabled by sharing embedding and generation computations, halving the required forward passes.Caching reduces the four forward passes needed by separate models to two. A shared model also simplifies load balancing, storage, and serving software compared with separate embedding and generation endpoints.Source: §1 · pp. 2–3
  • GRIT requires more finetuning compute than generative instruction tuning, but is generally cheaper than training separate generation and embedding models when pretraining compute is included.The unified approach requires one pretraining and finetuning process rather than separate processes for both models.Source: §1 · p. 3
  • Additional improvements include bidirectional attention with mean pooling, same-dataset in-batch negatives for embeddings, and mixed sample- and token-level loss aggregation for generation.The paper also proposes methods to reduce memory requirements during embedding training.Source: §1 · p. 3

2 GRIT

GRIT fine-tunes one pretrained language model for embedding and generation using consistent instruction formatting, task-specific attention patterns, and jointly optimized training objectives.

2 GRIT

  • Embedding uses bidirectional attention and mean-pooled final hidden states, trained with a contrastive objective and in-batch negatives.Pooling excludes instruction and format tokens, although they influence representations through self-attention. The contrastive objective uses cosine similarity and a temperature hyperparameter.Source: Figure 3 · p. 3, §2 · pp. 3–4
  • GRIT jointly fine-tunes a pretrained language model on embedding and generative instruction data, combining contrastive and next-token objectives with optional loss weights.A consistent input format supports both tasks and multi-turn conversations.Source: §2 · pp. 3–4, Figure 3 · p. 3
  • Generation uses causal attention and a generation-only language modeling head, with next-token loss computed only over response tokens and the ending token.Instruction and other input tokens are excluded from the generative loss.Source: Figure 3 · p. 3, §2 · p. 4
  • Generative loss aggregation controls weighting: sample-level aggregation weights samples equally, while token-level aggregation gives longer samples more influence.Sample-level aggregation can improve discriminative performance but bias generations shorter; token-level aggregation usually encourages longer responses, which evaluators such as AlpacaEval can favor. The schemes coincide for equal-length samples or batch size one, and mixed schemes can balance their trade-offs.Source: §2 · p. 4
  • GRIT permits different embedding and generative batch sizes, allowing larger embedding batches without increasing the generative batch size.Larger embedding batches can support strong embedding performance, but require more computation per training step.Source: §2 · p. 4

3 Experiments

GRITLM 7B leads prior open models on MTEB and generative models up to its size, while GRITLM 8X7B leads openly available models on the generative average. Joint training matches single-objective variants but requires more compute per training step.

3.1 Setup

  • Final models finetune Mistral 7B and Mixtral 8x7B using E5 augmented with S2ORC and filtered Tülu 2 [65].GRITLM 7B trains for 1253 steps with embedding/generative batches of 2048/256; GRITLM 8X7B uses embedding batches of 256, and separate embedding-triplet passes reduce memory requirements. Evaluation uses 56 main MTEB datasets and largely follows Ivison et al. [65] for generation, substituting HumanEvalSynthesize.Source: §3.1 · p. 4

3.2 Main Results

  • GRITLM 7B leads prior open models on MTEB and generative models up to 7B, while GRITLM 8X7B leads openly available models on the generative average.Joint training matches embedding-only and generative-only variants, although each step requires forward and backward passes on both data types. GRITLM 8X7B has slightly weaker embeddings than GRITLM 7B, plausibly because compute constraints reduced its embedding batch from 2048 to 256.Source: Table 1 · p. 5, Table 2 · p. 5, §3.2 · p. 6
  • Few-shot embedding generally worsens performance, with only marginal, inconsistent gains on some pair-classification tasks.Including few-shot embedding examples in around 5% of MEDI2 training samples still did not teach the model to use them effectively. Increasing retrieval beyond ten top-ranked documents is suggested as a possible improvement, at additional compute cost.Source: Table 4 · p. 7, §3.2 · p. 7
  • >10-point AlpacaEval gains accompany KTO alignment for both models, improving average generative performance while slightly reducing MTEB performance.KTO uses binarized UltraFeedback without further embedding training; the reported MTEB averages fall from 66.8 to 66.7 and from 65.7 to 65.2. GritLM-8x7B-KTO achieves an Alpaca 2.0 [39] length-controlled win rate of 18.5 with average output length 1662.Source: Table 5 · p. 7, §3.2 · p. 7

3.3 Ablations

  • Mistral 7B, filtered Tülu 2 [65], and E5 provide the strongest ablated backbone, generative dataset, and embedding dataset, respectively.Pretrained embedding performance does not predict post-finetuning performance; Tülu 2 wins every evaluated generative task, potentially reflecting task diversity or tuning to the shared evaluation setup. E5's advantage is attributed to hard negatives and task diversity; its unrestricted instructions require <|embed|> formatting, while query-only instructions enable document embeddings to be cached across tasks.Source: §3.3 · p. 8, Table 6 · p. 9
  • Bidirectional embedding attention and fully causal generation perform best, with position-weighted mean pooling preferred over last-token pooling for causal embeddings.Tülu 2 [65] formatting outperforms Zephyr β, whose additional end-of-sequence token encourages premature termination. Mixed sample/token generative loss performs best by a small margin: LRep/LGen initially exceeds 1 before both losses approach 1.0, and token-level loss across many samples helps preserve long generations and AlpacaEval performance.Source: §3.3 · pp. 8, 10, 11, Table 6 · p. 9
  • 1.0 gain on the embedding average results from scaling embedding batches to 4096, while generative performance remains stable.Down-projecting embeddings from 4096 to 1024 dimensions loses around 1% embedding performance and is omitted from final models; GRITLM 7B costs more to run than 335-million-parameter BGE Large and requires 4× its embedding storage. Same-dataset in-batch negatives improve the 15-dataset Retrieval average by 1.3; BF16 mixed precision suffices when pooling and similarity computations retain FP32, with no benefit from additional FP32 computation.Source: §3.2 · p. 6, Table 6 · p. 9, §3.3 · p. 10
  • 2048-token embedding-document training slightly improves embedding and generative performance over 512-token training, despite evaluating embeddings at 512 tokens.Longer training sequences likely preserve information needed to distinguish positive and negative contrastive pairs that truncation would remove. Final models support arbitrary-length embedding inputs, but performance beyond 512 tokens remains unknown because suitable benchmarks are lacking.Source: Table 6 · p. 9, §3.3 · pp. 10–11

4 RAG with GRIT

GRITLM enables RAG caching by unifying embedding and generation, reducing latency for long inputs. Doc Caching achieves the best evaluated answer quality but requires substantial storage, while combined query–document caching performs near No RAG.

Method

  • Query Caching reuses query key-value states between embedding and generation, saving one query forward pass; exact RAG equivalence requires a query-first prompt with causal attention.The experiments instead use bidirectional embedding attention, creating a mismatch with generation’s expected causal states. Embedding-format queries also differ from the formats seen during generative training.Source: §4 · pp. 11, 13

Performance

  • Doc Caching achieves the best Natural Questions match score, slightly improving on RAG, while Query Caching underperforms RAG but beats No RAG.Query Caching scores 25.46 versus 30.50 for RAG. Formatting RAG queries for embedding reduces its score to 29.36; the authors attribute most remaining degradation to attention mismatch.Source: Table 7 · p. 12, §4 · p. 13

Latency

  • 63% CPU and 31% GPU speed-ups occur with Doc Caching for 4000-token documents; Query Caching yields 54% CPU and 33% GPU speed-ups for 4000-token queries.Caching only 250 tokens provides negligible speed-up, motivating selection according to query or document length. The uncached RAG baseline uses GRITLM 7B for both embedding and generation.Source: Table 7 · p. 12, Figure 5 · p. 13, §4 · p. 14

Setup

  • Natural Questions evaluation uses answer-containing match scores rather than exact match to accommodate GRITLM’s sentence-length answers.The implementation adapts Izacard et al., uses the BEIR NQ corpus, and supplies no embedding instructions beyond formatting. Latency benchmarks generate 16 tokens on an Intel Xeon Platinum 8481C CPU and an NVIDIA H100 80GB GPU.Source: §4 · pp. 12–13, Table 7 · p. 12

Performance

  • Query-Doc and Doc-Query Caching avoid both forward passes but perform near No RAG, likely because separately cached queries and documents cannot attend to each other.Query-Doc Caching places the query first; Doc-Query Caching places the document first. Additional RAG finetuning is proposed as a possible remedy for attention and formatting mismatches.Source: §4 · pp. 11–14

Storage

  • 30TB of key-value states are required for 2,681,468 documents, trading storage for one avoided inference forward pass per retrieved document.GRITLM caches document states during indexing without additional computation; these can reside on disk, loading around 12.5MB per sample. Equivalence with RAG requires placing the document first in the prompt.Source: §4 · pp. 11, 14

5 Discussion

GRIT unifies embedding and generation without performance loss, with improved embedding performance under MEDI2 training. The discussion identifies multilingual and multimodal extensions and shared retriever–reader optimization as directions for further work.

Further unification

  • GRITLM supports non-English embedding and generation, but targeted multilingual improvements and unification of multimodal embedding and generation remain open directions.TyDi QA demonstrates non-English capability; the authors suggest multilingual data and architecture changes [42] could improve performance. Whether multimodal tasks can be unified as easily as text tasks remains unexplored.Source: §5 · p. 14

Why does GRIT work?

  • GRIT preserves embedding and generation performance, and adding the generative objective improves MEDI2 embedding performance over an otherwise equivalent embedding-only model.The authors hypothesize that both tasks share language understanding, with a small parameter subset switching representations between mean pooling for embeddings and the language-modeling head for generation.Source: §5 · pp. 14–15

Optimizing RAG with GRITLM

  • GRITLM could simplify joint retriever–reader optimization through shared parameters, potentially using one next-token objective to penalize irrelevant retrieval and poor context use.The reported RAG and caching variants use a frozen model, so this optimization remains a proposal. Prior work found joint optimization performed best, but Lin et al. used separate models and objectives.Source: §5 · p. 15

6 Related Work

Prior work progressively unified tasks within embedding and generation through prompting and instruction tuning. GRITLM extends this unification across both streams, using instructions without switching biases.

Embedding Models

  • Prompting unified symmetric and asymmetric embedding tasks, while detailed instructions extended coverage to practically any embedding task.Contextual models such as InferSent [28] and SBERT improved on word representations that struggled with sentences and passages [27]. Symmetric tasks share query–document distributions, whereas asymmetric tasks differ; MTEB revealed that SentT5 performed well only on symmetric tasks and GTR only on asymmetric tasks.Source: §6 · p. 15

Generative Models

  • Instruction finetuning eased applying LLMs to arbitrary generative tasks with strong results.McCann et al. unified generative tasks as question answering, but performance and generalization remained limited. Large-scale self-supervised pretraining enabled broad generative task coverage [13] [22] [11] [55] [1], although performance without careful prompting was often poor.Source: §6 · p. 15

Unifying Embedding and Generation

  • GRITLM unifies embedding and generation without switching biases, using instructions for symmetric and asymmetric tasks while addressing both streams’ full task breadth.SGPT changes only 0.01% of an LLM’s parameters through BitFit to produce embeddings. However, it requires separate symmetric and asymmetric models and does not address the full breadth of embedding tasks.Source: §6 · p. 15

7 Conclusion … C Loss Curves

GRIT unifies text embedding and generation in GRITLM without losing performance relative to equivalent single-purpose variants, while enabling improved reranking and faster retrieval-augmented generation.

7 Conclusion

  • 7B-parameter GRITLM achieves state-of-the-art MTEB performance among open models and beats generative models up to its size; GRITLM 8X7B leads tested open generative models.GRITLM 8X7B adds only 5B parameters at inference and outperforms tested models including Llama 2-based 70B models, while retaining strong embedding performance.Source: §7 · p. 16
  • GRITLM matches otherwise equivalent embedding-only and generation-only variants, unifying both capabilities without performance loss.Instructions distinguish embedding tasks from generation tasks within the unified model.Source: §7 · p. 16
  • > 60% inference speed-ups for long sequences are achieved without performance loss through GRIT Doc Caching, which unifies the RAG retriever and reader.Using one model enables caching operations that accelerate inference.Source: §7 · p. 16
  • Around 10% retrieval improvement comes from reusing GRITLM as a reranker, with improvements on 15 out of 16 retrieval datasets.The unified model serves as both the Bi-Encoder and Cross-Encoder rather than requiring a separate reranking model.Source: §7 · p. 16
  • Embedding ablations favor bidirectional attention and mean pooling, find pre-finetuning performance nonpredictive, and support BF16 finetuning with FP32 pooling and similarity computations.BF16 mixed-precision embedding finetuning can match FP32 finetuning. Generative instruction tuning should use some form of token-level loss.Source: §7 · p. 16

A Contributions

  • Niklas Muennighoff led execution, experiments, model development, and most writing; Hongjin Su created MEDI2, implemented and ran reranking, and helped with experiments.Liang Wang ran dataset ablations and contributed to framing; Nan Yang, Furu Wei, Tao Yu, Amanpreet Singh, and Douwe Kiela advised the project. All authors helped edit the paper.Source: A Contributions · p. 29

B Artifacts

  • The artifacts appendix distinguishes produced artifacts from used artifacts released by others.Table 8 identifies produced artifacts, while Table 9 identifies used artifacts released by others; the supplied captions do not enumerate them.Source: Table 8 · p. 29, Table 9 · p. 31

C Loss Curves

  • 0.9-weight exponential moving average smoothing is applied to the training-loss figures for GRITLM 7B and GRITLM 8X7B.Figure 6 covers GRITLM 7B, and Figure 7 covers GRITLM 8X7B; their supplied captions do not establish loss trends or convergence comparisons.Source: Figure 6 · p. 31, Figure 7 · p. 32

D Additional RAG results

Doc caching preserves regular RAG performance on additional datasets despite an attention mismatch. BGE embedding models yield substantially worse RAG performance than GRITLM.

D Additional RAG results

  • Document caching offers the greatest latency-reduction potential because documents usually exceed queries in length, and it also works with non-GRIT models.Constructing the cache takes longer for non-GRIT models.Source: D Additional RAG results · p. 32
  • Disabling bidirectional attention for embeddings resolves the attention mismatch but slightly reduces embedding performance.The authors state that this guarantees the same performance as not using RAG.Source: D Additional RAG results · p. 32
  • Doc caching maintains or slightly improves regular RAG performance on TriviaQA [74] and MMLU [61], despite an attention mismatch.These additional experiments use the same setup as Table 7.Source: Table 10 · p. 32, D Additional RAG results · p. 32
  • BGE embedding models produce substantially worse RAG performance than GRITLM; manual inspection suggests irrelevant retrieved passages confuse the generative model.The authors expect RAG performance to depend on both embedding and generative capabilities, while acknowledging that other smaller embedding models or other generators might perform better.Source: D Additional RAG results · p. 32

E Evaluation

GRITLM evaluation covers seven MTEB embedding tasks and six generative benchmark families, with task-specific metrics, instructions, and prompting protocols.

Embedding

  • Embedding evaluation uses seven MTEB tasks: classification, clustering, pair classification, reranking, retrieval, semantic textual similarity, and summarization.Classification uses F1; clustering uses V-measure; pair classification uses average precision; reranking uses mean average precision; retrieval uses nDCG@10; semantic textual similarity and summarization use Spearman correlation. Datasets span domains from scientific papers to casual conversations.Source: E Evaluation · p. 33
  • Embedding evaluation treats reranking, retrieval, and summarization as asymmetric query–document tasks, while the other four tasks are symmetric.Every dataset receives instructions specified in §P.1; some models receive different instructions for query and document embeddings on asymmetric tasks.Source: E Evaluation · p. 33

Generation

  • Generative knowledge and reasoning evaluation measures exact match on MMLU [61], GSM [25], and BIG-Bench Hard.MMLU uses no few-shot examples, GSM uses eight chain-of-thought examples, and BIG-Bench Hard uses three. Generation evaluation largely follows Tülu [65] using open-source frameworks.Source: E Evaluation · p. 33
  • Multilingual closed-book question answering uses TyDi QA [24] across six languages, evaluated in Gold Passage and no-context settings following Anil et al. [2].These settings distinguish evaluation with a supplied passage from answering without context.Source: E Evaluation · p. 33
  • 20 samples at temperature 0.2 determine pass@1 on HumanEvalSynthesize [17], following Muennighoff et al..The Python benchmark adapts HumanEval [17] to an instruction-following format, unlike the non-instruction HumanEval evaluation used by Ivison et al. [65].Source: E Evaluation · p. 33
  • Open-ended generation uses AlpacaEval 1.0 [40], with GPT-4 judging win rates against provided GPT-3 [13] answers.The maximum token length is reduced from 8192 to 6144 relative to Ivison et al. [65]. The authors report greedy evaluation throughout for reproducibility and exclude MT-Bench because of limitations discussed in Appendix M.Source: E Evaluation · p. 33

F Ablations Detailed Results

The detailed ablations examine attention, pooling, datasets, embedding architecture, and training settings. The supplied captions explain experimental variants but do not establish their relative performance.

Attention and pooling

  • Attention and pooling ablations cover unified, embedding-only, and generative-only models, including interleaved bidirectional instructions and causal answers.Pooling variants include position-weighted mean, last-token, and mean pooling. Interleaved attention permits answer KV-cache reuse for faster multi-turn generation.Source: Table 12 · p. 34, Table 13 · p. 34, Table 14 · p. 34

Datasets and instructions

  • Dataset ablations examine embedding instructions, negative sources, and generative datasets, with epoch counts identified for generative training.NNI excludes natural instructions; MEDI2 NNI II and MEDI2 NNI share a trained model but use Instructor-XL instructions [145] versus the new domain, intent, and unit structure at evaluation. MEDI2BGE uses MEDI2 with negatives from BGE’s MTP training dataset [171].Source: Table 16 · p. 35, Table 17 · p. 35, Table 18 · p. 36

Embedding head

  • 4096 to 1024 dimensions is the down-projection tested through a linear layer on the final hidden state, applied only to embedding tasks.The embedding-head ablation changes the representation dimensionality without applying that projection to generative tasks.Source: Table 19 · p. 36

Training settings

  • 256 and 4096 are the embedding batch sizes tested, with generative batches fixed at 256; other ablations examine precision, in-batch negatives, and base models.Precision variants are BF16 mixed precision and FP32. Base-model ablations use only 100 training steps and suboptimal settings, including Zephyr formatting, corrected in later ablations.Source: Table 15 · p. 35, Table 20 · p. 36, Table 21 · p. 36

Formats and token limits

  • Generative-format and maximum-token ablations distinguish embedding-document training limits from query-and-document embedding-evaluation limits.The maximum-token notation X:Y specifies the training limit first and the evaluation limit second.Source: Table 24 · p. 37, Table 23 · p. 37

Loss aggregation

  • 32 samples receive token-level loss before sample-level aggregation across 8 sub-batches in Mix (32 -> 8), giving a total batch size of 256.Loss-ablation notation also identifies the first-step embedding-to-generative loss ratio, LEmb/LGen.Source: Table 25 · p. 38

G GRITLM MTEB Full Results

Table 26 presents the full MTEB results corresponding to Table 1. The supplied excerpt reports Average and SummEval values without identifying their corresponding models.

G GRITLM MTEB Full Results

H Reducing Embedding Training Memory … N Dataset Composition

These appendices describe memory-saving embedding training, training settings, and evaluation choices involving instructions, chat formatting, and numerical precision. They also report MT-Bench sensitivity to evaluator choice and provide E5S and MEDI2 dataset composition tables.

H Reducing Embedding Training Memory

  • 3*bs samples require simultaneous forward-and-backward memory in naive embedding training because triplets and in-batch negatives prevent ordinary gradient accumulation.Training splits queries from positive and negative passages, requiring one additional gradient-free query forward pass. Backpropagating through only a subset of embeddings instead worsens performance.Source: Figure 8 · p. 39, H Reducing Embedding Training Memory · p. 40
  • Distributed training gathers negatives across GPUs, while GradCache [51] preserves in-batch negatives when enough GPUs are unavailable.Distribution can span up to bs GPUs. GradCache computes gradients one triplet at a time, effectively enabling contrastive gradient accumulation at the cost of additional forward passes.Source: H Reducing Embedding Training Memory · p. 40

I Hyperparameters

  • 1253 steps is the maximum full-parameter finetuning duration, using a 2e-5 learning rate, 3% linear warm-up, and linear decay to zero.Sequence lengths are 2048 for generative samples, 256 for embedding queries, and 2048 for embedding documents unless otherwise specified. Memory savings combine PyTorch FSDP, gradient checkpointing, BF16 mixed precision, and Appendix H’s strategies.Source: I Hyperparameters · p. 40

J Embedding Instruction for Generative Models

  • Embedding instructions slightly improve Mistral 7B Instruct’s MTEB performance, motivating their use for all instruction-tuned models in embedding evaluations.The instructed condition uses Instructor XL’s default instructions with Mistral Instruct’s prompt template.Source: Table 27 · p. 41, J Embedding Instruction for Generative… · p. 41

K HumanEval Format

  • Chat formatting yields approximately equivalent HumanEvalSynthesize performance for Tülu 2 7B [65], supporting chat-format evaluation for chat models.Chat-model evaluation follows HumanEvalSynthesize instructions and procedures. Non-chat models retain the original HumanEval continuation format proposed by Chen et al. [17].Source: Table 28 · p. 41, K HumanEval Format · p. 41

L Embedding in FP32 vs BF16

  • BF16 mixed precision performs comparably to FP32 on MTEB, while pooling and similarity computations remain FP32.All training and evaluations use BF16 mixed precision; the precision comparison tests raw Mistral 7B. The BF16 Cache condition stores pooled embeddings in BF16 and restores FP32 before similarity calculations.Source: L Embedding in FP32 vs BF16 · p. 41, Table 29 · p. 42

M Unreliability of MT-Bench

  • MT-Bench scores change significantly when GPT-4 Turbo replaces GPT-4 as judge, motivating comparison-based AlpacaEval [40] instead of absolute ratings.Zephyr 7B β and Llama 2 70B Chat often failed to provide ratings because they struggled with the prompt. The authors argue that pairwise comparisons are more stable than subjective absolute scores, while acknowledging shared limitations.Source: Table 30 · p. 42, M Unreliability of MT-Bench · p. 42

N Dataset Composition

  • Dataset composition tables organize E5S and MEDI2 by constituent dataset and sample count.The supplied table headings do not include individual dataset counts.Source: Table 31 · p. 43, Table 32 · p. 43

O Dataset Samples · P Evaluation Prompts · Q Hardware

These appendices present dataset samples and evaluation materials, including training-dataset-specific MTEB query and document instructions, one-shot embedding examples, and generative benchmark prompts. They also illustrate RAG examples, GRITLM 7B’s RAG prompt ordering and caching variants, and document training hardware, compute costs, and speed bottlenecks.

O Dataset Samples

P.1 Embedding Prompts

  • MTEB prompts follow Wang et al. for E5, use author-written instructions for MEDI2, and follow Instructor-XL from Su et al. for MEDI.Tables 33 and 34 list the evaluation instructions for E5-trained and MEDI2-trained models, respectively.Source: §P.1 · p. 47, Table 33 · p. 47, Table 34 · p. 49
  • E5 retrieval evaluation omits document instructions, while MEDI2 distinguishes query and document instructions for asymmetric datasets.E5 evaluation also uses the same instructions across all STS tasks.Source: Table 33 · p. 47, Table 34 · p. 49
  • Embedding instructions specify task-dependent matches: supporting evidence, search-query answers, likely cited papers, and sentences with the same meaning.FEVER and MSMARCO provide separate query and document instructions. SCIDOCS targets likely cited papers, while STS12 targets equivalent sentences.Source: §P.1 · pp. 51–52

P.2 Embedding Few-Shot Prompts

  • 1-shot examples for models trained on E5S and MEDI2 are appended to their respective instructions, separated by two newlines.Examples illustrate counterargument retrieval, scientific-paper retrieval, sentence paraphrases, and paired news summaries.Source: Table 35 · p. 53, §P.2 · pp. 54–55, Table 36 · p. 55

P.3 Generative Prompts

P.4 RAG Prompts

  • The RAG prompt materials include a no-RAG example and generation by GRITLM 7B.Source: Figure 19 · p. 63
  • GRITLM 7B examples demonstrate query-first and document-first RAG prompts, plus Query Caching, Doc Caching, and Query-Doc Caching.Doc Caching passes document embedding inputs through key-value states cached in the index.Source: Figure 20 · p. 63, Figure 21 · p. 64, Figure 22 · p. 64
  • 18.39, up from 14.13, was the match score when GRIT Doc-Query Caching prepended the bos token (“<s>”) to both query and document, unlike Doc Caching.The figure illustrates this variant with generation by GRITLM 7B.Source: Figure 24 · p. 65

Q Hardware

  • 3,072 GPU hours trained GRITLM 7B and 20,480 GPU hours trained GRITLM 8X7B, with roughly 100,000 GPU hours estimated overall.GRITLM 7B used 8 nodes with 8 NVIDIA A100 80GB GPUs each for 48 hours; GRITLM 8X7B used 32 nodes with 8 NVIDIA H100 80GB GPUs each for 80 hours; Gen.-only and Emb.-only models used 72 and 1760 H100 80GB GPU hours, respectively. Both GRITLM models trained for 1253 steps, taking several minutes per step because of large batches, large models, memory-saving strategies that sacrifice speed, and slow inter-node communication.Source: Q Hardware · p. 66

R Limitations and Future Work · S Version Control

Proposed GRIT extensions include self-initiated retrieval, pretraining from scratch, and more efficient formatting and training. The version history records expanded discussions, added experiments and resource reporting, and corrections.

R Limitations and Future Work

  • GritLM could initiate searches using its own embedding capability, potentially eliminating external retrieval plugins.Teaching self-invocation likely requires additional finetuning, as with external plugins. The illustrated sequence embeds a query, retrieves an answer, and generates a response.Source: R Limitations and Future Work · p. 66
  • Pretraining GRIT from scratch could use unsupervised embedding objectives or data-constrained training methods because labeled embedding data is likely too scarce.The experiments instead use an off-the-shelf pretrained language model. Proposed alternatives include RetroMAE, data augmentation [35], pruning, and multi-epoch training.Source: R Limitations and Future Work · p. 66
  • 13 tokens for embedding formatting and 15 tokens for generative formatting motivate using special tokens to make training and inference slightly cheaper.These counts describe the formatting overhead of the current format shown in Figure 3.Source: R Limitations and Future Work · p. 67
  • Packing embedding samples and reusing samples across embedding and generation could reduce GRIT training resource requirements.Packing would restrict attention to each respective sample. Joint reuse could apply generative loss to assistant responses and contrastive loss to representations following <|embed|>.Source: R Limitations and Future Work · p. 67

S Version Control

  • Revisions expanded training-improvement discussion, clarified the pretraining-inclusive compute rationale, and discussed intermediate-layer embeddings as a potential speed–performance trade-off.The version history also records added resource reporting for the Gen.-only and Emb.-only baselines in Appendix Q.Source: S Version Control · p. 67
  • Version updates added KTO experiments in §3.2 and Appendix D, corrected a figure reference in §4, and fixed the Figure 12 sample.Appendix D is listed as a separate addition; the version notes do not describe its contents.Source: S Version Control · p. 67

Related papers

Loading 2402.09906v3…