Knowledge injection

Overview

A pretrained language model knows only what was in its training data, up to a cutoff, and coverage is uneven across topics. Knowledge injection is the set of ways to add facts from a document collection afterwards. The problem is that the model's accuracy on questions about those documents must improve, and the method must be chosen against that goal.

The shared test used across this page is multiple-choice accuracy: a model "knows" a fact when it picks the right option more often than chance 2312.05934 §2. Because a good reasoner can score above chance by elimination, a gain on this test is not pure knowledge 2312.05934 §2 2312.05934 app. C.

Variants

Continued training on the documents (fine-tuning)

Start from a saved checkpoint and continue next-token training on the raw knowledge base, without labels. The paper uses a learning rate below the pre-training rate to limit forgetting, chunks text into fixed 256-unit pieces, and marks document edges with special tokens 2312.05934 §3.2 2312.05934 §5. Card: Fine-Tuning or Retrieval? Comparing Knowledge Injection in LLMs.

On the MMLU subsets, this variant gave mixed results: gains over the base model in most rows and drops below it in others 2312.05934 tab. 1. On the new current-events questions, training on the original text was not competitive; for Llama2-7B it dropped accuracy from 0.353 to 0.219 2312.05934 tab. 2.

Paraphrase-based continued training

The same training, but on rewrites of the chunks that keep the same facts, instead of on the original text. For current events, each chunk got ten paraphrases 2312.05934 §4.4 2312.05934 §6. This was the only fine-tuning form that improved on the base model for all three tested models 2312.05934 tab. 2. Mistral-7B went from 0.481 to 0.588 and Orca2-7B from 0.456 to 0.566 2312.05934 tab. 2. The authors report that accuracy rose with the number of paraphrases for every model tested, but they call the need for repetition a hypothesis and the paraphrase counts appear only in a plot 2312.05934 §6 2312.05934 fig. 4.

Retrieval-augmented generation (RAG)

No training. Embed each document once, store the vectors, and at query time retrieve the K nearest documents by dot product and place them before the question 2312.05934 §3.3. The paper embeds with bge-large-en and stores vectors in FAISS 2312.05934 §5. The main card for this variant is the retrieval-augmented-generation concept.

Across the MMLU subsets, base plus RAG beat fine-tuning alone in 28 of the 30 rows; the two exceptions are Llama2-7B 0-shot on astronomy and on college chemistry 2312.05934 tab. 1. On current events, base plus RAG reached 0.875 for Mistral-7B and 0.876 for Orca2-7B, against 0.481 and 0.456 for the base models 2312.05934 tab. 2. The authors say RAG did especially well there because each question came from a chunk that is in the retrieval knowledge base 2312.05934 §5 2312.05934 §4.3.

Fine-tuning and RAG combined

The fine-tuned model is used as the generator inside the RAG pipeline 2312.05934 §5. The gain over RAG alone was not consistent. The authors call it unstable 2312.05934 §5. On current events, FT-par plus RAG reached 0.830 for Mistral-7B and 0.826 for Orca2-7B, against 0.875 and 0.876 for RAG alone on the base models 2312.05934 tab. 2.

When to use which

Constraints that change the choice:

• Latency and per-query cost. RAG adds an embedding index and extra input tokens on every query 2312.05934 §3.3. The paper's main tables do not state which K was used, and K was not stable across tasks 2312.05934 app. A 2312.05934 tab. 3.
• Training effort. The fine-tuning runs used 4 A100 GPUs, a learning-rate search from 1e-6 to 5e-5 and at most 5 epochs 2312.05934 §5.
• Forgetting. The authors suspect catastrophic forgetting for the weak fine-tuning results, but did not measure it 2312.05934 §5. See catastrophic-forgetting.
• Data generation. Paraphrases and questions were written with GPT-4 2312.05934 §4.3 2312.05934 §4.4. A different generator may change the results; the paper does not test one.

The results cover three 7B models, Wikipedia text, a few MMLU subjects and 910 current-events questions 2312.05934 §8. Do not assume the ranking holds for larger models, other domains or corpora that were not chunked the same way.

Implementation sketch

The two variants share one evaluation step and differ only in where the facts go. Pseudocode in the style of the paper cards; comments cite the source.

# Evaluation of one multiple-choice question  [arxiv:2312.05934 §5] [eq. 4]
# Input:  model, question with options, optional retriever with K
# Output: 1 if the model's pick is correct, else 0

def score_question(model, question, retriever, k):
    # 1. Optional RAG: prepend the k nearest documents  [arxiv:2312.05934 §3.3]
    prompt = question.text
    if retriever is not None:
        docs = retriever.top_k(question.text, k)  # dot product
        prompt = concat(docs, question.text)  # separator not stated
    # 2. Score each option by its log-probability  [arxiv:2312.05934 eq. 4]
    scores = [log_prob(model, prompt + opt) for opt in question.options]
    # 3. Pick the best option; length normalisation not stated
    return int(argmax(scores) == question.answer)

# Paraphrase-based continued training  [arxiv:2312.05934 §3.2] [arxiv:2312.05934 §4.4]
# Input:  base checkpoint, documents, paraphrases per chunk
# Output: fine-tuned checkpoint

def finetune_on_paraphrases(model, documents, n_par):
    # 1. Cut text into 256-unit chunks; mark document edges  [arxiv:2312.05934 §5]
    chunks = pack(documents, length=256, mark_edges=True)
    # 2. Add several rewrites of each chunk that keep the facts  [§4.4]
    data = [p for c in chunks for p in paraphrase(c, n_par)]
    # 3. Plain next-token training, no labels; small learning rate  [§3.2] [§5]
    for epoch in range(5):  # at most 5 epochs in the paper
        for batch in make_batches(data, size=64):
            model = apply_update(model, next_token_loss(model, batch))
    return model

Contradictions

None known. Only one card lists this slug. Note that the paper's own K-ablation gives K from 1 to 5, while its stated grid runs from 0 to 5 2312.05934 §5 2312.05934 tab. 3; the main-table K is not stated.

Papers

• Fine-Tuning or Retrieval? Comparing Knowledge Injection in LLMs: compares unsupervised continued training with RAG on 7B models, on MMLU subsets and new current-events questions; RAG wins in nearly all settings and paraphrases help fine-tuning.