Learning by Distilling Context

Charlie Snell; Dan Klein; Ruiqi Zhong

Source-linked AI summary

Learning by Distilling Context

Charlie Snell, Dan Klein, Ruiqi Zhong

arXiv:2209.15189v1Sep 30, 2022cs.CLcs.AI

TL;DR

Language models benefit from prompts, examples, explanations, and scratchpads but do not retain those gains once the context is removed. Context distillation trains the same model to turn detailed teacher-context completions into direct predictions under minimal context, and the paper reports successful internalization across instructions, explanations, examples, and reasoning. The method also supports overwriting old updates and combining contexts beyond the context-window size.

Problem
Language models lose context-dependent performance when instructions or scratchpads are removed, while those contexts can be much longer than task inputs and exceed the context window.
Method
Context distillation generates answers from detailed teacher prompts, then fine-tunes the same model to predict extracted final answers from simpler student prompts.
Results
Context distillation internalizes abstract instructions and explanations, concrete training examples, and step-by-step reasoning across the reported settings.
Takeaways & Limitations
The framework supports iterative parameter updates, overwriting previous associations, and combining more training context than fits in one context window.
Takeaways & Limitations
With the current model, not all tasks benefit from internalizing natural-language explanations.

Abstract

from arXiv · show

Language models significantly benefit from context tokens, such as prompts or scratchpads. They perform better when prompted with informative instructions, and they acquire new reasoning capabilities by generating a scratch-pad before predicting the final answers. However, they do not \textit{internalize} these performance gains, which disappear when the context tokens are gone. Our work proposes to apply context distillation so that a language model can improve itself by internalizing these gains. Concretely, given a synthetic unlabeled input for the target task, we condition the model on ``[instructions] + [task-input]'' to predict ``[scratch-pad] + [final answer]''; then we fine-tune the same model to predict its own ``[final answer]'' conditioned on the ``[task-input]'', without seeing the ``[instructions]'' or using the ``[scratch-pad]''. We show that context distillation is a general method to train language models, and it can effectively internalize 3 types of training signals. First, it can internalize abstract task instructions and explanations, so we can iteratively update the model parameters with new instructions and overwrite old ones. Second, it can internalize step-by-step reasoning for complex tasks (e.g., 8-digit addition), and such a newly acquired capability proves to be useful for other downstream tasks. Finally, it can internalize concrete training examples, and it outperforms directly learning with gradient descent by 9\% on the SPIDER Text-to-SQL dataset; furthermore, combining context distillation operations can internalize more training examples than the context window size allows.

1 INTRODUCTION

Context tokens improve language-model performance but their benefits disappear when the context is removed. Context distillation addresses this by practicing with informative contexts and fine-tuning the model to reproduce the resulting answers with minimal context.

Context tokens improve zero-shot learning, in-context learning, and complex reasoning when they provide task definitions, examples, explanations, or scratchpads.These gains arise from informative prompts and step-by-step reasoning before the final answer.
Removing context tokens eliminates these gains and can impose substantial inference overhead because instructions and scratchpads may exceed task-input length by more than 10x.The context-window limit also makes it unclear how to use contexts whose total length exceeds the available window.
Context distillation adapts repeated human-like practice by generating answers with detailed instructions and reasoning, then fine-tuning the same model to predict answers from simpler prompts.The method is illustrated using synthetic addition questions and is intended to distill capabilities into model parameters.
The framework is presented as a general training method spanning abstract instructions, concrete examples, and step-by-step reasoning.The introduction highlights applications to instruction learning, SPIDER Text-to-SQL, and reasoning capabilities.
On SPIDER Text-to-SQL, context distillation outperforms direct gradient-descent learning by 9% for 8-shot adaptation and continues improving when distilling more examples than fit in the context window.The introduction also reports internalization of instructions, explanations, and reasoning skills.
The authors anticipate greater usefulness as models improve at in-context learning and future tasks require longer descriptions and reasoning chains.This is presented as a future expectation rather than a demonstrated result.

2 CONTEXT DISTILLATION

Context distillation trains a model to convert teacher-side context into parameters: it generates completions from detailed prompts, extracts answers, and fine-tunes the same model under a leaner prompt. Multiple distillation operations can be combined to accumulate or overwrite learned contexts.

