One email each morning — yesterday's AI, sortedGet it in your inbox

METAL LAB

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

arXiv:2608.202812026-08-19

Teaching LLMs to answer document questions from memory, without search, using a three-stage training recipe

This paper introduces IAR (Inject, Align, Recover), a three-stage post-training method that lets language models answer questions about a fixed set of documents without retrieving them at answer time. It separates dense document exposure (Inject), question-answer format adaptation (Align), and weight merging to restore general skills (Recover). Tested across two datasets and four model families, IAR generally raises both domain accuracy and general capability compared to standard fine-tuning.

What they did

  1. Defines the problem of 'document knowledge internalization': converting a fixed document collection into knowledge stored in a model's parameters so it can answer questions without retrieval at inference time
  2. Inject stage exposes the model to documents through three tasks (continuation, rewrite, instruction-conditioned reconstruction) that give denser training signal than question-answer pairs alone
  3. Align stage fine-tunes only on the answer portion of QA pairs to adapt the injected model to a question-answering interface; Recover stage merges the adapted model with the original instruction model using techniques like SLERP and TIES to restore lost general abilities
  4. Compared to standard supervised fine-tuning, IAR improved all four reported metrics in 7 of 8 dataset-model settings, with average gains of 3.6 percentage points in domain QA accuracy and 12.1 percentage points in general performance across IFEval, MMLU, and MSBench
  5. Scaling Qwen3 from 8B to 32B parameters, IAR stayed within 1.1 points of peak domain accuracy while recovering 14.9 to 24.1 points of general performance, showing the pattern holds across model sizes
