Learning by Distilling Context
Authors: Charlie Victor Snell
Dan Klein
Ruiqi Zhong
Venue: Submitted to ICLR 2023 (2023)
Keywords: language models
NLP
prompting
distillation

Abstract: Language models significantly benefit from context tokens, such as prompts or scratchpads. They perform better when prompted with concrete training examples and abstract statements about the target task (instructions), and they acquire new capabilities to perform complex tasks by generating step-by-step reasoning (scratch-pad) before predicting the final answers. However, they do not internalize these performance gains, which disappear when the context tokens are gone. Consequently, we always need to pay extra computation for this gain, and it is unclear how to transfer the capabilities acquired by context tokens to other tasks, or how to leverage the context tokens when their length exceeds the context window size. Our work proposes to apply context distillation so that a language model can internalize these gains. Concretely, given an input for the target task, we let the model use all relevant context tokens to generate the output, using ``[instructions] + [task-input]'' to predict ``[scratch-pad] + [final answer]''; then we fine-tune the same model to predict the above ``[final answer]'' conditioned on the ``[task-input]'', without seeing the ``[instructions]'' or using the ``[scratch-pad]''. This incentivizes the model to behave as if the context were present, hence updating the parameters to internalize the context information. We show that context distillation can be used as a general method for learning. In particular, we demonstrate that context distillation can effectively internalize 3 types of contexts: 1) abstract task instructions and natural language explanations of why an output is correct or incorrect on Natural-Instructions-V2; 2) step-by-step reasoning on 8-digit addition questions, where we show the model can apply this newly acquired capability to downstream question answering tasks; and 3) concrete training examples on the SPIDER Text-to-SQL dataset, where context distillation outperforms directly learning with gradient descent by 7%.

Under review as a conference paper at ICLR 2023
LEARNING BY DISTILLING CONTEXT
Anonymous authors
Paper under double-blind review
ABSTRACT
Language models significantly benefit from context tokens, such as prompts or seratchpads. 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 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 internal- izing these gains. Concretely, given 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 predict its own "[final an- swer]" 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 types of training signals. First, it can internalize abstract task instructions along with explanations, so we can recursively update the model parameters with new instructions and overwrite old ones. Second, 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 multiple context distillation operations can internalize more training examples than what the context window size allows. Finally, we show preliminary evidence that it can internalize step-by-step reasoning on 8-digit addition, and such a newly acquired capability proves to be useful for other downstream tasks.
INTRODUCTION
Recent work has shown that language models significantly benefit from context tokens. When prompted with task definitions, language models can perform zero-shot learning (Wei et al. 2022a Sanh et al.] 2022), and the performance further improves with additional in-context examples and explanations (Chen et al. 2022 Scheurer et al. 2022). They also acquire the capability to perform more complex tasks by generating step-by-ste stesteppspppeeeocp reasoning in the context window before predicting the final answer (Nye et al. 2021b Wei et al. 2022b Zhou et al., 2022).
However, language models cannot internalize these performance gains, which disappear when the context tokens are gone. Consequently, we always need to pay extra computation for running infer- ence on context tokens; this is undesirable, as sometimes the task instructions and the ccratch-pad can be more than 10x longer than the actual task inputs. Furthermore, it is unclear how to leverage the context tokens when their total length exceeds the context window size. These shortcomings are analogous to how humans are slow at performing complex cognitive tasks (Wason & Evans 1974) and can hold only a limited amount of information the working memory (Baddeley) 1992).
Humans get around this by practicing. Consider, for example, learning to type your friends' phone numbers. The first few times you type it, you need to consciously recall the number using working memory and slowly decide which button to press. After repeatedly typing the same number, it be- comes a habit and you can type the number quickly without conscious reasoning. Through repeated practice, the knowledge of your friend's 's phone number is "distilled" into your muscle memories ! This mechanism for distilling knowledge critical for learning complex tasks because allows us
See declarative learning procedural learning for a friendly but more in-depth discussion. https:
7/en. wiki pedia www.or//wwkk... becl arat inpative__eeee ve_learning

---

