Elamur Research⌕
☰
home / papers / 2023 / In-context Autoencoder for Context Compression in a Large Language Model

In-context Autoencoder for Context Compression in a Large Language Model

Paper · 2023 · Tao Ge, Jing Hu, Lei Wang et al.

Summary

Transformer LLMs struggle with long inputs because self-attention cost grows with sequence length. Architectural fixes reduce that cost but often lose quality on long texts. This paper takes a different route: it shrinks the input rather than changing the model. The In-context Autoencoder (ICAE) uses a LoRA-adapted copy of Llama as an encoder. The encoder reads a context with k learnable memory tokens appended and outputs k "memory slots." The unmodified Llama then acts as the decoder and conditions on these slots in place of the original text. Training has two stages. First comes self-supervised pretraining on the Pile, combining a reconstruction (autoencoding) objective with a text-continuation objective. Second comes instruction tuning on PwC, a new collection of (context, prompt, response) triples. Representing 512-token inputs with 128 slots gives 4× compression while adding roughly 1% extra parameters, which lowers inference latency and GPU memory use. The approach can complement other long-context methods. The authors also link it to human working memory, noting that heavy pretraining improves how faithfully the slots store information.

Key points

  • Architecture: The encoder applies LoRA (rank 128) to the query and value projections of Llama's attention, plus an embedding table for memory tokens. The frozen target LLM serves as the decoder.
  • Default setup: k = 128 memory slots and a maximum of 512 tokens per input, giving 4× compression.
  • Data: Pretraining uses the Pile. The PwC (Prompt-with-Context) dataset, released with the paper, has 240k training and 18k test samples.
  • Reconstruction quality: With Llama-7b, autoencoding loss stays below 0.05. Reconstruction is almost perfect (near 100% BLEU and Exact-Match) for contexts up to about 300 tokens, and quality drops once inputs exceed about 400 tokens.
  • Effect of pretraining: In GPT-4–judged comparisons, a non-pretrained ICAE made factual errors that the pretrained ICAE avoided, such as "three" instead of "30 years" and a wrong job title. Code, data, and models are available at github.com/getao/icae.

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

🏢MicrosoftRedmond · United States378 works · 161,630 citations

 Fields, topics and concepts

📚 Shelves 5The 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 worksComputer Vision and Pattern Recognition · 3,486 works

🧭 Topics 2Research subjects and languages the work is about.Language models & pretraining · 3,894 worksEfficient models & compression · 297 works

Cite

Cite as Tao Ge, Jing Hu, Lei Wang, Xun Wang, Si-Qing Chen and Furu Wei. 2023. In-context Autoencoder for Context Compression in a Large Language Model. arXiv:2307.06945.

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

@misc{ge2023context,
  title = {{In-context Autoencoder for Context Compression in a Large Language Model}},
  author = {Tao Ge, Jing Hu, Lei Wang, Xun Wang, Si-Qing Chen and Furu Wei},
  year = {2023},
  doi = {10.48550/arxiv.2307.06945},
  eprint = {2307.06945},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2307.06945}
}

