Learning by Distilling Context

By Charles Burton Snell, Daniel Klein, Ruiqi Zhong. 2022

104 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence, Materials Chemistry, Representation & Self-Supervised Learning

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Materials Chemistry

In Representation & Self-Supervised Learning

Citations in the corpus

Large Language Models Can Self-Improve (2022, neutral): “Recent studies (Zelikman et al., 2022; Snell et al., 2022; Eisen-stein et al., 2022) show that dark knowledge within LLMs can be retrieved with more computation at inference time, such as adding informative instructions into the input sequence and output CoT generation (Wei et al., 2022c; Kojima et…”
MIMIC-IT: Multi-Modal In-Context Instruction Tuning (2023, background): “One potential reason for the zero-shot performance gain by instruction tuning is that it internalizes the context [40], which is preferred in user interactions especially when user input skips commonsense context.”
Distilling Reasoning Capabilities into Smaller Language Models (2022, background): “Snell et al. (2022) demonstrated the usefulness of providing instruction that can help models achieve better reasoning skills.”
Distilling Multi-Step Reasoning Capabilities of Large Language Models into Smaller Models via Semantic Decompositions (2022, background): “Snell et al. (2022) demonstrated the usefulness of providing instruction that can assist the models in achieving better reasoning abilities.”
Honest Students from Untrusted Teachers: Learning an Interpretable Question-Answering Pipeline from a Pretrained Language Model (2022, background): “Concurrent research uses chain-of-thought prompting in a student-teacher setup, similar to our architecture (Snell et al., 2022).”
Learning to Compress Prompts with Gist Tokens (2023, methodology): “2 This is precisely the approach taken by recent work (Askell et al., 2021; Choi et al., 2022; Snell et al., 2022), including Wingate et al.”
Adapting Language Models to Compress Contexts (2023, background): “Context distillation A related line of work considers context distillation (Askell et al., 2021; Snell et al., 2022), in which in-context information, e.”
The Instruction Hierarchy: Training LLMs to Prioritize Privileged Instructions (2024, methodology): “To generate training data, we leverage two principles: synthetic data generation and context distillation (Askell et al., 2021; Snell et al., 2022).”