2.1 INTUITION AND MAIN COMPONENTS: Unlike classical distillation, context distillation begins with the same teacher and student model parameters and drives learning through differences between their prompts.Classical distillation instead typically starts with distinct teacher and student parameters.
2.1 INTUITION AND MAIN COMPONENTS: The framework is defined by a raw input distribution D, teacher and student templates, an answer extractor f, and teacher parameters used to generate training completions.D supplies core task inputs; the templates transform them into teacher and student prompts.
2.1 INTUITION AND MAIN COMPONENTS: The framework samples raw task inputs, generates teacher completions from detailed prompts, extracts final answers, and fine-tunes the same model to predict them from student prompts.The teacher prompt may contain instructions, explanations, and examples, while the student prompt contains substantially less context.
2.2 FORMAL DESCRIPTION: The formal objective samples x from D, samples a teacher completion y conditioned on TTEACHER(x), and maximizes student likelihood of f(y) conditioned on TSTUDENT(x).Throughout the process, the teacher parameters remain fixed.
2.3 COMBINING MULTIPLE UPDATES: Simultaneous distillation sums multiple context-distillation losses, making it useful when independent instructions and examples collectively exceed the context window.The combined objective represents K separate distillation operations.
2.3 COMBINING MULTIPLE UPDATES: Sequential distillation initializes each update from the preceding student parameters, enabling incremental updates and overwriting previous updates; the two variants can be composed.A recursive variant can additionally make the updated student the next teacher.
2.4 IMPLEMENTATION DETAILS: Directly training on extracted hard labels wastes token-logit information and produces noisy gradients, motivating token-level KL-divergence optimization with approximations for large vocabularies.The full soft-label distribution can require substantial memory because language-model vocabularies are often 50–100k tokens.

3 EXPERIMENTS

Experiments show that context distillation internalizes abstract instructions and explanations, concrete examples, and step-by-step reasoning, while supporting task updates and downstream transfer.

Experimental setup: The experiments use Natural-Instructions-V2 and few-shot-sampled raw inputs to test internalization of abstract task instructions, explanations, examples, and reasoning skills.The student and teacher differ in which contextual signals they receive, allowing the experiments to measure what context distillation transfers.
Learning from abstract instructions and explanations: r = 0.75, p = 1%: Distillation gains correlate positively with the benefit of natural-language explanations, indicating that the method learns explanations when they help the teacher.Not all tasks benefit from internalizing explanations for the current model.
Learning from abstract instructions and explanations: Mixed simultaneous distillation prevents the student from associating task inputs rather than task ids with the required task, successfully overwriting a prior task-id association.The naïve input distribution permits this shortcut, whereas the mixed variant successfully overwrites the past association.
Learning from concrete examples: 8.6%: On SPIDER Text-to-SQL, context distillation outperforms gradient descent on four examples in exact-set match accuracy, with the margin increasing for eight examples.Simultaneous distillation also enables learning from examples whose combined length exceeds the context window.
Learning from step-by-step reasoning: 0% to 94.7%: After distillation, direct 8-digit addition improves from 0% to 94.7%, showing that step-by-step reasoning can be internalized without scratch-pads.Both transfer-learning and multitask-learning baselines perform more than 20% worse than context distillation.
Learning from step-by-step reasoning: 1% to 17%: Direct 8-digit addition improves after distillation while Natural Instructions performance remains roughly stable, from RougeL 57 to 58.On related synthetic addition tasks, performance also increases from 17% to 30%, indicating transfer to downstream applications.

4 RELATED WORK

Prior work shows that language models learn from task definitions, explanations, examples, and scratch-pad reasoning, while context distillation builds on these context-based learning signals.

Prompting and Instruction Tuning: Language models can learn from abstract task definitions, natural-language explanations, and concrete in-context examples.
Scratch Pad: Scratch-pad prompting improves language-model performance by eliciting chains of reasoning before final answers.
Distillation: Context distillation is related to neural-network knowledge distillation but focuses on transferring capabilities induced by context.

5 CONCLUSION

The paper presents context distillation as a general learning method that internalizes instructions, examples, and reasoning, with applications to behavioral control and factual editing.