Under review as a conference paper at ICLR 2023
Model Input
Model
Sample
Completions
Teacher
Eesatrsction wish Extra Informatiion
[Bocaschopad
[Rav Inpot)
Row Task Input
Fieedne topredkt
Student
[Minina] Instroction) Taxk Isput]
[Einel onawee]
Teacher's Template
Teacher's Output
outputring your
Teacher
Exaeple Ootguts D ab 4 L. the final, asawer is 4
Exseple Toputs
Oz
Conpletion:
Sample
>>
The
tinal assveE Se
X 2
Bow odd she following twD suxbecs.
Inpat:
Random
Addition
Expressions
Template Tinart
Student's Output fly)
Pactora the folleving caleslation.
Finetone preditt
Student
Inpatt
1 +
2
teacherss warpe))
Figure 1: Top. An overview our context distillation framework. We sample a raw task input, form the teacher's prompt by pre-prending detailed instruction that might contain more examples and explanations, and ask the language model to conditionally sample a scratch-pad and final answer. Then we fine-tune the same language model to directly predict the final answer with a minimal instruction. We formalize this framework mathematically in Section Bottom. An instantiation of our framework that internalizes step-by-step reasoning for 8-digit addition.
to incrementally build up our knowledge and skills, so that we can learn to accomplish increasingly complex tasks.
We propose to apply similar method, context distillation, to fine-tune language models. For exam- ple, as shown in Figure [ to make language models internalize the step-by-step addition capability, we first synthesize large number of "practice" addition questions; we then ask the model to follow the more informative instruction to reason step-by-step before generating the target answer; finally, fine-tune the language model to directly predict the answer conditioned on simpler student prompt. As a result, by practicing on a lot of addition problems, the ability to add distilled into its parameters. We formally state our generalized context distillation framework in Section
Section 3 shows that we can apply context distillation to a wide range of settings: learning from abstract statements, learning from concrete examples, and learning from step-by-step reasoning. Section 3.1 (Figure ( shows that context distillation can effectively internalize task instructions along with natural language explanations from natural-Instructions ditionally, we can teach the student to associate numerical indices with certain tasks, and then we can recursively re-assign these task indices, overwriting the student's past associations. Section 3.2 (Figure 3) shows that context distillation can be used to internalize Text-to-SQL training examples from the SPIDER dataset (Yu et al.) 2018 into Incoder (Fried al. 2022). and it outperforms di- rectly learning with gradient descent by 9% for 8-shot adaptation; additionally, we show that as we distill more training examples than can in the context window, we observe continual improve- internalize step-by-ster reasoning to perform 8-digit addition, and such a capability can transfer to downstream question answering tasks; hope this preliminary results generalize to more complex and realistic tasks that larger models can perform with chain-or thain-of-tee ofinn-of-hhuut ttouughtt reasoning (Wei et al.) 2022b Zhou et al. 2022).
(Wang et al.) (022b): ad-
ments in performance. Section 3.3 (Figure shows that
Overall, context distillation demonstrates promising potential as general method for learning, As discussed in Section 4 we predict that future models will be better able to learn from context than today's 's models, and researchers will these models tackle increasingly complex tasks that
2

---

Under review as a conference paper at ICLR 2023
require more background knowledge and longer reasoning chains. Therefore, anticipate method to be increasingly useful the future.
2 CONTEXT DISTILLATION
We introduce the main components and the intuition of our context distillation framework in Section 2.1 describe our algorithm for single round distillation in Section explain how to distill multiple contexts recursively simultaneously Section 2.3 and describe various implementation details to make it efficient and stable in Section
2.1 INTUITION AND MAIN COMPONENTS
We explain our method by contrasting it with the classical distillation methods (Hinton et al. 2a) These classical methods ask the teacher model with parameter TEACHER to generate label for given imput z, and train the student @STUDENT to mimic the teacher by predicting y conditioned Typically, θπAссER * OSTUDENT when the algorithm starts, and the distillation process is driven by the difference between their parameters. In contrast, under context distillation, θTEaChER STUDENT when the training starts, and the distillation process instead driven by the differences in the x and y that they see and predict.
æ.
To design such a difference that drives the distillation process, our framework requires the model developers to provide four components: a student template TSTUDENT, and an answer extractor f. We introduce them below.
task input distribution D, teacher template TTeacher
Raw Task Input Distribution D. DD is distribution of strings, which are typically the "core" inputs of the target task of interest. For example, if the target task to classify movie review sentiment, the input distribution could be defined as random movie reviews. Generally, there are many ways to define a raw task input distribution: we can use rule-based method to generate random strings, sample from pool of unlabeled data, or conditionally sample from a language model. We explicitly distinguish raw task inputs from the whole "input prompt" to the language model, which will be obtained after applying the templates below.
Except for the synthetic experiments in Section 3.3 all experiments our paper create via few- shot prompting and do not assume access to pool of unlabeled data.
Teacher Template TTeACheR. TTeaCher is mapping from strings to strings, which transforms raw task inputs to the input prompts for the teacher model. As shown in Figure [ the teacher template usually contains more detailed instructions, explanations, and examples about the task.
Student Template TSTUDENT- TSTUDENT is mapping from strings to strings, which transforms raw task inputs to the input prompts for the student model. As shown in Figure this template usually still contains minimal information about the task so that the request in the prompt not under- specified. However, compared to the teacher prompt, it incorporates far fewer explanations and training examples the task, and such difference transfers this useful context information into the student parameters.
Answer Extractor f. ff is a mapping from token sequences to token sequences, which extracts the final answer (a sub-sequence of tokens) from the full teachers' generation. As shown in Figure [ strips away the intermediate reasoning process, and the students need to internalize the hep-by-step reasoning process directly predict what the teacher predicts at the end.
We now describe context distillation formally using the mathematical terms we just introduced.
2.2 FORMAL DESCRIPTION
Our algorithm first samples ø from D and ask the language model to sample completion conditioned on TTeAcher (z); we then fine-tune the language model to predict f(y) conditioned on TSTUDENT(2)) (z)). Throughout the distillation process, OteaCheR is fixed.
3

---

Under review as a conference paper at ICLR 2023
Formally, let OSTUDENT and TEACHER be the parameters of language model, and define (0(||RRMPPT] to be the probability distribution of the completions conditioned on the prompt. We optimize
D..Tvraasti Cannsseriiaaee P 1.8rssmux (SSTUDENT) = Evz/D (Ey mosmmn (TtzacuR)) [lg Po (f(y)) STTUDENT (x))
(1)
Notice that the definition of & depends on five variables D, TTTUDENT TEACHEr f. and θeAChhR To keep the notation uncluttered, will only include the necessary subscripts if the rest be inferred from the surrounding text.
2.3 COMBINING MULTIPLE UPDATES
We now introduce simultaneous distillation and recursive distillation, two straightforward variants that combine multiple context distillation operations, allowing us to internalize ensembles or se- quences of contexts, enabling a form of learning that is not possible with just a single prompt.
Simultaneous Distillation. To simultaneously perform K different context distillation operations represented by DI... TEACIEK/STUDENT.L Ji.... Ji..k,
optimize the total loss:
K
LTOTAL
LD T.
(2)
Simultaneous distillation is especially useful when the prompts contain independent instructions and in-context training examples, but their total length exceeds the language model context window size.
Recursive Distillation. To perform R rounds of context distillation operations recursively, inductively define 0 as the initial language model, and 0k+1 to be the parameters after fine-tuning with the loss function
:= DD. TruneTTn.r... TRAAEER..
(3)
with 0, as the initialization. Essentially, for each round of context distillation we can obtain new student, which then becomes the teacher and the student's initialization for the next round. This allows language model to recursively update itself via context distillation? 2
2.4 IMPLEMENTATION DETAILS
naive method to optimize Equation I is to sample y and directly fine-tune the student to predict the hard label of [y): such a method wastes the token logit information and results in noisy gradi- ents. Instead, we minimize the token-level KI. divergence between the student and the teacher (i.e., learn from the soft labels). However, this leads to another issue: the vocabulary space of language models is often on the order of 50k, and the full soft labels consume as much as 200KB memory per decoding step. Therefore, approximate the it with empirical distribution of 100 token samples. Such a technique saves memory by a factor of 500. Appendix 8.1 shows more details.
EXPERIMENTS
Overview. We apply context distillation three types of settings. Section 3.1 applies context dis- tillation to internalize abstract instructions and natural language explanations; additionally, language models can recursively update themselves via context distillation to associate each task instruction with a task id. Section 3.2 applies context distillation to internalize concrete training examples, and it can outperform directly learning with gradient descent on the SPIDER Text-to-SQL dataset; ad- ditionally, we show that simultaneous distillation can be used to internalize more training examples than the context window can fit. In Section 3.33 we apply context distillation internalize the ability to perform 8-digit addition step-by-step additionally, such a capability can transfer to other down- stream question-answeren
tasks even when the scratch-pad is not present. In all cases, the student's input length is significantly shorter than the teacher's, and we may reduce the context length by as much as times, saving us a large amount compute at inference time.
Choi al. 2022) first proposed this recursive formulation but only evaluated it qualitatively.
4

---

Under review as a conference paper at ICLR 2023
will aovic
Samps
Prompt Generate Raw Task Input
ouspus:
total vasta of Lime.
Righly recomsended.
Exeppi teput:
Exaeple Outpus:
You need to classity zevieva. zeans good, nenne not good.
Seample Input:
impressed by the plnt betore
Teacher's Template
Chaptez 130, but everything betveen 131-139 trass.
Sxample Cusput:
Comphtioo
Explanation: anavce Khen it in anbiyuous.
clangity following review,
Ispus: have mixed realing.
Outguct
lapredet
Student's
have mixeed feellag.
[Rask Index *0")
Template-A
Fie-ti
Student's
Template-8 B
have milsed fealing.
explanations.]
[Iaxtruction Nithout
We use context distillation to internalize abstract task instructions and associate each of them with task id. For example, after distilling the context with the teacher's template and student template-A, the student should perform sentiment classification whenever it sees the index "[O]". We perform additional ablation study by using the student template-B without the explanation verify that context distillation can be used to learn from natural language explanations. The raw task inputs are sampled from the same teacher model via few-shot prompting (top).
Figure 2:
Baselines and Oracles. At the beginning of the distillation, the teacher and the student networks have identical parameter values. However, the teacher has an advantage because it has access to additional context while the student does not. The student's goal is to mimic the teacher to inter- nalize the teacher's context. Therefore, distillation to support the claim that context distillation is successful. The teacher's performance is generally an upper bound on the student's performance.
student's performance needs to improve after context
Naming Conventions. We use the following names when reporting the results.
"Teacher" initial language model's (LM) performance with additional context.
"Pre-distill Student" the initial LM's performance without additional context.
"Post-distill Student" the distilled LM's performance without additional context.
Note that whether system called teacher or student is not based on its parameter values (c.g..
initial vs. distilled) but on whether it uses context.
3.1 LEARNING FROM ABSTRACT INSTRUCTIONS AND EXPLANATIONS
We apply context distillation to internalize abstract instructions. In all experiments, unless men- tioned otherwise, the answer extractor is the identity function and we use few-shot prompting to sample raw task inputs (see Figure]
Dataset and Language Models. We use Natural--Insrrutionss V2 for all the experiments on this section. Natural Instructions a dataset of 1600+ diverse language tasks, where each task is asso- ciated with natural language task description, input-output examples, and explanations about why certain outputs are correct or incorrect (Wang et al. 2022b). We trained our teacher language model (TK- Ins-tustuuct) by fine-tuning LM-adapted T5-1IB (Raffel et al., 2020) its training set, and the details can seen in Appendix 8.5 For evaluation, we use their dataset's s official metric, Rouge-L, calculate the performance averaged across the 10 tasks selected.
5

---

Under review as a conference paper at ICLR 2023
Design Choice. We select tasks from their evaluation split where the teacher can most signifi- cantly improve the performance when prompted with extra natural language explanations compared to task description only, and 5 where they improve the least. We chose these tasks because they also allow us to investigate how well context distillation can internalize natural language explanations in Appendix 8.2 Junde different teacher's performance.
Hypothesis 1: context distillation can internalize abstract task instructions. To test this hypoth- esis, we defined the student template as the identity mapping, i.e., the student only sees the raw task input, while the teacher template contains the task instruction, which consists of a task description, 2 positive in-context examples, and 2- 2-negative in- in-context examples (Figure SigureZteacher's teacher's template). The teacher's ' performance is 43.4 Rouge-L, establishing an upper bound for the student. Before context distillation, the student's 's performance 9.0, since does not know what task it should perform. Af- ter context distillation, the student's ss performance significantly increases to 34.7. Context distillation successfully internalized abstract task instructions. Finally, we used 11.1 times fewer inference time tokens when evaluating the student verses the teacher.
In light of emerging trends to learn from from natural language explanations (Scheurer et al. 2022
Lampinen et al. 2022), we present experiments on distilling in-context explanations in Appendix 8.2
Hypothesis 2: recursive distillation can overwrite past updates. We study recursive distilla- tion using four classification tasks, auperglue_copa_tex. cop..text.compllioo meta_ www.task.classiiiccc.... wweetqa.casssifice
and rocstories title classification. We define a new task id association chal- lenge: each task is associated with index, and when the student model sees index, it needs to perform the corresponding task without looking at its instruction; we train the student model to do this via simultaneous context distillation, where the student sees the task index while the teacher sees the task instruction (Figure 2 student template A). After we use context distillation to train the student to associate each task-id with the corresponding instruction, we shuffle the task-id associ- ation, perform context distillation again, and then evaluate the model's performance on the newly defined association to measure the ability of context distillation to overwrite previous updates.
We define two metrics for this task id association challenge. First, we will measure the average accuracy (correct association accuracy) of the student on the four tasks when they are prompted with the corresponding index. Second, it is plausible that the student can learn to associate the task input distributions, rather than the index, with the corresponding instructions. For example, suppose that id "[0]" corresponds to classifying sentiment and the task input distribution is movie reviews, while "[1]" corresponds to whether it is sports related and the raw input distribution is news articles, then the student might cheat by always classifying sentiment whenever it sees movie review, regardless of what id it sees. Therefore, we also measure the average accuracy when we prompt the model with the wrong task id, and want this number to be low (wrong association accuracy): for example, the model sees "[O]" and a news articles, it should have low accuracy at classifying whether it corresponds to sports, because "[O]" corresponds sentiment classification.
We experimented with two variants of simultaneous distillation: 1) the "naive" one, where each Dr contains only the input distribution for one single task, and 2) the "mixed" one, where we define each Dj as mixture of all the raw task input distributions. As shown in Table under the "naïve" input distribution, the student model can cheat by associating the task input distribution, rather than the task id, with the task it needs to perform; on the other hand, the "mixed" variant successfully writes the past task id association.
Can the language model recursively update itself based on new task-id associations? It not obvious that it would work, since it plausible that the model forgets how to follow the instruction properly after being fine-tuned to associate task with Ids. We recursively apply context distillation for three rounds and find in Table T right that the performance does not drop. Along with findings from (Ramasesh al.) 2022). these results suggest that catastrophic forgetting might pose little trouble to recursively applying context distillation on larger instruction -unnidd models.
3.2 LEARNING FROM CONCRETE EXAMPLES
We show that context distillation can internalize training examples, and therefore can be potential alternative to directly learning from these examples with gradient descent; additionally, simultancous distillation allows the model to learn from more in-context examples when their total length exceeds

---

Under review as a conference paper at ICLR 2023
Table I: We apply context distillation to train the student to associate task instructions with task ids, and then recursively apply context distillation repeatedly to overwrite the previous associa- tion. We report the correct association accuracy ("correct") and the wrong association accuracy score ("wrong") over all four tasks we associate. For all entries we report the 95% confidence in- terval. Left: Performance significantly improves after context distillation, and mixed distillation successfully prevents the model from cheating by memorizing the input distribution. Right: Task-id
association performance does not drop after several rounds recursive updates.
model
correct
wrong
model
correct 1
Teacher
81 +
Round r
68 +
Pre-distill Student
49 i
48 5
Round 2
66 + 5
"Naive" Post-distill Student (r = 2)
68 5
61 5
Round r = 3
68 + 5
"Mixed" Post-distill Student (r 2)
66 +
16 4
Cename cracc proplectiie shiiing AAE 1steperiy achona
Beee?
Now old
Eow
[izampla Zopat
pesngent
[Gchamo]
gle
Teacher's
BEeapte Seport Derpot
Teacher's Template
Template
SAAE *Eres*
Student's
SE0202 All people
Student's
Teacher's
Template-A
urane 1930 "Brex*
Template-
TemplatesA
The teacher template A contains additional training examples 1-4 compared to the student template A. and context-distilingg them outperforms directly learning from them with gradient de- scent. Additionally, by simultaneously applying the teacher's templates A and B (Section 2.3 more in-context examples (5-8) can be distilled, even though their total length might exceed the context window size. The raw task inputs are sampled from the same teacher model via few-shot prompting (top).
Figure 3:
the context window length. In all experiments, the answer extractor is the identity function and we use few-shot prompting to sample raw task inputs (Figure 3).
Dataset and Language Models. We use the SPIDER text-to-SQL. dataset Yu al. 2018) for our experiments. Each database in SPIDER is associated with schema and list of English questions annotated with SQL query. The task is to predict the SQL query given the schema and the question. For the teacher language model, we chose Incoder-6.7B (Fried et al. (202),, which is pre-trained on code. For each experiment, we randomly sampled eight text-to-SQL pairs for each database as in-context examples and evaluate the rest by calculating the exact set match accuracy.
Design Choice. We chose this task because prior work (Rajkumar et al. 2022 has demonstrated that SOTA models are able to effectively utilize in-context examples for this task; therefore, our approach is more likely to outperform gradient descent. More generally, in-context learning perfor- mance improves when we scale up (Chen et al., 20222 Brown et al.) 2020), and hence we conjecture that context distillation can outperform gradient descent more tasks for future models.
Hypothesis 3: context distillation sometimes outperform directly learning with gradient de- scent. We define the student template to be database schema followed directly by the question, whereas the teacher template contains the database schema followed by four in-context examples (see Figure 3 teacher's 's template-A), As we see in Table 2 context distillation outperforms learning via gradient descent four examples by 8.6% in exact-set match accuracy, margin which further increases by 0.4% when training on eight examples. Therefore, context distillation can be used as an alternative to directly learning from training examples with gradient descent.
7

---

Under review as a conference paper at ICLR 2023
Model
Examples 27.7
Examples
Teacher
28.2
Pre-distill Student
0.3
0.3
Post-distill Student
22.1
27.9
Direct Gradient Descent
13.4
18.9
Table 2: Comparing context distillation to gradient descent on the SPIDER text-to-SQL validation set. We see that context distillation outperforms directly learning via gradient descent on four exam- ples by 8.6% (exact set match); this margin further increases when learning from eight examples.
Hypothesis 4: context distillation enables learning from more training examples when their total length exceeds the context window size. For this experiment, we select four databases from SPIDER training set which have particularly long database schema, such that it is possible to fit four training examples into Incoder's context window but not eight. Therefore, we include 4 training examples in the student template (Figure 3 student's template B.). which would lead to the best in-context learning performance ex-ante given the context window size. To internalize more training examples, we perform simultaneous distillation and sample teacher templates by including random subset of four in-context examples out of eight (Figure [ teacher's template A and B). After context distillation, our student achieves an exact set match of 16.2 + 0.6, which improves over the pre-distillation student performance of 14.22 + 0.8, hence confirming our hypothesis.
3.3 PRELIMINARY EVIDENCE FOR LEARNING FROM STEP-BY- REASONING
We show that context distillation can effectively internalize skills acquired by generating step-by- step reasoning, and such a skill can benefit the model for other downstream tasks. Unless otherwise mentioned, in this section we use an f that will extract the final answer from the teacher's full output, as shown in Figure bottom.
Dataset. We define the input distribution D to be uniform over all possible through digit addition questions, identical those used by Nye et al. (2021a). We report addition accuracy for each experiment. We hope that these results can generalize to more complex and realistic tasks that larger models can perform with chain-of tf-thughhss reasoning (Wei et al. 2022b Zhou et al. 2022).
Hypothesis 5: context distillation internalize step-by-step reasoning. To test this hypoth- esis, we obtain a teacher model that can use scratch-pad to perform step-by-step reasoning by fine- tuning the I.M-aaptedd T5-small dataset of 500 addition expressions, where the model needs to first generate a scratch-pad before predicting the final answer. We then perform context-distillation with an that extracts the final answer from the teacher's output as shown in Figure bottom. As shown in Table 3] after distillation, the ability of the student to perform direct additions (with- using scratch-pad) improves from 0% 94.7%, implying that context distillation internalizes step-by-step reasoning.
We compare this several other transfer learning and multi-task learning baselines that use the same amount of training data in Table 3 Under transfer learning, we first fine-tune the student model to predict the scratch pad, and then fine-tune to directly predict the final answer. Under multi- task learning, we fine-tune the model to predict scratch-pad and the final answer independently. Both variants perform significantly worse (> 20%) than context distillation. In addition to the gains in reasoning ability over baselines, we also saved inference time compute in our evaluations: specifically we used 8.0 times fewer tokens at inference time when evaluating the student compared to the teacher.
Hypothesis 6: the reasoning cability internalized by scratchpad distillation can transfer to other related reasoning tasks. To test this hypothesis, we distill the addition scratch-pads on our IK-Instruct model, and evaluate capability transfer to other tasks. We obtain the teacher language model by fine-tuning TK-Instruct on a distribution of both natural instructions data and 500 addition scratchpad examples (see Appendix 8.9). For the context distillation training, we define the student template to be a description of the addition task followed by two direct answer in-context exam- ples, and similarly the teacher template contains the task description and two in-context examples
8

---

Under review as a conference paper at ICLR 2023
Teach
Pre-Dist
Post-Dist
Sc-+Dir
Sc+Dir
Digit Addition Accuracy %
93
0
95
72
61
Distilling addition scratchpads T5-small. "Teach" refers to the teacher LM's per- formance using scratch-pad. "Pre-Dist" refers to the student's performance before distillation; "Post-Dist" refers to the student's performance of direct addition (without scratch-pad) after refers to our transfer/multi-rask learning baseline. Context Dis- tillation performs the best for direct addition.
Table 3:
with scratch-pad answer. To prevent the student from catastrophically forgetting its in-context learn- ing ability during distillation, we mix our 10k distillation data-points with a distribution of 65536 randomly selected examples from Natural Instruction-V?
The student's ss accuracy on directly answering 8-digit addition questions improves from 1% to 17%, while the student's performance Natural Instructions remains roughly the same (from Rougel. 57 before distillation to 58 after), implying that the student did not lose its original capabilities to follow instructions. Additionally, we evaluate the student's performance on of related reasoning tasks. We then use a template to synthesize simple questions that require knowledge to add two numbers, for example, "A has turkies. B has turkies. How many turkies do they have altogether?". On this synthetic dataset, the student's performance significantly increases from 17% to 30% after context distillation, implying that the capability perform direct addition can transfer to other related applications.
4 RELATED WORK
Prompting and Instruction Tuning. Many recent works show that language models can learn from abstract task definitions (Zhong al. 202T Mishra al., 2022, Wei et al. 2022a), natural language ex explanations (Scheurer et 2002 Wang et al.) 2022b). and concrete in-context examples (Min et al. 2022; Chen et 2002)) We anticipate the in-context learning performance to improve further in the future (Kaplan et al. 2020), thus increasing the upper bound of what context distillation can achieve.
Scratch Pad. Many recent works show that language models perform better when it is required to generate a chain of reasoning steps before outputting the final answer (Zhou al.. 2021 Nye et al., [2021b] [obbe et al., 2021; Wei et al., [202b)] Zhou et al.) 2022; Lewkowycz et al., (022). We anticipate context distillation to be increasingly useful, as the research community starts to tackle more difficult problems, which require more sophisticated skills and have longer problem descriptions and reasoning chains.
Distillation. There has been large literature on distilling knowledge in a neural network (Hinton
al. [0155b] Adriana al. 2015 Liu al. 2019 Yang et al. 2020 Xie al. 2020). Most related to our work, different sub-variants of context distillation have been independently discovered by different researchers. Waan et al. (00T) emphasizes the aspect of creating a dataset without any human annotations and uses a language model to generate task inputs and their labels. Choi et al. 2002 ; Askell et al. (2021) formulated the method context distillation (also referred to as prompt injection), which distills fixed input prompt; their method is special case of our framework with an identity student template and an identity output selector. Additionally, they focused more on the benefit of saving computational resources, while considered it as general learning method.
5 CONCLUSION
We present context distillation general learning method, which can internalize abstract state- ments, concrete examples, and step-by-step reasoning. Given that 1) is general and delivers strong performance, 2) future models will have stronger in-context learning capability, and 3) future tasks will have longer descriptions and reasoning chains, we anticipate our methods to be increasingly useful in the future.