Figure 1: The overview of IAR. Vanilla SFT learns from generated QA pairs and covers only the facts selected by those questions. CPT+SFT models full document token streams before a separate QA alignment stage. IAR injects document knowledge through continuation, rewrite, and instruction-conditioned reconstruction objectives, aligns the injected model with answer-only QA supervision, and recovers general ability via post-hoc model merging. The final model is selected by balancing retrieval-free domain QA performance and general capability retention.
Figure 1: The overview of IAR. Vanilla SFT learns from generated QA pairs and covers only the facts selected by those questions. CPT+SFT models full document token streams before a separate QA alignment stage. IAR injects document knowledge through continuation, rewrite, and instruction-conditioned reconstruction objectives, aligns the injected model with answer-only QA supervision, and recovers general ability via post-hoc model merging. The final model is selected by balancing retrieval-free domain QA performance and general capability retention.
Table 1: Main comparison for RQ1. Scores are reported as percentages. Within each model block and dataset block, bold denotes the best result and underlining denotes the second-best result for each metric. CPT+SFT starts from the corresponding released Base checkpoint; “/” indicates that no such release is available, as for Phi-4-mini.
ModelMethodCCCCI
Dom.IFEvalMMLUMSB.Dom.IFEvalMMLUMSB.
Llama-3.2-3BBase Instruct11.277.150.854.529.277.150.854.5
Vanilla SFT35.554.211.221.553.061.222.531.5
CPT+SFT38.326.43.713.553.724.24.317.5
IAR36.560.235.030.555.361.333.236.5
Phi-4-miniBase Instruct13.177.261.062.027.377.261.062.0
Vanilla SFT24.447.851.032.540.247.853.831.5
CPT+SFT////////
IAR34.149.057.043.039.751.650.244.0
Qwen3-4BBase Instruct34.384.865.877.070.684.865.877.0
Vanilla SFT42.451.18.851.075.145.626.349.5
CPT+SFT49.631.818.860.569.029.612.858.0
IAR50.559.819.563.076.376.164.570.0
SmolLM3-3BBase Instruct15.377.147.763.034.177.147.763.0
Vanilla SFT32.135.610.525.052.341.716.731.5
CPT+SFT37.126.94.826.048.924.14.219.5
IAR37.540.325.729.053.957.446.847.0
Figure 2: BudgetMatch-to-IAR movement. Right is higher domain accuracy; up is a higher mean over IFEval, MMLU, and MSBench. IAR moves up and right in three of four settings; for CC Llama, it trades 3.9 domain points for an 11.0-point gain in mean general performance.
Figure 2: BudgetMatch-to-IAR movement. Right is higher domain accuracy; up is a higher mean over IFEval, MMLU, and MSBench. IAR moves up and right in three of four settings; for CC Llama, it trades 3.9 domain points for an 11.0-point gain in mean general performance.
Table 2: Token-budget ablation for RQ2 (higher is better; scores are percentages). BudgetMatch uses setting-specific QA-only epochs matched to the Inject+Align budget. IAR improves both domain accuracy and mean general performance in three settings; CC Llama is the remaining domain–general trade-off. Supplementary Material, Section C gives the full accounting.
ModelMethodCCCCI
Dom.IFEvalMMLUMSB.Dom.IFEvalMMLUMSB.
Llama-3.2-3BVanilla SFT35.554.211.221.553.061.222.531.5
BudgetMatch40.453.417.322.053.045.322.524.0
IAR36.560.235.030.555.361.333.236.5
Qwen3-4BVanilla SFT42.451.18.851.075.145.626.349.5
BudgetMatch46.849.431.845.072.256.129.552.5
IAR50.559.819.563.076.376.164.570.0
Figure 3: Pre-recovery domain gains from Inject+Align. CCI Qwen3-4B’s +0.4 pp is the high-base-prior boundary case.
Figure 3: Pre-recovery domain gains from Inject+Align. CCI Qwen3-4B’s +0.4 pp is the high-base-prior boundary case.
Table 3: Complete Qwen3 scaling ablation on CC. Bold and underline mark the best and second-best result within each model block. Across 8B/14B/32B, IAR stays within 1.1 points of Best IA domain accuracy while recovering 14.9–24.1 points in mean general performance.
ModelMethodDomain (%)IFEval (%)MMLU (%)MSBench (%)
Qwen3-8BBase Instruct38.587.665.382.5
Vanilla SFT48.756.414.052.5
Best IA57.550.618.548.5
IAR (TIES d=0.3)56.862.226.773.5
Qwen3-14BBase Instruct40.490.072.581.5
Vanilla SFT54.862.954.557.0
Best IA60.553.540.342.0
IAR (TIES d=0.3)59.667.567.273.5
Qwen3-32BBase Instruct47.287.574.884.5
Vanilla SFT56.458.544.056.5
Best IA63.953.063.044.5
IAR (TIES d=0.3)62.867.074.572.5
Figure 4: Recover candidate frontiers for the main dataset–model settings. Each panel plots held-out test performance of the fixed Recover candidates by retrieval-free domain QA accuracy and mean general performance across IFEval, MMLU, and MSBench. The black line marks non-dominated points in this two-dimensional projection, and the black ring highlights the candidate selected by the validation protocol for the main IAR table. The displayed test frontier is diagnostic and is not used to choose the selected candidate.
Figure 4: Recover candidate frontiers for the main dataset–model settings. Each panel plots held-out test performance of the fixed Recover candidates by retrieval-free domain QA accuracy and mean general performance across IFEval, MMLU, and MSBench. The black line marks non-dominated points in this two-dimensional projection, and the black ring highlights the candidate selected by the validation protocol for the main IAR table. The displayed test frontier is diagnostic and is not used to choose the selected candidate.
Table 4: Dataset construction contract. The train/test split and retrieval-free input column define the experimental setting: models must answer from internalized parameters rather than from retrieved source passages. The files named eval_750.jsonl and eval_575.jsonl in the repository are treated as held-out test files in this paper.
DatasetSplitCountInput at train timeInput at test timePrimary use
CCtrain14,258 QAquestion, answer, derived document fieldsquestion onlymixed-domain internalization
test750 QA/question onlyheld-out domain test
CCItrain10,926 QAquestion, answer, derived document fieldsquestion onlyChinese-domain internalization
test575 QA/question onlyheld-out domain test
Table 5: Document-derived QA construction prompts. The schemas show how chunks are converted into self-contained questions and grounded answers while preventing deictic questions that require access to the original document.
StepPrompt schemaOutput contract
Anchor extractionGiven a text chunk, extract at most K independently referable core objects. Prefer explicit concepts, methods, mechanisms, modules, devices, or technical terms appearing in the text. Do not output deictic objects such as “this method” or “the above mechanism.”JSON list of anchors.
Type applicabilityGiven a question-type description and the text chunk, decide whether the chunk can support a question of that type. The supported types are factual extraction, mechanism explanation, design rationale, condition/constraint, limitation/trade-off, and comparison/relation.yes/no.
Question generationGiven the chunk, a selected anchor, a question type, the type description, and an expected answer schema, generate one natural question. The anchor must be explicitly named; the question must be understandable without the source document; deictic expressions such as “this,” “above,” or “according to the text” are forbidden; output only one question ending with a question mark.Plain question text.
Question validationCheck whether the question is independently understandable, avoids document/deictic references, is semantically clear, and has an answer direction.JSON with valid and reason.
Answer generationGiven the source chunk and generated question, answer strictly from the chunk. The answer must be faithful, accurate, professional, directly answer the question, avoid document/deictic references, and be written as a natural paragraph rather than a template with section headings.JSON with answer.
Table 6: QA-generation stage accounting. “Files filtered” gives the count and percentage of input documents rejected at file-level filtering; the chunk, question, and QA columns give retained/considered counts and micro rates. “Experiment QA” is the train-plus-test total after exact-question deduplication and fixed chunk-group sampling.
DatasetInput docsFiles filteredChunks keptValid questionsQA keptExperiment QA
CC4,0011,944 (48.6%)5,324/6,593 (80.8%)37,397/77,224 (48.4%)16,674/37,397 (44.6%)15,008
CCI7,000310 (4.4%)11,407/11,793 (96.7%)100,096/140,148 (71.4%)62,670/100,096 (62.6%)11,501
Table 7: Inject objective definitions. All three objectives use the instruction model’s chat template and assistant-target loss; the loss mask excludes every system/user token. Recipe ratios control the relative counts of the three objective streams, while realized shares can differ slightly after tokenizer-specific length filtering.
ObjectiveUser input uAssistant target yMasked tokensRecipe roleIntended exposure
ContinuationContinue/complete instruction plus a document prefix.Held-out suffix.System prompt, instruction, and prefix.single or mixedPrefix-conditioned document exposure.
RewriteReconstruction instruction plus a generated summary, outline, or knowledge skeleton.Full cleaned document.System prompt, instruction, and compressed representation.single or mixedRecover document content from a compressed representation.
Instruction-formatted reconstructionShort generic reading instruction.Full cleaned document.System prompt and instruction.1:0:0 or mixedDense exposure through a full-document target.
Table 8: Shared optimization settings recovered for the completed 3B/4B training runs and extended CC baselines. The effective global batch is per-device batch 1×8 accumulation steps ×8 GPUs.
ParameterValueParameterValue
OptimizerAdamW (PyTorch)Learning rate5×10−5
Adam β1,β2,ϵ.9,.999,10−8Scheduler / warmupcosine / ratio .05
Weight decay / max grad norm.01/1.0Precision / max lengthBF16 / 4096
Batch per GPU1Gradient accum. / GPUs8/8
Effective global batch64 examples/stepDeepSpeedZeRO-2, no offload
Termination / checkpointepoch based / final epochGradient checkpointingmodel-level enabled
Table 9: Stage-specific training settings. Unless stated as an exception, each row uses Table 8. BudgetMatch epoch counts are setting-specific rather than a shared 13-epoch approximation.
Method / stageEpochsData and objectiveMethod-specific setting
Vanilla SFT3QA; answer-onlyOriginal Instruct initialization
BudgetMatch14/17/21/11Same QA and loss as VanillaCC Llama/Qwen; CCI Llama/Qwen order
Inject3Three assistant-target document-generation objectivesSelected mixtures and counts in Table 10
Align3QA; answer-onlyInitializes from the Inject final epoch
SDFT3Model-specific synthetic QA; answer-only14,258 CC examples per model
LoRA3CC QA; answer-onlyr=16, α=32, dropout .05; merged for evaluation
Replay375% domain QA + 25% general instructionEqual-size replacement; construction seed 42
CPT16Raw-document causal LMMatched Base initialization; 4,425 CC / 10,769 CCI rows
CPT+SFT3QA; answer-onlyInitializes from the CPT final epoch
Table 10: Selected Inject configurations for the eight main settings. Counts are realized post-tokenization training rows; tokenizer-specific length filtering explains the small differences among nominally equal mixtures.
DatasetModelSelected Inject recipeInject rowsAlign QA rows
CCLlama-3.2-3BMixed 1:1:219,00014,258
Phi-4-miniMixed 1:1:219,00014,258
Qwen3-4BMixed 1:1:118,96814,258
SmolLM3-3BMixed 1:1:118,96014,258
CCILlama-3.2-3BMixed 1:1:219,00010,926
Phi-4-miniMixed 1:1:219,00010,926
Qwen3-4BReconstruction 1:0:010,00010,926
SmolLM3-3BMixed 1:1:219,00010,926
Table 11: Realized token accounting for BudgetMatch, in millions of non-padding model tokens. BM/IA compares the completed QA-only BudgetMatch run directly with the corresponding Inject+Align token volume.
SettingInject recipe (samples)InjectAlignIA totalBudgetMatchBM/IA
CC Llama-3.2-3BMixed 1:1:2 (19k)45.73612.91758.65360.280102.8%
CC Qwen3-4BMixed 1:1:1 (18,968)46.01710.19156.20857.752102.7%
CCI Llama-3.2-3BMixed 1:1:2 (19k long-doc.)57.2969.73467.03068.085101.6%
CCI Qwen3-4B1:0:0 (10k)19.0817.15126.23226.20299.9%
Table 12: Evaluation metrics and their roles in operating-point selection. Domain QA is the primary retrieval-free internalization metric, while IFEval, MMLU, and MSBench serve as general-capability guardrails.
MetricDirectionEvaluation typeReported unit
Domain accuracyhigher is betterV2 multi-judge correctnesspercentage
IFEval inst-stricthigher is betterinstruction-following evaluatorpercentage
MMLU accuracyhigher is bettermultiple-choice benchmarkpercentage
MSBench accuracyhigher is betterLLM-judge benchmarkpercentage
Table 13: Domain QA judge reliability. Binary agreement collapses scores at ≥.5. The full audit additionally contains 600 samples from a smaller CCI evaluation set. These statistics quantify judge reliability over the evaluated samples.
ScopeRecordsExact agreeBinary agreeBinary κThird judge
CC test artifacts163,347.732.854.691.274
CCI test artifacts78,308.656.836.670.345
Full audit242,255.707.848.691.297
Table 14: Per-judge score distributions before aggregation. Vote counts differ because panel configurations vary across evaluations and later judges are invoked adaptively; the marginal means are therefore descriptive and should not be interpreted as a controlled judge ranking.
JudgeVotesMean scoreP⁡(s≥.5)P⁡(s=1)
gpt-oss-120b242,031.387.454.321
MiniMax2.5_Local156,812.303.434.172
deepseek-v3.2131,172.296.465.127
deepseek-v3.124,697.378.652.103
Table 15: Recover candidate grid. Every selected IAR row is chosen from this fixed set of post-hoc merge candidates rather than from an unreported per-row search space. Candidate choice is made on the validation split before held-out test reporting.
Operator familyHyperparameter gridCandidate count
SLERPt∈{0.2,0.3,0.4}3
Task Arithmeticw∈{0.3,0.5,0.7}3
TIESd∈{0.3,0.5,0.7}3
DAREdr∈{0.1,0.3,0.5}3
Totalfixed grid per IA checkpoint12
Table 16: Selected Recover settings for the main IAR rows. The table reports held-out test scores for the concrete operating point chosen from the fixed Recover grid by the validation protocol, making the domain-primary selection rule explicit rather than treating recovery as an unreported hyperparameter search.
DatasetModelSelected checkpointDomain (%)IFEval (%)MMLU (%)MSBench (%)Selection note
CCLlama-3.2-3BTIES d=0.336.560.235.030.5domain-primary feasible point
Phi-4-miniTask Arithmetic w=0.734.149.057.043.0balanced feasible point
Qwen3-4BTIES d=0.350.559.819.563.0domain-primary feasible point
SmolLM3-3BTIES d=0.337.540.325.729.0domain-primary feasible point
CCILlama-3.2-3BTIES d=0.355.361.333.236.5domain-primary feasible point
Phi-4-miniTIES d=0.339.751.650.244.0boundary trade-off point
Qwen3-4BTask Arithmetic w=0.776.376.164.570.0domain-primary feasible point
SmolLM3-3BTask Arithmetic w=0.753.957.446.847.0domain-primary feasible point
Table 17: Conditional BPB on identical CCI source-document continuations. Corpus BPB is byte weighted; the primary effect first averages paired Qwen-minus-peer BPB differences within each of seven domains and then weights domains equally. Negative Δ means lower BPB for Qwen. All intervals exclude zero and all seven domain means agree in direction.
Initial checkpointCCI QA (%)Doc. mean BPBCorpus BPBQwen–peer macro Δ [95% CI]
Qwen3-4B-Instruct70.60.7440.729
Llama-3.2-3B-Instruct29.21.0210.998−0.273 [−0.285, −0.261]
Phi-4-mini-instruct27.30.9680.948−0.220 [−0.231, −0.208]
SmolLM3-3B34.10.8380.819−0.093 [−0.102, −0.084]
Table 18: Officially disclosed Qwen3 dense-family base-model scores from the Qwen3 technical report. These are public general and multilingual benchmark results, not CCI QA results, and are included only to contextualize the high-base-prior interpretation.
ModelMMLUMMLU-ProBBHMGSMMMMLUINCLUDE
Qwen3-1.7B-Base62.6336.7654.4750.7163.2745.57
Qwen3-4B-Base72.9950.5872.5967.7471.4256.29
Qwen3-8B-Base76.8956.7378.4076.0275.7259.40
Qwen3-14B-Base81.0561.0381.0779.2079.6964.55
Qwen3-32B-Base83.6165.5487.3883.0683.8367.87
Table 19: Qwen3 CCI high-base-prior diagnostics. The 1.7B and 4B rows show that Qwen3 starts unusually high on CCI, so this setting is better interpreted as preserving and recovering a strong prior than as creating a large new domain gain.
ModelSettingCCI domain (%)
Qwen3-1.7BBase Instruct60.0
Vanilla SFT57.0
Mixed 1:1:1 + Stage256.7
Qwen3-4BBase Instruct70.6
Vanilla SFT75.1
Context-aware SFT65.0
Mixed 1:1:1 + Stage273.7
Continue + Stage272.9
1:0:0 + Stage275.5
Rewrite + Stage274.4
Mixed 1:1:2 + Stage272.5
Replay66.4
Table 20: Qwen3-4B CCI Recover sweep. Multiple merge operators stay in the high-domain range, with Task Arithmetic and TIES reaching the main selected value; this supports the boundary-case reading of the Qwen3 CCI row.
OperatorCoefficientCCI domain (%)OperatorCoefficientCCI domain (%)
DAREdr=0.174.1SLERPt=0.269.0
DAREdr=0.374.3SLERPt=0.372.7
DAREdr=0.573.7SLERPt=0.473.2
Task Arithmeticw=0.371.1TIESd=0.376.3
Task Arithmeticw=0.574.1TIESd=0.572.5
Task Arithmeticw=0.776.3TIESd=0.775.0
Table 21: Complete CC extended-baseline comparison for Llama-3.2-3B and Phi-4-mini.
ModelMethodDomain (%)IFEval (%)MMLU (%)MSBench (%)
Llama-3.2-3BVanilla SFT35.554.211.221.5
SDFT39.958.423.533.5
LoRA22.564.053.230.0
Replay30.856.730.737.5
Vanilla-FAPM22.375.351.862.0
IA-FAPM22.474.553.759.0
IAR36.560.235.030.5
Phi-4-miniVanilla SFT24.447.851.032.5
SDFT32.943.258.745.5
LoRA18.070.659.339.5
Replay24.748.728.848.5
Vanilla-FAPM17.156.458.264.5
IA-FAPM16.953.861.363.0
IAR34.149.057.043.0
Table 22: Complete CC extended-baseline comparison for Qwen3-4B and SmolLM3-3B.
ModelMethodDomain (%)IFEval (%)MMLU (%)MSBench (%)
Qwen3-4BVanilla SFT42.451.18.851.0
SDFT44.155.914.053.0
LoRA31.368.248.553.0
Replay42.869.734.766.5
Vanilla-FAPM41.283.861.383.5
IA-FAPM44.182.665.585.0
IAR50.559.819.563.0
SmolLM3-3BVanilla SFT32.135.610.525.0
SDFT36.549.844.325.5
LoRA15.566.322.542.0
Replay33.653.128.034.0
Vanilla-FAPM24.771.035.568.5
IA-FAPM26.169.449.067.0
IAR37.540.325.729.0
Table 23: Available Instruct-initialized CPT+SFT diagnostics. CC Llama exposes a domain–general trade-off rather than uniform IAR dominance, while the CC Phi and CCI Llama diagnostics remain below the corresponding best IA domain scores. Coverage is limited to completed archived runs and is not extrapolated to Qwen3 or SmolLM3.
DatasetModelDomain (%)IFEval (%)MMLU (%)MSBench (%)
CCLlama-3.2-3B38.743.59.819.5
CCPhi-4-mini31.243.340.731.0
CCILlama-3.2-3B52.929.08.323.0
Table 24: FAPM domain results at sparsity 0.9. The domain-only comparison shows that pruning-based recovery often sacrifices internalized document knowledge, even when it is useful for general-capability restoration.
DatasetModelVanilla SFT (%)Vanilla-FAPM (%)IA-FAPM (%)
CCLlama-3.2-3B35.522.322.4
Phi-4-mini24.417.116.9
Qwen3-4B42.441.244.1
SmolLM3-3B32.124.726.1
CCILlama-3.2-3B53.038.640.7
Phi-4-mini40.227.328.7
Qwen3-4B75.171.568.9
SmolLM3-3B52.343.841.6

