Inject, Align, Recover: Staged Post-Training for Retrieval-Free Document Knowledge Internalization

The Problem: Forgetting What You Learned

Large language models struggle to answer questions about documents they've memorized during training if those documents aren't provided at inference time. The challenge: how do you teach a model to internalize a fixed corpus of documents as permanent knowledge—without needing to retrieve the source text when answering questions later? Standard fine-tuning on QA pairs tends to work well on domain questions but erodes the model's general capabilities. IAR tackles this head-on by treating it as a staged post-training problem rather than a simple supervised fine-tuning task.

How IAR Works: Three Deliberate Stages

The authors separate the learning process into three phases, each solving a different subproblem.

Inject converts raw documents into training objectives designed to encode structural knowledge. Instead of just showing the model question-answer pairs, this stage creates three kinds of learning signals: continuation (predicting what comes next in a document), rewrite (paraphrasing document content), and instruction-conditioned reconstruction (answering "given this instruction, reproduce key content from the document"). This multi-faceted approach embeds document knowledge more deeply than naive continuation pretraining.

Align then fine-tunes the injected model using answer-only QA supervision—the model sees questions paired with correct answers from the corpus, but without explicit document retrieval. This grounds the internalized knowledge in the actual downstream task.

Recover is the crucial stabilization step. Rather than leaving the model in a domain-specialized state, the authors merge the domain-adapted model back toward the base instruction model. This hybrid approach retains the internalized document knowledge while restoring general capabilities that fine-tuning typically damages.

What It Achieves

IAR was tested on two corpora (Common Corpus and CCI) across four model families (Llama, Phi, Qwen, SmolLM). The results show consistent wins: in 7 of 8 dataset-model settings, IAR outperformed vanilla supervised fine-tuning on all four metrics. On average, domain QA accuracy improved by 3.6 percentage points, and general performance (measured across IFEval, MMLU, and MSBench) improved by 12.1 percentage points compared to standard fine-tuning.

Critically, while simpler methods like LoRA or FAPM can win on individual general-capability metrics, IAR maintains one of the strongest profiles when requiring both strong domain internalization and preserved general knowledge—the real bottleneck in practice.

Why It Matters

This work reframes a practical constraint—inference without retrieval—into a solvable design problem. Many deployed systems can't afford to retrieve documents at inference time due to latency or cost. The three-stage framework shows that knowledge internalization and capability retention aren't locked in zero-sum tradeoff if you architect the post-training pipeline carefully. The recovery stage is conceptually simple but empirically powerful: explicit model merging to restore forgotten abilities while keeping domain knowledge intact.