Elamur Research⌕
☰
home / papers / 2022 / Learning by Distilling Context

Learning by Distilling Context

Paper · 2022 · Charlie Snell, Dan Klein and Ruiqi Zhong

Summary

Language models do better when given rich context, such as detailed instructions, worked examples or a step-by-step scratch-pad. These gains vanish once that context is removed, so every inference call pays for the long prompt, and some knowledge cannot fit into the context window at all. The authors propose context distillation to make the model absorb such gains into its weights.

The method works on synthetic or unlabeled task inputs. The model first acts as a teacher: it receives an elaborate prompt and produces reasoning plus an answer. An extractor keeps only the final answer. The same model, as a student, is then fine-tuned to produce that answer from a minimal prompt. Teacher and student start with identical parameters, so learning comes only from the difference in prompts. Multiple distillations can be run sequentially or simultaneously.

Using TK-Instruct (11B) and Incoder, the method internalizes instructions, explanations, SPIDER training examples and scratch-pad addition. It beats plain gradient descent on Text-to-SQL. This suggests a general way for models to improve themselves from context.

Key points

  • Four components define the setup: a raw input distribution, a teacher template, a student template and an answer extractor. A simultaneous variant sums several distillation losses, and a sequential variant chains them to update or overwrite earlier learning.
  • Instructions and explanations: Instructions and explanations from Natural-Instructions-V2 (1600+ tasks) were distilled into a TK-Instruct model built on T5 LM-adapted 11B. The more an explanation helped the teacher, the more the student gained. Associations between task IDs and instructions could be reassigned and overwritten.
  • SPIDER Text-to-SQL: With Incoder, 8-shot context distillation outperformed direct gradient-descent fine-tuning by 9%. Distilling more examples than the context window holds kept improving accuracy.
  • 8-digit addition: The model internalized scratch-pad reasoning for 8-digit addition, and this skill transferred to downstream QA. The teacher reached 97% accuracy on 2000 held-out problems.
  • Typical training setup: 4096 distillation examples, 1 epoch, batch size 16, AdamW, on 32 TPU-V3 cores with a Jax codebase.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Institutions

🎓University of California, BerkeleyBerkeley · United States209 works · 244,782 citations

 Fields, topics and concepts

📚 Shelves 4The field of science this work belongs to (OpenAlex fields).Computer Science · 18,614 worksArtificial Intelligence · 17,965 worksArts and Humanities · 17,450 worksLanguage and Linguistics · 17,296 works

Cite

Cite as Charlie Snell, Dan Klein and Ruiqi Zhong. 2022. Learning by Distilling Context. arXiv:2209.15189.

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.48550/arxiv.2209.15189 · arXiv 2209.15189

@misc{snell2022learning,
  title = {{Learning by Distilling Context}},
  author = {Charlie Snell, Dan Klein and Ruiqi Zhong},
  year = {2022},
  doi = {10.48550/arxiv.2209.15189},
  eprint = {2209.15189},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2209.15189}
}

References