References

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

  1. Amanda Askell, Yuntao Bai, Anna Chen, Dawn Drain, Deep Ganguli, Tom Henighan, Andy Jones, Nicholas Joseph, Ben Mann, Nova DasSarma, et al. A general language assistant as a laboratory for alignment. arXiv preprint arXiv:2112.00861, 2021. find
  2. Alan Baddeley. Working memory. Science, 255(5044):556–559, 1992. find
  3. Mohammad Bavarian, Heewoo Jun, Nikolas Tezak, John Schulman, Christine McLeavey, Jerry Tworek, and Mark Chen. Efficient training of language models to fill in the middle. arXiv preprint arXiv:2207.14255, 2022.
  4. Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05150, 2020.
  5. Amanda Bertsch, Uri Alon, Graham Neubig, and Matthew Gormley. Unlimiformer: Long-range transformers with unlimited length input. Advances in Neural Information Processing Systems, 36, 2023.
  6. Aydar Bulatov, Yury Kuratov, and Mikhail Burtsev. Recurrent memory transformer. Advances in Neural Information Processing Systems, 35:11079–11091, 2022. find
  7. Aydar Bulatov, Yuri Kuratov, and Mikhail S Burtsev. Scaling transformer to 1m tokens and beyond with rmt. arXiv preprint arXiv:2304.11062, 2023. — [arXiv:2304.11062](https://arxiv.org/abs/2304.11062) find
  8. Howard Chen, Ramakanth Pasunuru, Jason Weston, and Asli Celikyilmaz. Walking down the memory maze: Beyond context limit through interactive reading. arXiv preprint arXiv:2310.05029, 2023. — [arXiv:2310.05029](https://arxiv.org/abs/2310.05029) find
  9. Alexis Chevalier, Alexander Wettig, Anirudh Ajith, and Danqi Chen. Adapting language models to compress contexts. arXiv preprint arXiv:2305.14788, 2023.
  10. Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509, 2019.
  11. Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamás Sarlós, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, David Belanger, Lucy J. Colwell, and Adrian Weller. Rethinking attention with performers. ArXiv, abs/2009.14794, 2020.
  12. Grégoire Delétang, Anian Ruoss, Paul-Ambroise Duquenne, Elliot Catt, Tim Genewein, Christopher Mattern, Jordi Grau-Moya, Li Kevin Wenliang, Matthew Aitchison, Laurent Orseau, et al. Language modeling is compression. arXiv preprint arXiv:2309.10668, 2023.
  13. Jiayu Ding, Shuming Ma, Li Dong, Xingxing Zhang, Shaohan Huang, Wenhui Wang, and Furu Wei. Longnet: Scaling transformers to 1,000,000,000 tokens. arXiv preprint arXiv:2307.02486, 2023.
  14. Randall W Engle, Stephen W Tuholski, James E Laughlin, and Andrew RA Conway. Working memory, short-term memory, and general fluid intelligence: a latent-variable approach. Journal of experimental psychology: General, 128(3):309, 1999. find
  15. K Anders Ericsson, William G Chase, and Steve Faloon. Acquisition of a memory skill. Science, 208 (4448):1181–1182, 1980. find
  16. Rinon Gal, Yuval Alaluf, Yuval Atzmon, Or Patashnik, Amit H Bermano, Gal Chechik, and Daniel Cohen-Or. An image is worth one word: Personalizing text-to-image generation using textual inversion. arXiv preprint arXiv:2208.01618, 2022. find
  17. Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. The Pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00027, 2020. find
  18. Tao Ge, Hu Jing, Li Dong, Shaoguang Mao, Yan Xia, Xun Wang, Si-Qing Chen, and Furu Wei. Extensible prompts for language models on zero-shot language style customization. In A. Oh, T. Neumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (eds.), Advances in Neural Information Processing Systems,… find
  19. Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. arXiv preprint arXiv:2106.09685, 2021.
  20. Huiqiang Jiang, Qianhui Wu, Chin-Yew Lin, Yuqing Yang, and Lili Qiu. Llmlingua: Compressing prompts for accelerated inference of large language models. arXiv preprint arXiv:2310.05736, 2023a.
  21. Zhiying Jiang, Matthew Yang, Mikhail Tsirlin, Raphael Tang, Yiqin Dai, and Jimmy Lin. “low-resource” text classification: A parameter-free classification method with compressors. In Findings of the Association for Computational Linguistics: ACL 2023, pp. 6810–6828, 2023b. find
  22. Mark A. Kramer. Nonlinear principal component analysis using autoassociative neural networks. Aiche Journal, 37:233–243, 1991. find
  23. Patrick Lewis, Ethan Perez, Aleksandra Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Küttler, Mike Lewis, Wen-tau Yih, Tim Rocktäschel, et al. Retrieval-augmented generation for knowledge-intensive nlp tasks. Advances in Neural Information Processing Systems, 33: 9459–9474, 2020.
  24. Nelson F. Liu, Kevin Lin, John Hewitt, Ashwin Paranjape, Michele Bevilacqua, Fabio Petroni, and Percy Liang. Lost in the middle: How language models use long contexts, 2023.
  25. Eleanor A Maguire, Elizabeth R Valentine, John M Wilding, and Narinder Kapur. Routes to remembering: the brains behind superior memory. Nature neuroscience, 6(1):90–95, 2003. find
  26. Jesse Mu, Xiang Lisa Li, and Noah Goodman. Learning to compress prompts with gist tokens. arXiv preprint arXiv:2304.08467, 2023.
  27. OpenAI. Gpt-4 technical report. ArXiv, abs/2303.08774, 2023.
  28. Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback. Advances in Neural Information Processing Systems, 35: 27730–27744, 2022.
  29. Kishore Papineni, Salim Roukos, Todd Ward, and Wei Jing Zhu. Bleu: a method for automatic evaluation of machine translation. 10 2002. doi: 10.3115/1073083.1073135.
  30. Guangyue Peng, Tao Ge, Si-Qing Chen, Furu Wei, and Houfeng Wang. Semiparametric language models are scalable continual learners. arXiv preprint arXiv:2303.01421, 2023. — [arXiv:2303.01421](https://arxiv.org/abs/2303.01421) find
  31. Guanghui Qin and Benjamin Van Durme. Nugget: Neural agglomerative embeddings of text. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of M… find
  32. Jack W Rae, Anna Potapenko, Siddhant M Jayakumar, and Timothy P Lillicrap. Compressive transformers for long-range sequence modelling. arXiv preprint arXiv:1911.05507, 2019.
  33. Charlie Snell, Dan Klein, and Ruiqi Zhong. Learning by distilling context. arXiv preprint arXiv:2209.15189, 2022.
  34. Woomin Song, Seunghyuk Oh, Sangwoo Mo, Jaehyung Kim, Sukmin Yun, Jung-Woo Ha, and Jinwoo Shin. Hierarchical context merging: Better long context understanding for pre-trained llms. In The Twelfth International Conference on Learning Representations, 2024. find
  35. Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aur’elien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. Llama: Open and efficient foundation language models. ArXiv, abs/23…
  36. Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288, 2023b.
  37. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
  38. Yizhong Wang, Yeganeh Kordi, Swaroop Mishra, Alisa Liu, Noah A Smith, Daniel Khashabi, and Hannaneh Hajishirzi. Self-instruct: Aligning language model with self generated instructions. arXiv preprint arXiv:2212.10560, 2022. find
  39. Zhenhailong Wang, Shaoguang Mao, Wenshan Wu, Tao Ge, Furu Wei, and Heng Ji. Unleashing the emergent cognitive synergy in large language models: A task-solving agent through multi-persona self-collaboration. arXiv preprint arXiv:2307.05300, 2023.
  40. Jason Wei, Maarten Bosma, Vincent Y Zhao, Kelvin Guu, Adams Wei Yu, Brian Lester, Nan Du, Andrew M Dai, and Quoc V Le. Finetuned language models are zero-shot learners. arXiv preprint arXiv:2109.01652, 2021.
  41. Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V Le, Denny Zhou, et al. Chain-of-thought prompting elicits reasoning in large language models. Advances in neural information processing systems, 35:24824–24837, 2022.
  42. David Wingate, Mohammad Shoeybi, and Taylor Sorensen. Prompt compression and contrastive conditioning for controllability and toxicity reduction in language models. arXiv preprint arXiv:2210.03162, 2022. find
  43. Yuhuai Wu, Markus N Rabe, DeLesley Hutchins, and Christian Szegedy. Memorizing transformers. arXiv preprint arXiv:2203.08913, 2022. find
  44. Yadong Zhang, Shaoguang Mao, Tao Ge, Xun Wang, Yan Xia, Man Lan, and Furu Wei. K-level reasoning with large language models. arXiv preprint arXiv:2402.01521, 2024. — [arXiv:2402.01521](https://arxiv.org/abs/2402.01521) find
  45. Hao Zhao, Maksym Andriushchenko, Francesco Croce, and Nicolas Flammarion. Long is more for alignment: A simple but tough-to-beat baseline for instruction fine-tuning. arXiv preprint arXiv:2402.04833, 2024. — [arXiv:2402.04833](https://arxiv.org/abs/2402.04833) find
  46. Lin Zheng, Chong Wang, and Lingpeng Kong. Linear complexity randomized self-attention mechanism. In International Conference on Machine Learning, 2022. find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/ge-2023-context-autoencoder-context-compression-large/