---

Under review as a conference paper at ICLR 2023
ETHICS STATEMENT
Our method is still far from perfect for internalizing context information and should not be applied high stake scenarios.
REPRODUCIBILITY
We include more experiment details in the appendix and release code for reproducing all experi- ments.
REFERENCES
Romero Adriana, Ballas Nicolas, K Samira Ebrahimi, Chassang Antoine, Gatta Carlo, and Yoshua. Fitnets: Hints for thin deep nets. Proc. ICLR, 2, 2015.
Amanda Askell, Yuntao Bai, Anna Chen, Dawn Drain, Deep Ganguli, Tom Henighan, Andy Jones, Nicholas Joseph, Ben Mann, Nova DasSarma, Nelson Elhage, Zac Hatfield-Dodds, Danny Her- nandez, Jackson Kernion, Kamal Ndousse, Catherine Olsson, Dario Amodei, Tom Brown, Jack Clark, Sam McCandlish, Chris Olah, and Jared Kaplan. A general language assistant as labora- tory for alignment, 2021. URL https: www.aaavvv..... http://twwwxxxx.aaaaooo// http:///aaaa/a/.......... 0002
Alan Baddeley. Working memory. Science, 255(5044):556555)) 1992.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal
Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman- and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http:
//github. com/ commgoogle//aa.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot leamners Advances in neural information processing systems, 33:1877-1901, 2020.
Yanda Chen, Ruiqi Zhong, Sheng Zha, George Karypis, and He He. Meta-learning via lan- guage model in-context tuning. In Proceedings of the 60th Anmual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers). pp. 719-730, Dublin, Ireland, May 2022. Association for Computational Linguistics. doi: 10. 10.186555/1120222al--ooogg https://acllantho www.aatthhoooyyy..oooooo wwtpoo.............///o/ www.22222.cllonnog..
URL
Eunbi Choi, Yongrae Jo, Joel Jang, and Minjoon Seo. Prompt injection: Parameterization of fixed inputs, 2022. URL Intpp::////// //tp///aavv/////////a////a//aaaa http//aawwrraa/aaa///////a/aa/aa
11349.
Karl Cobbe, Vineet Kosaraju, Mohammad
Bavarian, Jacob Hilton, Relichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems. arXiv preprint arXiv:2/ 14168, 2021.
Nicola De Cao, Wilker Aziz, and Ivan Titov. Editing factual knowledge language models, 2021.
URL https: wwwp://////aba.oooooooooo.se 2204. 081684
Daniel Fried, Armen Aghajanyan, Jessy Lin, Sida Wang, Eric Wallace, Freda Shi, Ruiqi Zhong, Wen-tau Yih, Luke Zettlemoyer, and Mike Lewis. Incoder: A generative model for code infilling and synthesis. arXiv preprint arXiv:2204. 05999, 2022.
Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in neural network, 2015a. URL
02531
Geoffrey Hinton, Oriol Vinyals, Jeff Dean, et al. Distilling the knowledge in a neural network. arXiv preprint
2(7), 2015b.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arX.v.2001.08361, 2020.
10