42 references, 18 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Romero Adriana, Ballas Nicolas, K Samira Ebrahimi, Chassang Antoine, Gatta Carlo, and B Yoshua. Fitnets: Hints for thin deep nets. Proc. ICLR, 2, 2015. find
  2. Anonymous. Large language models can self-improve. In Submitted to The Eleventh International Conference on Learning Representations, 2023. URL https://openreview.net/forum? id=NiEtU7blzN. under review. ↗
  3. 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 Hernandez, Jackson Kernion, Kamal Ndousse, Catherine Olsson, Dario Amodei, Tom Brown, Jack Clark, Sam McCandlish, Chris Olah, a… find
  4. Alan Baddeley. Working memory. Science, 255(5044):556–559, 1992. find
  5. James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http://github.com/google/jax. ↗
  6. Yanda Chen, Ruiqi Zhong, Sheng Zha, George Karypis, and He He. Meta-learning via language model in-context tuning. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 719–730, Dublin, Ireland, May 2022. Association for Computational… find
  7. Eunbi Choi, Yongrae Jo, Joel Jang, and Minjoon Seo. Prompt injection: Parameterization of fixed inputs, 2022. URL https://arxiv.org/abs/2206.11349. — [arXiv:2206.11349](https://arxiv.org/abs/2206.11349) ↗
  8. Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14168, 2021.
  9. Nicola De Cao, Wilker Aziz, and Ivan Titov. Editing factual knowledge in language models, 2021. URL https://arxiv.org/abs/2104.08164. ↗
  10. 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. find
  11. Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in a neural network, 2015a. URL https://arxiv.org/abs/1503.02531. ↗
  12. Geoffrey Hinton, Oriol Vinyals, Jeff Dean, et al. Distilling the knowledge in a neural network. arXiv preprint arXiv:1503.02531, 2(7), 2015b.
  13. 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 arXiv:2001.08361, 2020.
  14. Aitor Lewkowycz, Anders Andreassen, David Dohan, Ethan Dyer, Henryk Michalewski, Vinay Ramasesh, Ambrose Slone, Cem Anil, Imanol Schlag, Theo Gutman-Solo, et al. Solving quantitative reasoning problems with language models. arXiv preprint arXiv:2206.14858, 2022.
  15. Xiaodong Liu, Pengcheng He, Weizhu Chen, and Jianfeng Gao. Improving multi-task deep neural networks via knowledge distillation for natural language understanding. arXiv preprint arXiv:1904.09482, 2019. find
  16. Andrew L. Maas, Raymond E. Daly, Peter T. 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, J…
  17. Kevin Meng, David Bau, Alex Andonian, and Yonatan Belinkov. Locating and editing factual associations in gpt, 2022. URL https://arxiv.org/abs/2202.05262. ↗
  18. Sewon Min, Mike Lewis, Luke Zettlemoyer, and Hannaneh Hajishirzi. MetaICL: Learning to learn in context. In Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 2791–2809, Seattle, United States, July 2022… find
  19. 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 1: Long Papers), pp. 3470–3487, Dublin, Ireland, May 2… find
  20. Eric Mitchell, Charles Lin, Antoine Bosselut, Chelsea Finn, and Christopher D. Manning. Fast model editing at scale, 2021. URL https://arxiv.org/abs/2110.11309. ↗
  21. Maxwell Nye, Anders Johan Andreassen, Guy Gur-Ari, Henryk Michalewski, Jacob Austin, David Bieber, David Dohan, Aitor Lewkowycz, Maarten Bosma, David Luan, Charles Sutton, and Augustus Odena. Show your work: Scratchpads for intermediate computation with language models, 2021a. URL https://arxiv.org/… ↗
  22. Maxwell Nye, Anders Johan Andreassen, Guy Gur-Ari, Henryk Michalewski, Jacob Austin, David Bieber, David Dohan, Aitor Lewkowycz, Maarten Bosma, David Luan, et al. Show your work: Scratchpads for intermediate computation with language models. arXiv preprint arXiv:2112.00114, 2021b.
  23. 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 https://arxiv.org/abs/1910.10683. ↗
  24. 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. J. Mach. Learn. Res., 21(140):1–67, 2020.
  25. Nitarshan Rajkumar, Raymond Li, and Dzmitry Bahdanau. Evaluating the text-to-sql capabilities of large language models. arXiv preprint arXiv:2204.00498, 2022. find
  26. Amit Sabne. Xla : Compiling machine learning for peak performance, 2020. find
  27. 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, Debajyoti Datta, Jonathan Chan… — Multitask Prompted Training Enables Zero-Shot Task Generalization
  28. 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 on Learning with Natural Language Supervision, 2022. find
  29. Yizhong Wang, Swaroop Mishra, Pegah Alipoormolabashi, Yeganeh Kordi, Amirreza Mirzaei, Anjana Arunkumar, Arjun 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, Ma… find
  30. Yizhong Wang, Swaroop Mishra, Pegah Alipoormolabashi, Yeganeh Kordi, Amirreza Mirzaei, Anjana 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. find
  31. Zirui Wang, Adams Wei Yu, Orhan Firat, and Yuan Cao. Towards zero-label language learning. arXiv preprint arXiv:2109.09193, 2021. — [arXiv:2109.09193](https://arxiv.org/abs/2109.09193) find
  32. Peter C Wason and J St BT Evans. Dual processes in reasoning? Cognition, 3(2):141–154, 1974. find
  33. Jason Wei, Maarten Bosma, Vincent Zhao, Kelvin Guu, Adams Wei Yu, Brian Lester, Nan Du, Andrew M. Dai, and Quoc V Le. Finetuned language models are zero-shot learners. In International Conference on Learning Representations, 2022a. URL https://openreview.net/ forum?id=gEZrGCozdqR. ↗
  34. 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 arXiv:2201.11903, 2022b.
  35. Qizhe Xie, Minh-Thang Luong, Eduard Hovy, and Quoc V. Le. Self-training with noisy student improves imagenet classification. In 2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 10684–10695, 2020. doi: 10.1109/CVPR42600.2020.01070. find
  36. 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. find
  37. 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 Confer…
  38. Eric Zelikman, Yuhuai Wu, Jesse Mu, and Noah D. Goodman. Star: Bootstrapping reasoning with reasoning, 2022. URL https://arxiv.org/abs/2203.14465. ↗
  39. 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 Association for Computational Linguistics: EMNLP 2021, pp. 2856–2878, Punta Cana, Dominican Republic, November 2021. Association f…
  40. Denny Zhou, Nathanael Schärli, Le Hou, Jason Wei, Nathan Scales, Xuezhi Wang, Dale Schuurmans, Olivier Bousquet, Quoc Le, and Ed Chi. Least-to-most prompting enables complex reasoning in large language models. arXiv preprint arXiv:2205.10625, 2022.
  41. Pei Zhou, Karthik Gopalakrishnan, Behnam Hedayatnia, Seokhwan Kim, Jay Pujara, Xiang Ren, Yang Liu, and Dilek Hakkani-Tur. Think before you speak: Using self-talk to generate implicit commonsense knowledge for response generation. arXiv preprint arXiv:2110.08501, 2021. positive neutral sent 94.0 0.5… — [arXiv:2110.08501](https://arxiv.org/abs/2110.08501) find
  42. Pre-distillation Student sent RougeL ent 0.24 0.6 2.5 0.24 0.6 2.5 Post-Distillation Student sent RougeL ent 0.96 4.8 14.4 0.34 7.0 14.2 find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/snell-2022-learning-distilling-context/