Why it matters

For real-world settings like company manuals or legal documents where retrieval systems are impractical due to latency or privacy concerns, this offers a concrete method to bake document knowledge directly into a model. It also demonstrates a systematic way to reduce the common side effect of domain fine-tuning where models lose their ability to follow instructions or perform general tasks.

Terms in this paper

  • retrieval-free QA · answering questions using only knowledge stored in the model, without fetching the source document at answer time
  • supervised fine-tuning (SFT) · further training a model on labeled example data so it learns a specific desired behavior
  • continued pretraining (CPT) · having a model keep reading domain-specific text to absorb its style and facts
  • catastrophic forgetting · losing previously learned skills while a model is being trained on something new
  • model merging · mathematically combining the weights of two differently trained models into a single model

Original abstract (English)

Large language models often fail to answer questions about a bounded document collection when the source documents are not retrieved at inference time. We study this setting as document knowledge internalization: converting a fixed corpus into usable parametric knowledge for retrieval-free question answering. We propose IAR (Inject, Align, and Recover), a three-stage post-training framework that separates structured document knowledge injection, QA behavior alignment, and general ability recovery. Unlike conventional continued pretraining, Inject converts source documents into continuation, rewrite, and instruction-conditioned reconstruction objectives. Align then adapts the injected model with answer-only QA supervision, while Recover merges the domain-adapted model with the base instruction model to recover general capabilities. Across Common Corpus (CC) and CCI, and across Llama, Phi, Qwen, and SmolLM model families, IAR improves the domain-primary domain-general frontier for retrieval-free document internalization. In the main comparison, IAR improves over Vanilla SFT on all four reported metrics in 7 of 8 dataset-model settings, with average gains of 3.6 percentage points in domain QA accuracy and 12.1 percentage points in mean general performance across IFEval, MMLU, and MSBench. Extended CC baselines show that LoRA and FAPM can win individual general metrics, but among methods that also reach leading or near-leading domain internalization, IAR retains one of the strongest general profiles.

Authors · Qian Kou

Read on arXiv

Latest papers

All papers →

Latest from METAL LAB

Figures: Qian Kou et al., arXiv:2608.20281, CC BY 4.0