---

Under review as a conference paper at ICLR 2023
Andrew K Lampinen. Ishita Dasgupta, Stephanie CY Chan, Kory Matthewson, Michael Henry Tessler, Antonia Creswell, James McClelland, Jane X Wang. and Felix Hill. Can language models learn from explanations in context? arXiv preprint arXiv:2204.02329, 2022.
Aitor Lewkowyez, Anders Andreassen, David Dohan, Ethan Dyer, Henryk Michalewski, Vinay Ra- masesh, Ambrose Slone, Cem Anil, Imanol Schlag, Theo Gutman- et al. Solving quantitative reasoning problems with language models. arXiv preprint arXiv:2206. 14858, 2022.
Xiaodong Liu, Pengcheng He, Weizhu Chen, and Jianfeng Gao. Improving multi-task deep neu- ral networks via knowledge distillation for natural language understanding. arXiv preprint arXiv:1904.09448 094822 2019.
Andrew L. Maas, Raymond E. Daly, Peter Pham, Dan Huang, Andrew Y. Ng. and Christopher Potts. Learning word vectors for sentiment analysis. In Proceedings of the 49th Annual Meeting of the Association for Computational Linguistics: Human Language Technologies, pp. 142-150, Portland, Oregon, USA, June 2011. Association for Computational Linguistics. URL https:
//aclanthology. wtt.l/oyy.rrgg.1l11o crt/p11-11011
Kevin Meng, David Bau, Alex Andonian, and Yonatan Belinkov. Locating and editing factual associations in gpt, 2022. URL
httg:/ab/222000 2202. 05262
Sewon Min, Mike Lewis. Luke Zettlemoyer, and Hannaneh Hajishirzi. MetalCI.: Learning to learn in context. In Proceedings of the 2022 Conference of the North American Chapter of the Asso-
ciation for Computational Linguistics: Human Language Technologies, pp. 2791-2809, Seattle, United States, July 2022. Association for Computational Linguistics. doi: 10.18653/v 1/2022. naacl-main.201 201. [ httph: //ta//aaaaaatooooooo org.2ogggyyoog2222ooo..// wtt. naacl main.
Swaroop Mishra, Daniel Khashabi, Chitta Baral, and Hannaneh Hajishirzi. Cross-task generalization via natural language crowdsourcing instructions. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume Long Papers), pp. 3470-3487, Dublin, Ireland, May 2022. Association for Computational Linguistics. doi: 10. 244. URL http:s/s///cannooooooooo//////22 wwtp://2222...11ooogg. 244
Eric Mitchell, Charles Lin, Antoine Bosselut, Chelsea Finn, and Christopher D. Manning. Fast model editing at scale, 2021. URL http:////aaaiiioooraa///////////
Maxwell Nye, Anders Johan Andreassen, Guy Gur-Ari, Henryk Michalewski, Jacob Austin, David Bieber, David Dohan, Aitor Lewkowyez, Maarten Bosma, David Luan, Charles Sutton, and Au- gustus Odena. Show your work: Scratchpads for intermediate computation with language models, 2021a. URL
00114
Maxwell Nye, Anders Johan Andreassen, Guy Gur-Ari, Henryk Michalewski, Jacob Austin, David Bieber, David Dohan, Aitor Lewkowycz, Maarten Bosma, David Luan, Show your work: Scratchpads for intermediate computation with language models. arXiv preprint arXiw:2112.00114. 2021b.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer, 2019. URL.
www.oorgggbbbga. abl./a/...1.........../1..111111 10683.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, Peter J Liu, et al. Exploring the limits of transfer learning with a unified text-to-text transformer. Mach. Learn. Res., 21(110):1 2020.
Nitarshan Rajkumar, Raymond Li. and Dzmitry Bahdanau. Evaluating the text-to-sql capabilities of large language models. arXiv preprint arXiv:2204.00455 2022.
Vinay Venkatesh Ramasesh, Aitor Lewkowycz, and Ethan Dyer. Effect of scale on catastrophic forgetting neural networks. In International Conference on Learning Representations, 2022.
URL
wwwparevieweneeeeeeee mmet//oormmicn GhVS8 yPeEa
Amit Sabne. Xla Compiling machine learning for peak performance, 2020.
11