Conclusion: Context distillation can internalize abstract statements, concrete examples, and step-by-step reasoning as a general learning method.
Factual Knowledge Editing: The method offers a way to edit factual knowledge by distilling prompts that state new or modified declarative facts.
Controlled Text Generation: Distilling prompts that specify positive sentiment enables a language model to generate positive movie reviews without retaining the full teacher instruction.

A.2 FACTUAL KNOWLEDGE EDITING WITH CONTEXT DISTILLATION

Context distillation applies prompt-based factual edits to Counterfact, recovering teacher paraphrase performance but underperforming on neighborhood prompts because the teacher itself performs poorly there.

Evaluation: The Counterfact evaluation tests whether edited facts generalize across paraphrases while preserving unrelated knowledge under neighborhood prompts.
Evaluation: Table 5 reports paraphrase and neighborhood performance using average score and average magnitude based on correct-versus-incorrect object probabilities.
Evaluation: Context distillation largely recovers the teacher’s fact-editing performance and is comparable in absolute score to ROME and MEND, despite using more computational resources.
Limitation: Poor teacher performance on neighborhood prompts leads to poor student performance, limiting the demonstrated editing quality in that setting.

B EXPERIMENT DETAILS

The experiments fine-tune an 11B T5 language model on Natural-Instructions-V2, a broad instruction-tuning dataset containing task descriptions, examples, and explanations.

Experiment Setup: The study fine-tunes an LM-adapted T5-11B model on Natural-Instructions-V2, which contains 1600+ language tasks.

B.2 FINETUNING LANGUAGE MODEL IMPLEMENTATION DETAILS

Experiments use 32 TPU-V3 cores with model parameters and optimizer states sharded equally across the cores. The implementation uses Jax and PJIT for model-parallel training and inference.

Experiments run on 32 TPU-V3 cores.
Model parameters and optimizer states are sharded equally across all cores.
The codebase uses Jax with PJIT for model-parallel training and inference.

B.3 GENERAL EXPERIMENT DETAILS

Unless otherwise noted, distillation uses a fixed 4096-example, one-epoch setup with batch size 16 and AdamW. The experiments also examine explanation margins and task-id association accuracy.

Distillation uses 4096 examples for 1 epoch with batch size 16 and AdamW unless otherwise noted.The learning rate is 1e-5 with TK-Instruct and 1e-4 with Incoder, with 0 weight decay.
Figure 5 compares In-context Margin, measuring explanation gains for the teacher, with Distillation Margin, measuring student gains after distillation.
Table 6 evaluates task-id association using correct-association accuracy and wrong-association accuracy across four task instructions.Mixed students generally successfully distill the task-id associations.

B.4 DISTILLING ABSTRACT STATEMENTS DETAILS

The abstract-statement experiments study natural-language explanations, task-id associations, and gradient-descent and long-context comparison settings. These experiments measure explanation transfer, association accuracy, and performance under alternative training configurations.

Learning from natural language explanations: The explanation experiment plots In-context Margin against Distillation Margin in Figure 5.
Distilling task id associations: The task-id association experiment reports detailed per-task results in Table 6.
Gradient descent: The gradient-descent baseline fits all training examples into one batch for 25 AdamW epochs and selects the epoch with the highest average exact set match score across databases.
Distilling long contexts: The long-context experiment averages y probabilities under 8 teacher prompts and estimates teacher performance using greedy decoding over an ensemble of 8 prompts.For each database, two students with different in-context examples are distilled and their performance is averaged.

B.6 DISTILLING STEP-BY-STEP REASONING DETAILS

The step-by-step reasoning experiments compare scratchpad-training baselines with transfer distillation from a teacher initialized to preserve prompt responsiveness. Evaluation uses unseen addition problems and mixed training data.

Distilling scratchpads with T5-small: Scratchpad baselines train for 1000 epochs, while Scratchpad then Direct trains 1000 epochs on scratchpads followed by 1000 epochs on direct answers.All use batch size 8, learning rate 3e-4, and AdamW, with performance reported on 10k unseen addition problems.
Transfering step-by-step reasoning: The transfer teacher is fine-tuned on 500 scratchpad examples mixed with 4096 Natural Instructions-V2 examples so it retains prompt responsiveness.
Transfering step-by-step reasoning: The transfer experiment initializes the teacher with scratchpad addition ability while mixing Natural Instructions-V2 data to prevent loss of its ability to respond to prompts.