---

Under review as a conference paper at ICLR 2023
Victor Sanh, Albert Webson, Colin Raffel, Stephen Bach, Lintang Sutawika, Zaid Alyafeai, Antoine Chaffin, Arnaud Stiegler, Arun Raja, Manan Dey, M Saiful Bari, Canwen Xu, Urmish Thakker, Shanya Sharma Sharma, Eliza Szczechla, Taewoon Kim, Gunjan Chhablani, Nihal Nayak, De- bajyoti Datta, Jonathan Chang, Mike Tian-Jian Jiang. Han Wang, Matteo Manica, Sheng Shen, Zheng Xin Yong. Harshit Pandey, Rachel Bawden, Thomas Wang. Trishala Neeraj, Jos Rozen, Abheesht Sharma, Andrea Santilli, Thibault Fevry, Jason Alan Fries, Ryan Teehan, Teven Le Scao, Stella Biderman, Leo Gao, Thomas Wolf, and Alexander M Rush. Multitask prompted training enables zero-shot task generalization. In International Conference on Learning Repre- sentations, 2022. URL https: httpp:://pperevieececc wttp:reeiiieeetee//teeeeee wow.t///oooooodd///a
Jérémy Scheurer, Jon Ander Campos, Jun Shern Chan, Angelica Chen, Kyunghyun Cho, and Ethan Perez. Learning from natural language feedback. In ACL Workshop Learning with Natural Language Supervision, 2022.
Yizhong Wang, Swaroop Mishra, Pegah Alipoormolabashi, Yeganch Kordi, Amirreza Mirzaci, An- jana Arunkumar, Arjum Ashok, Arut Selvan Dhanasekaran, Atharva Naik, David Stap, Eshaan Pathak, Giannis Karamanolakis, Haizhi Gary Lai, Ishan Purohit, Ishani Mondal, Jacob Anderson, Kirby Kuznia, Krima Doshi, Maitreya Patel, Kuntal Kumar Pal, Mehrad Moradshahi, Mihir Par- mar, Mirali Purohit, Neeraj Varshney, Phani Rohitha Kaza, Pulkit Verma, Ravsehaj Singh Puri, Rushang Karia, Shailaja Keyur Sampat, Savan Doshi, Siddhartha Mishra, Sujan Reddy, Sumanta Patro. Tanay Dixit, Xudong Shen, Chitta Baral, Yejin Choi. Noah A. Smith, Hannanch Hajishirzi, and Daniel Khashabi. Benchmarking generalization via in-context instructions on 1,600+ lan- guage tasks, 2022a. htRtsttss://att
org/ 2tt.bb///2220.00005555
Yizhong Wang. Swaroop Mishra, Pegah Alipoormolabashi, Yeganeh Kordi, Amirreza Mirzaei, An- jana Arunkumar, Arjun Ashok, Arut Selvan Dhanasekaran, Atharva Naik, David Stap, et al. Benchmarking generalization via in-context instructions on 1,600+ language tasks. arXiv preprint arXiv:2204.07705. 2022b.
Zirui Wang, Adams Wei Yu, Orhan Firat, and Yuan Cao. Towards zero-label language learning. arXiv preprint arXiv:2 100.00919, 2021.
Peter C Wason and J St BT Evans. Dual processes in reasoning? Cognition, 3(2):141-15, 1974.
Jason Wei, Maarten Bosma, Vincent Zhao, Kelvin Guu, Adams Wei Yu, Brian Lester, Nan Du, Andrew M. Dai, and Quoe V Le. Finetuned language models are zero-shot learners. In Interna- tional Conference on Learning Representations, 2022a. URL htttps://epenreeiieee
forum?i d=g.2..
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Ed Chi, Quoc Le, and Denny Zhou. Chain of thought prompting elicits reasoning in large language models. arXiv preprint MXX::2201.11903 2022b.
Qizhe Xie, Minh-Thaan Thang Luong, Eduard Hovy, and Quoe V. Le. Self-training with noisy student improves imagenet classification. In 2020 JÉEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 10684-10695 2020. doi: 10. 10119900CRR2220000200000000
Ze Yang, Linjun Shou, Ming Gong, Wutao Lin, and Daxin Jiang. Model compression with two- stage multi-teacher knowledge distillation for web question answering system. In Proceedings of the 13th International Conference on Web Search and Data Mining, pp. 690-698, 2020.
Tao Yu, Rui Zhang, Kai Yang, Michihiro Yasunaga Dongxu Wang, Zifan Li, James Ma, Irene Li, Qingning Yao. Shanelle Roman, Zilin Zhang, and Dragomir Radev. Spider: A large-scale human-labeled dataset for complex and cross-domain semantic parsing and text-to-SQL task. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, pp. 3911-3921, Brussels, Belgium, guistics. doi: 10.18653/v1/DI8. 1-251--111.......
2018. Association for Computational Lin- wwwp:yyy.rgDD--1122..0
Ruiqi Zhong, Kristy Lee, Zheng Zhang, and Dan Klein. Adapting language models for zero- shot learning by meta-tuning on dataset and prompt collections. In Findings of the Associa- tion for Computational Linguistics: EMNLP 2021, pp. 2856-2878, Punta Cana, Dominican Re- public, November 2021. Association for Computational Linguisties. doi: 10. 10.18653/w1/2011
12

---

Under review as a conference paper at ICLR 2023
findings-emnlp...
URL
httes://eclaattoooooy.oo//2222.2 wtw.:/2eee..dddnoeeemeee....p
244
Denny Zhou, Nathanael Schärli, Le Hou, Jason Wei, Nathan Scales, Xuezhi Wang, Dale Schu- urmans, Olivier Bousquet, Quoc Le, and Ed Chi. Least-to-most prompting enables complex reasoning in large language models. arXiv preprint arXiv:2205.10625, 2022.
Pei Zhou, Karthik Gopalakrishnan, Behnam Hedayatnia, Seokhwan Kim, Jay Pujara, Xiang Ren, Yang Liu, and Dilek Hakkani-Tur. Think before you speak: Using self-ralk to generate implicit commonsense knowledge for response generation. arXiv preprint 1000v:2110.08500) 2021.
13

---

Under review as a conference paper at ICLR 2023
pe 16 =
Token Ap 0.8)
Taken 6 (p 1)
M" 14 80)
Ground Truth Teacher Output Distribution (auto-regressive notation)
Tolee Cpp 0.54)
Taken ip = 0.21
Taken (p= 0.33]
[.C.D.CCCCC............
samples Pet
Tokan (p 0.01) (p 0)
lp
AG
Soft Labels,
Approximate soft labels, represented tokens for each time step.
Empirical Distribution Teacher's Output (n=6)
reprosented as real-valued logits all vocabe for each time
BC BF
Hard Labels, represented one token each stop
Figure 4: A concrete example used illustrate how we performed context distillation by approx- imately optimizing token level KL divergence. The red ones are hard labels, which are noisy. The green ones are soft labels that are less noisy, but consumes a lot of memory. Therefore, we create a list of samples from the soft labels distribution, and use the empirical distribution to approximate the soft label distribution.
8 APPENDIX
8.1 OPTIMIZING APPROXIMATE TOKEN-LEVEL KL DIVERGENCE
We illustrate the intuition in greater details using concrete example shown in Figure The ground truth teacher's s distribution over all output sequence is shown at the top. Ideally, we want the student's output token distribution pe to match the teacher's distribution p" (also referred to the "ground truth" in this section); however, optimizing this object exactly requires enumerating all possible output sequences, which computationally inefficient. Therefore, we sample n (=6 in the figure, 100 in the actual experiments) output sequences from the teacher output distribution. and our goal to use these output samples and the logit information to approximate p" and use it to efficiently train Po.
First, naïvely training the student to predict the empirical distribution (red) is noisy and sample inefficient, since it might not contain lower probability sequences, e.g., the sequence BD, which occurs only 2% of the time. wastes the logit information the teacher has provided us. Therefore, we should optimize
((A) List Token (f) + Lznd Token (0),
where
TiseToken (0) = KL(p" (t)))lp(ttz)).
and
2 KL(p" (t2l11 A) A)) + (t B)llpo B)),
Land Token (8)
(4)
and we got the coefficient 4/6 for the KL term for t because of 6 the samples in the empirical distribution starts with token A.
Since our hardware does not have enough memory to hold two large models the same time, we need to compute the teacher's 's logits on large set of inputs, store them, and then load the student model. load the teacher's s logits, and train the student accordingly. However, exactly computing ()((tet
B)) requires the full information of the probability distribution of (-1111 = B), which requires 200 KB in memory Therefore, the teacher's logit will incur a large storage, which makes it impossible fit them the specialized hardware (i.e., GPU or TPU);
LL
Using float32 representation and Incoder takenizer
14

---

Under review as a conference paper at ICLR 2023
Distillig Natural Language Explanations
Distillation Margin
5
-10
10
In-context Margin
Figure 5: Distilling explanations for 10 tasks from the Natural Instructions V2 test set. Each point corresponds to a task in the figure. "In-context Margin" quantifies how much the explanations help the teacher by measuring the difference in RougeL of TK-Instruct with and without explanations on each task, and "Distillation Margin" quantifies how much the student learns from the teacher by measuring the difference between the RougeL. score of the student after distillation and the Rougel. score of III-tuuttttt without explanations on each task. We observe positive correlation between the utility of the explanations to the teacher and how much our student learns.
consequently, we might need to store these information on Secondary memory, which incurs large communication cost. To compress the information of (p" = B) to fit on our special- ized hardware, need to use compact and unbiased representation; therefore, approximate (p"(-tt = B) with an empirical distribution of 100 sample tokens from this distribution (as shown in pa Figure 4 which now only takes 0.4KB
8.2 LEARNING FROM NATURAL LANGUAGE EXPLANATIONS
Hypothesis 7: context distillation can learn from natural language explanations when they benefit the teacher. In this experiment, we use the same teacher template as the experiment above; in contrast, in this experiment, we define a student template that exactly the same as the teacher, sttudet but without explanations (Figure student template B). We run context distillation to internalize the
effect of natural language explanations, using the task's 's training split as D.
Ideally, want the student to fully internalize the effect of natural language explanations and match the teacher's 's performance. To measure how much the student internalizes, we define two quantities for each task: the in-context margin, which measures the performance increase after add the explanations to the context, and 2) the distillation margin, which measures the performance increase after we perform context distillation. We plot the margins for each task in Figure [ and observe a positive and statistically significant correlation (r = 0.75, p = 1%). Therefore, context distillation can learn from natural language explanations when they benefit the teacher.
Notice that for our current model, not all tasks benefit from internalizing natural language explana- tions. However, we anticipate that future language models be more capable of responding natural language explanations (Kaplan et al. 2020 Scheurer et al. 2022). hence improving the performance of context distillation.
8.3 CONTROLLED TEXT GENERATION WITH CONTEXT DISTILLATION
By distilling prompts that describe desired behavior (e.g. don't 't output toxic language or only gen- erate positive sentiment text), we can instill form of control in language models. To test this, we prompt our In--tuttt model complete negative movie reviews from the IMDB sentiment classification dataset (Maas et al. 2011).
$10 int32
15

---

Under review as a conference paper at ICLR 2023
teacher
student w/o
student w
sent
RougeL.
ent
sent
RougeL.
ent
sent
RougeL
ent
positive
94.0
4.8
13.0
0.24
0.6
2.5
0.96
4.8
14.4
neutral
0.54
3.6
9.9
0.24
0.6
2.5
0.34
7.0
14.2
Table Controlled generation via-context distillation. The "sent" column reports the average sen- timent score of generations from the model, and the "ent" column refers to the model's 's estimated output entropy nats. We see that By injecting positive control directions into the model, we can obtain a model which outputs greater postive sentiment text without significantly sacrificing output coherence (RougeL) or output diversity (entropy).
The teacher template contains directions to complete the movie review specifically to end the review on positive note, and the student template contains directions just complete the review without any other specification. After distillation, the student should learn to internalize the abstract control preference that we expressed to the teacher model. We evaluate our models by querying GPT-3 for a sentiment classification on the model's 's output. We generate questions using a few-shot prompt to TK-instruct, and take 64 distillation steps with batch size 16 and AdamW optimizer with le-5 learning rate and O weight decay.
We see in Table that by applying this procedure we are able to successfully control our language
model.
8.4 FACTUAL KNOWLEDGE EDITING WITH CONTEXT DISTILLATION
Context distillation also enables a very natural way to edit the factual knowledge implicitly inter- nalized by language models, by distilling a prompt that states a new or edited declarative fact. This in contrast prior works (Mitchell et al. 2021 De Cao et al.) 2021, Meng et al.] [2022] on fact editing, which instead perform a constrained optimization procedure that directly learns an edit to the model parameters, corresponding to the factual knowledge update.
We use the challenging Counterfact dataset from Meng et al. (Meng et al. 2022 to test our method's fact editing ability. Each instance of the Counterfact task involves editing the object of a given fac- tual relation. Ideally the language model's should consistently apply the new fact under significant paraphrases to the original relation and context changes; the model should also not edit unrelated knowledge. To measure this, the Counterfact dataset provides set of paraphrase prompts (signifi- cant paraphrases of the original fact, which the LM should consistently edit), neighborhood prompts (un-related facts that share the same object the original pre-edit fact, which should not change), and attribute prompts (un-related facts which share the same object ad the new post-edit fact, which should not change).
We perform fact editing experiments on our TK-Instruct model. For a randomly selected fact cor- responding to each of the 34 unique relations in the Counterfact dataset, we synthesize teacher template, which includes a description of the fact to be edited and instructions to not edit unrelated facts. We also use GPT-3 to help us generate new attribute, paraphrase, and neighborhood prompt for each fact edit to use as demonstrations of desired behavior in the prompt, alongside a natural language explanation of why the fact edit or is not applied in each case. We generate the inputs P(r) using few-shot prompt to TK- Instrucct..
In Table [5] we evaluate our model on the set of paraphrase and neighborhood prompts in the dataset. We report both the average score - 1 [c object) > P(incorrect object)] and the average magnitude - P(correct object) Pincorreec object) under the language model, where P(correct object) and P(incorrect object) are the probability of the correct and incorrect objects under the model, when conditioned on the relevant subject and relation. We see that context distil- lation is largly able to recover the fact editing performance of the teacher, and preforms comparably in absolute score to current SOTA approaches to fact editing - ROME (Meng et al.) 2022) and MEND (Mitchell et al., 2021) even though method makes slightly different assumptions and therefore isn't 't directly comparable (i.e. our prompt contains more information than these methods use).
16

---

Under review as a conference paper at ICLR 2023
method
paraphrase
neighborhood
score
magnitude
score
magnitude
Teacher
73
29
58
Pre-distill student
34
80
Post-distill student
79
28
48
-2
GPT-3
75
17
MEND
65
12
38
-12
ROME
89
33
74
Table 5: Performance of context distillation on fact editing. We see that context distillation is able to recover the paraphrase score of the teacher, but slightly under-performs in neighborhood score.
However, our prompted TK-Instruct teacher model generally performs poorly on neighborhood prompts, which also carries context distilling larger and more capable language models. To demonstrate this, we evaluate GPT- 3 this task with the same prompt, which can see in Table performs much better on these neighborhood prompts. While we cannot perform context distillation on GPT-3 due to limitations in OpenAT's API, we expect these improvements to also carry over the distilled model.
distilled model. We expect this issue be largely resolved by
8.5 INSTRUCTION-TUNDD T5-1lB.
Following the procedure of (Wang et al. 2022a), we fine-tuned the 11B T5 LM-adapted model (Raf-) fel et al. 2019), Natural Instructions V2 (Wang al. (0022a). large dataset of 1600+ language tasks which includes, for each task, a task description, positive and negative in-context examples, and natural language explanations of why the output is right wrong for each of the in-context examples. Prior work (Wang et al. 2022a) has used this dataset to train instruction-tunel models on prompts consisting of only 2-positive examples planations, To maximize the flexibility of our instruction-turr tuned model, we instead instruction-uus -unded on distribution of randomized prompts, which consist of randomly chosen ples, 0 to 3 negative examples, and whether there is an explanation or not. We trained the model for 9728 steps with a batch size of 16 and AdamW optimizer on 32 TPU-V3 cores. The model achieves Rougel, score of 58 on the 2 positive, 2 negative with explanation test split, of unseen natural instructions tasks.
only 2-positive and negative examples with
positive exam-
8.6 FINETUNING LANGUAGE MODEL IMPLEMENTATION DETAILS
We run all experiments on 32 TPU-V3 cores, with the model parameters and optimizer states for fine-tuning sharded equally across all cores. Our codebase is built in Jax (Bradbury et al. 2018 Sabne 2020) using the PJIT function to handle the model parallel training and inference.
8.7 GENERAL EXPERIMENT DETAILS
For all experiments, except where denoted otherwise, we distill on 4096 examples for epoch with batch size 16 with AdamW optimizer. We use learning rate of le-5 with TK-Instruct and le-4 with Incoder. We use 0 weight decay for all experiments.
8.8 DISTILLING CONCRETE EXAMPLES DETAILS
Gradient descent details (hypothesis 4). For gradient descent we fit all training examples into a single batch and train for 25 epochs using the AdamW optimizer with a learning rate le-5. We report the performance of the epoch with the highest average exact set match score across all databases.
Distilling long contexts details (hypothesis 5), To estimate the per- pok-token log- proaabbbabslities of the y sampled from the teacher ensemble for distillation, average the probabilities of each y under teacher prompts. We estimate the teacher performance by performing greedy decoding on the ensemble of 8 prompts uniformly sampled from the set of all 4 choose 8 teacher prompts. For each
17

---

Under review as a conference paper at ICLR 2023
database, we distill two students with different in-context examples in the prompt, and report the average exact-set match accuracy for both of these students.
8.9 DISTILLING STEP-BY-STER REASONING DETAILS
Distilling scratchpads with T5-small (hypothesis 6). All baselines were trained for 1000 epochs except "Scratchpad then Direct" which was trained for 1000 epochs to predict scratchpads and then 1000 epochs to directly predict the answer with batch size of learning rate of 3e-4, and AdamW optimizer. We report performance at the end of training on 10k unseen addition problems.
Transfering step-by-step reasoning details (hypothesis 7). Since IK-Instruct cannot success- fully do scratchpad addition from a few shot prompt, to initialize the teacher, we first fine-tune TK- Instruct on the same distribution of 500 scratchpad examples from the previous experiment mixed in with 4096 randomly selected examples from the "2 positive" split of Natural Instructions-VV V2 train- ing set, such that the teacher doesn't t lose its ability to respond to prompts. We train the teacher for 32 epochs, at which point it achieves 97% accuracy on 2000 held-out scratchpad addition problems.
18


Reviews:
[{"review_id": "z9cdAgTM-o", "reviewer": "Reviewer_ipLo", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This work addresses some of the shortcomings of few-shot learning paradigms with large language models such as learning from instructions and explanations, scratch-pad/chain-of-though reasoning and learning from demonstrations. The authors point out that the following issues are prevalent with these methods - (i) The context window can be long, which results in inference being computationally expensive (ii) The inability to fit more examples than the context window allows and (iii) The gains disappear after the demonstrations or scratchpad disappear, i.e., the model doesn’t internalize these gains. This work proposes ‘context distillation’ as a solution. A student model is trained to match the distribution of outputs of a teacher model, where the student model uses a simpler and shorter context, such as only a simple instruction without examples/explanations/scratch-pad. Experiments attempt to show the advantage of the proposed method under the different few-shot learning paradigms described above.", "summary_of_the_review": "The authors propose a neat idea which is applicable across a variety of learning paradigms with large language models. However, the experimental results are not comprehensive enough to understand the benefit of the proposed approach. ", "strength_and_weaknesses": "Pros\n* The paper considers an interesting problem - Internalizing the performance gains from methods known to improve few-shot learning performance, but without paying the additional price of having to process a long context.\n* An interesting solution is proposed.\n* Experiments show some evidence of the advantage of the proposed approach.\n\nCons\n* Technical details are missing/not clearly described.\n* Experimental results are sparse and not fully convincing.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "i4uO8Suspa-", "reviewer": "Reviewer_Krod", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper introduces a new distillation method, with a teacher model given a long, detailed instruction as well as in-context examples, and a student model given a concise instruction. The motivation is that, while the model performs the task well by reading a detailed instruction and examples, it is often expensive at inference time, because it has a significantly longer context. Therefore, having a student model that operates with shorter instruction but works as well as the model given a long instruction and in-context examples is beneficial. The paper evaluates the proposed model on a range of scenarios, effectively showing that the student model effectively learns from the teacher model with long instruction or explanation, and outperforms baselines such as naive instruction tuning without distillation.", "summary_of_the_review": "Overall, the paper is well-written, and the idea of dilstilling from the model reading detailed instruction to the model reading more concise instruction is interesting and new. However, I am concerned with overclaiming in the paper, missing critical baselines and inconsistency in evaluation of the model.\n", "strength_and_weaknesses": "**Strengths**\n* The paper tackles an interesting problem – building a model that learns from a concise instruction and performs the task as much as the model given a detailed instruction and examples does. The problem is well-motivated – reading a long instruction and examples is often very expensive and inefficient, and is even strictly bounded by the maximum sequence length of the Transformer.\n* The idea of distilling from the model reading a long instruction to the model reading a shorter instruction is new.\n* The paper specifies a series of hypotheses and shows empirical results that verify each hypothesis very clearly.\n\n**Weaknesses**\n* Overclaiming:\n    * The paper made a strong emphasis that the proposed model is able to “internalize” the task, and made this claim many times overall in the paper. However, to my point of view, “internalizing the task” is exactly “fine-tuning”, and in fact, the proposed method is all about “better fine-tuning through distillation”. In this regard, I think the paper is overall overclaiming.\n    * Some details that indicate the proposed model is based on more assumptions than the baselines are not explicitly mentioned. For instance, based on my understanding, the proposed model is trained on the unlabeled portion of the test datasets, whereas baselines are not. Is it correct? (I am not super sure about this and asked this question in the “clarity” section.) If it is true, it is a fairly strong assumption to assum the unlabeled data of the test datasets, and should be mentioned very clearly. Not making this clear seems overclaiming.\n* Missing the most critical baseline\n    * I think the most important baseline is a student model that is instruction-tuned on a particular template (either teacher’s template or student’s template) with no distillation. To my understanding, this baseline is not provided, except for one experiment (Table 2). Instead, the paper claims that the proposed model effectively learns from the teacher model because it is better than the pre-distill student model. However, to my understanding, the pre-distill model is not instruct-tuned, thus it is unclear if this gain is coming from instruction tuning, or distillation from a teacher model. For this reason, I strongly believe that the paper should have mainly compared with the model instruction-tuned but without distillation, and should not claim the effectiveness of the distillation method without this baseline. \n    * Side note: I think “Direct Gradient Descent” is basically the model that is instruction-tuned but without distillation. Is this correct? I think the name is not very intuitive. To my understanding, it is using the template for the teacher model (long instruction, in-context examples) – and if it is correct, I think it’s also good to include the model using the template for the student model (concise instruction). This will support the assumption in the paper that longer instruction is more helpful than concise instruction, which this paper actually never explicitly showed (nor cited relevant prior work).\n* Experiments are not very consistent with no specific reason specified. For instance:\n    * In the experiments with Natural Instruction v2, the result is evaluated only on 5 datasets with the largest gains from instruction tuning and other 5 datasets with the smallest gains from instruction tuning. It is not clear why the paper chooses to do so instead of reporting all datasets, especially because the paper did not report these two sets of 5 datasets separately and discuss how their results differ.\n\nMinor:\n* While the paper claims that prior work in distillation is mostly focusing on reducing the model size, I don’t think it is true – previous work has shown that distillation helps even when the teacher model and the student model have the same size of the model.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "g_YagyiEW", "reviewer": "Reviewer_fidG", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes “context distillation”, a method to “internalize” the information provided by abstract instructions, concrete demonstration examples, or scratchpad (model intermediate output that helps reasoning). They first use a language model to sample some input (e.g., movie reviews), and then they use the teacher model to generate pseudo answers given inputs as well as demos, instructions and generated scratch-pad. The pseudo input-answer pairs are used to fine-tune the student model (minimizing the KL divergence using teacher logits), without the extra information. The benefit of internalizing the knowledge is that there is no longer a need to use the extra input and output, making the use of LLM more efficient. The difference to traditional distillation is that the teacher and student models are the same one, and it is the “extra information” that gets distilled into the student model. The authors conducted three experiments:\n\nExp#1: learning from instructions. They first fine-tune the model on Natural-Instruction-V2 (a T5-11B). They show that they can significantly improve the student model’s performance without instruction or explanation.\n\nExp#2: learning from examples. The authors use SPIDER text-to-SQL as a test bed. They show that context distillation can effectively internalize the examples in the context and outperform direct fine-tuning. \n\nExp#3: learning from step-by-step reasoning. The authors test the model on 8-digit addition questions. The experiment shows that the student model can even internalize the step-by-step reasoning by fine-tuning on teacher model’s output, without using scratchpad.\n\nEven though the concept of “context distillation” was already proposed by Choi et al.; Askell et al., this paper did a good job of systematically testing it on three distinctive aspects, and the experiment results show that context distillation is an effective way to “internalize” knowledge from the “context”.\n", "summary_of_the_review": "Overall the paper is well written and the conducted research is insightful and the results are strong. The paper should be accepted. The quality of the presentation can be further improved per the weakness section. ", "strength_and_weaknesses": "# Strength\n\nThe authors conduct a systematic study of “context distillation” in three aspects: instructions, demonstrations, and step-by-step reasoning. The method is intuitive and the results are strong. It is not only a way to improve the efficiency of using in-context learning and instruction prompting, but the fact that the knowledge can be “internalized” is very interesting. \n\n# Weakness\n\nThe experiment setup can be significantly improved. First, the paper can be improved on how they demonstrate the results. There are very few tables/figures on the results in the main paper and most numbers scatter in the text. Second, the clarity of the experiment setup (what datasets they use for exp#1, how many pseudo input examples they generate, etc. sorry if I missed it) can be improved. For Exp#2, it is very limited to only test context distillation on the SQL dataset. It would be better to show other datasets too, like NLU, QA, etc. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "fk-zrIXIvFy", "reviewer": "Reviewer_7ryR", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors present an alternative paradigm to prompting: context distillation. The idea is that the model, instead of utilizing examples and specific prompts just for predicting without learning anything from these, tries to internalize this information. The authors propose a teacher-student architecture, where both start from the same parameters' state, and the updates of the student network are driven by the differences in x and y, fed and output respectively, by each network. While the teacher network receives a complete input, with examples and explanations, and outputs the target plus scratch-pad text, the student network generally receives some simple task explanation and the raw input, and is trained to output the raw target.\n\nThroughout the paper, the authors come up with several hypotheses, aiming at proving these contextualization capabilities, and perform different tests to assess the correctness of these hypotheses. In all cases the student network, after distillation, outperforms the original student network, proving the potential of this approach for incorporating valid contextual information into the networks. The authors test the presented methodology in three different contexts, which requires different datasets and downstream tasks.", "summary_of_the_review": "Overall, I believe the paper presents really interesting ideas, as well as tests and results to support the main claims. However, I still feel that the concept of knowledge vanishing has not been tackled, and it would be important to understand how it works in order to really provide a methodology that can perform in multi-task environments. I believe the authors should further discuss this, in order to improve the overall quality.", "strength_and_weaknesses": "The authors propose a really interesting idea, where distillation is used to infuse a pre-trained network with further context that is associated with diverse downstream tasks, boosting their performance. The approach considered, where different hypotheses are laid out and assessed with different tests, is interesting, and helps grounding step by step the possibilities of the methodology. Most of the results are clearly presented, and successfully justify the hypothesis postulated, even though further tests on other downstream tasks would be necessary to really confirm them. \n\nStill, I believe the authors do not tackle a really important point, critical to assess the full potential of the methodology: how the knowledge distilled in the network vanishes, or not, as it is trained on more and more tasks. Therefore, it would have been interesting seeing results of this phenomena, as another hypothesis laid out on this premise. I believe such hypothesis would have been more relevant to evaluate than some of the ones considered by the authors. The concern is that the knowledge distilled vanishes as the network is trained on more tasks, which could be indicated by the decrease in performance of those tasks. If this is the case, the demanding fine-tune required would be less justifiable, and more ad-hoc networks or other fine-tuning methods might be a better alternative.\n\nBesides, it would be valuable to have confidence intervals in the reported accuracies, in order to understand how the distillation process can be affected by the randomness in the set of provided inputs, as most experiments are just trained on a few epochs. \n\nFinally, I have two other questions:\n- The explanation given in Section 2.4, for the simplification done on the soft-labels, is not really clear to me. Could you further elaborate? \n- In table 1, my intuition tells me that, when training on a single task, i.e. naive post-distill, the performance should be better than when mixing several. Is not that right? Or is it the case that the single task for naive training is the most complicated one, among all the ones used in the Mixed training? However, according to table 6, the performance of all tasks increases in the mixed version. Could you provide some further intuition on this?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject