Elamur Research⌕
☰
home / papers / 2024 / Long Context Compression with Activation Beacon

Long Context Compression with Activation Beacon

Paper · 2024 · Peitian Zhang, Zheng Liu, Shitao Xiao et al.

Summary

Long inputs are costly for transformer LLMs. Self-attention cost grows quadratically, and holding the KV cache for the whole sequence takes a lot of GPU memory. Earlier context compression methods have several weaknesses:

  • they squeeze the input into a few soft tokens in a single pass,
  • they must re-encode those tokens before generation,
  • they fix the ratio in advance,
  • or they rely on the user's question to decide what to discard.

The authors propose Activation Beacon, a plug-in module that stores condensed context in the per-layer keys and values of special beacon tokens. The input is cut into fixed-size chunks. Inside each chunk, one beacon token follows every unit of α tokens, so α sets the compression ratio. Chunks are encoded one after another. Raw-token activations are dropped, while beacon activations are kept and reused for later chunks.

Training uses compression-based auto-regression on a 1B plain-text corpus plus 30K instruction samples (up to 20K tokens), with the ratio randomly drawn per chunk. On Llama-2 and Qwen-2, the method nearly matches uncompressed models on tasks up to 128K tokens. It runs about twice as fast and shrinks the KV cache eightfold, making long-input processing much cheaper.

Key points

  • The recommended setting is x8 compression. It gives roughly 2x end-to-end speedup over the uncompressed Full-FT baseline at 128K and an 8x smaller KV cache, with generation quality close to lossless.
  • The efficiency test was a multi-turn Multi-Needle-in-a-Haystack task (NeedleBench; 32K for Llama-2, 128K for Qwen-2, 3 needles). Activation Beacon was 1.8x faster than AutoCompressor, and 9.3x and 3.6x faster than LongLLMLingua and SnapKV across three turns. The gap with the last two comes from Activation Beacon not depending on the query.
  • Short-context skills are largely kept, likely because the original weights stay frozen. MMLU went from 47.5 to 46.6 (Llama-2-7B) and from 70.1 to 69.1 (Qwen-2-7B); ARC-C, BoolQ and GSM8K are essentially unchanged.
  • The ablation used Qwen-2-7B on LongBench Single-Doc QA (32K, x4), where the default scores 40.5. Scores dropped to:
    • 35.2 without fine-grained units,
    • 37.7 with a per-instance instead of per-chunk ratio,
    • 34.9 without pre-training,
    • 35.5 without fine-tuning.
  • Data, model and code are released in the FlagOpen/FlagEmbedding GitHub repository.

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

🎓Renmin University of ChinaBeijing · China114 works · 16,118 citations🔬Beijing Academy of Artificial IntelligenceBeijing · China32 works · 7,511 citations

 Fields, topics and concepts

📚 Shelves 5The field of science this work belongs to (OpenAlex fields).Computer Science · 16,667 worksArtificial Intelligence · 16,194 worksArts and Humanities · 16,039 worksLanguage and Linguistics · 15,891 worksComputer Vision and Pattern Recognition · 2,869 works

🧭 Topics 1Research subjects and languages the work is about.Efficient models & compression · 242 works

Cite

Cite as Peitian Zhang, Zheng Liu, Shitao Xiao, Ninglu Shao, Qiwei Ye and Zhicheng Dou. 2024. Long Context Compression with Activation Beacon. arXiv:2401.03462.

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

@misc{zhang2024long,
  title = {{Long Context Compression with Activation Beacon}},
  author = {Peitian Zhang, Zheng Liu, Shitao Xiao, Ninglu Shao, Qiwei Ye and Zhicheng Dou},
  year = {2024},
  doi = {10.48550/arxiv.2401.03462},
  eprint = {2401.03462},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2401.03462}
}

References

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

  1. Ntk-aware scaled rope, 2023. URL https://www.reddit.com/r/LocalLLaMA/comments/ 14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_have/. ↗
  2. Joshua Ainslie, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. GQA: training generalized multi-query transformer models from multi-head checkpoints. In Houda Bouamor, Juan Pino, and Kalika Bali (eds.), Proceedings of the 2023 Conference on Empirical Methods i…
  3. Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, Yuxiao Dong, Jie Tang, and Juanzi Li. Longbench: A bilingual, multitask benchmark for long context understanding. arXiv preprint arXiv:2308.14508, 2023.
  4. Yushi Bai, Jiajie Zhang, Xin Lv, Linzhi Zheng, Siqi Zhu, Lei Hou, Yuxiao Dong, Jie Tang, and Juanzi Li. Longwriter: Unleashing 10,000+ word generation from long context llms. CoRR, abs/2408.07055, 2024. doi: 10.48550/ARXIV.2408.07055. URL https://doi.org/ 10.48550/arXiv.2408.07055.
  5. Iz Beltagy, Matthew E. Peters, and Arman Cohan. Longformer: The long-document transformer. CoRR, abs/2004.05150, 2020. URL https://arxiv.org/abs/2004.05150. ↗
  6. Sumithra Bhakthavatsalam, Daniel Khashabi, Tushar Khot, Bhavana Dalvi Mishra, Kyle Richardson, Ashish Sabharwal, Carissa Schoenick, Oyvind Tafjord, and Peter Clark. Think you have solved direct-answer question answering? try arc-da, the direct-answer AI2 reasoning challenge. CoRR, abs/2102.03315, 20… — [arXiv:2102.03315](https://arxiv.org/abs/2102.03315) find
  7. William Brandon, Mayank Mishra, Aniruddha Nrusimha, Rameswar Panda, and Jonathan Ragan-Kelley. Reducing transformer key-value cache size with cross-layer attention. CoRR, abs/2405.12981, 2024. doi: 10.48550/ARXIV.2405.12981. URL https://doi.org/10. 48550/arXiv.2405.12981. find
  8. Aydar Bulatov, Yuri Kuratov, and Mikhail S. Burtsev. Scaling transformer to 1m tokens and beyond with RMT. CoRR, abs/2304.11062, 2023. doi: 10.48550/ARXIV.2304.11062. URL https://doi.org/10.48550/arXiv.2304.11062. — [doi:10.48550/arxiv.2304.11062](https://doi.org/10.48550/arxiv.2304.11062) ↗
  9. Shouyuan Chen, Sherman Wong, Liangjian Chen, and Yuandong Tian. Extending context window of large language models via positional interpolation. arXiv preprint arXiv:2306.15595, 2023. find
  10. Yukang Chen, Shengju Qian, Haotian Tang, Xin Lai, Zhijian Liu, Song Han, and Jiaya Jia. Longlora: Efficient fine-tuning of long-context large language models. arXiv preprint arXiv:2309.12307, 2023.
  11. Alexis Chevalier, Alexander Wettig, Anirudh Ajith, and Danqi Chen. Adapting language models to compress contexts. In Houda Bouamor, Juan Pino, and Kalika Bali (eds.), Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, EMNLP 2023, Singapore, December 6-10, 2023, p…
  12. Christopher Clark, Kenton Lee, Ming-Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. Boolq: Exploring the surprising difficulty of natural yes/no questions. In Jill Burstein, Christy Doran, and Thamar Solorio (eds.), Proceedings of the 2019 Conference of the North American Chapte…
  13. Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems. CoRR, abs/2110.14168, 2021. URL https://arxiv.org/abs/2110.1416… ↗
  14. Together Computer. Redpajama: An open source recipe to reproduce llama training dataset, 2023. URL https://github.com/togethercomputer/RedPajama-Data. ↗
  15. Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. CoRR, abs/2307.08691, 2023. doi: 10.48550/ARXIV.2307.08691. URL https://doi.org/10. 48550/arXiv.2307.08691.
  16. DeepSeek-AI. Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model, 2024. find
  17. Yiran Ding, Li Lyna Zhang, Chengruidong Zhang, Yuanyuan Xu, Ning Shang, Jiahang Xu, Fan Yang, and Mao Yang. Longrope: Extending llm context window beyond 2 million tokens, 2024.
  18. Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, Anirudh Goyal, Anthony Hartshorn, Aobo Yang, Archi Mitra, Archie Sravankumar, Artem Korenev, Arthur Hinsvark, Arun Rao, Aston Zhang, Aurelien Rodriguez,… — The Llama 3 Herd of Models
  19. 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, 2020. find
  20. Tao Ge, Jing Hu, Lei Wang, Xun Wang, Si-Qing Chen, and Furu Wei. In-context autoencoder for context compression in a large language model, 2024. find
  21. gkamradt, 2023. URL https://github.com/gkamradt/LLMTest_NeedleInAHaystack. ↗
  22. Chi Han, Qifan Wang, Wenhan Xiong, Yu Chen, Heng Ji, and Sinong Wang. Lm-infinite: Simple on-the-fly length generalization for large language models. CoRR, abs/2308.16137, 2023. doi: 10.48550/ARXIV.2308.16137. URL https://doi.org/10.48550/arXiv.2308.16137. — [doi:10.48550/arxiv.2308.16137](https://doi.org/10.48550/arxiv.2308.16137) ↗
  23. Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL https://…
  24. 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, 2023.
  25. Huiqiang Jiang, Qianhui Wu, Xufang Luo, Dongsheng Li, Chin-Yew Lin, Yuqing Yang, and Lili Qiu. Longllmlingua: Accelerating and enhancing llms in long context scenarios via prompt compression, 2023.
  26. Huiqiang Jiang, Yucheng Li, Chengruidong Zhang, Qianhui Wu, Xufang Luo, Surin Ahn, Zhenhua Han, Amir H. Abdi, Dongsheng Li, Chin-Yew Lin, Yuqing Yang, and Lili Qiu. Minference 1.0: Accelerating pre-filling for long-context llms via dynamic sparse attention, 2024. URL https://arxiv.org/abs/2407.02490… ↗
  27. Ziyan Jiang, Xueguang Ma, and Wenhu Chen. Longrag: Enhancing retrieval-augmented generation with long-context llms. CoRR, abs/2406.15319, 2024. doi: 10.48550/ARXIV.2406. 15319. URL https://doi.org/10.48550/arXiv.2406.15319. — [doi:10.48550/arxiv.2406](https://doi.org/10.48550/arxiv.2406) ↗
  28. Jang-Hyun Kim, Junyoung Yeom, Sangdoo Yun, and Hyun Oh Song. Compressed context memory for online language model interaction, 2024. find
  29. Wojciech Kryściński, Nazneen Rajani, Divyansh Agarwal, Caiming Xiong, and Dragomir Radev. Booksum: A collection of datasets for long-form narrative summarization, 2022. find
  30. Mo Li, Songyang Zhang, Yunxin Liu, and Kai Chen. Needlebench: Can llms do retrieval and reasoning in 1 million context window?, 2024. URL https://arxiv.org/abs/2407.11963. — [arXiv:2407.11963](https://arxiv.org/abs/2407.11963) ↗
  31. Yuhong Li, Yingbing Huang, Bowen Yang, Bharat Venkitesh, Acyr Locatelli, Hanchen Ye, Tianle Cai, Patrick Lewis, and Deming Chen. Snapkv: Llm knows what you are looking for before generation, 2024. URL https://arxiv.org/abs/2404.14469. ↗
  32. Jesse Mu, Xiang Lisa Li, and Noah D. Goodman. Learning to compress prompts with gist tokens. CoRR, abs/2304.08467, 2023. doi: 10.48550/ARXIV.2304.08467. URL https://doi.org/10.48550/arXiv.2304.08467. ↗
  33. OpenAI, Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, Red Avila, Igor Babuschkin, Suchir Balaji, Valerie Balcom, Paul Baltescu, Haiming Bao, Mohammad Bavarian, Jeff Belgum, Irwan Bello, Ja… — GPT-4 Technical Report
  34. Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. Yarn: Efficient context window extension of large language models. arXiv preprint arXiv:2309.00071, 2023. find
  35. Jianlin Su, Yu Lu, Shengfeng Pan, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. CoRR, abs/2104.09864, 2021. URL https://arxiv.org/abs/2104.09864. ↗
  36. Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller,… — Llama 2: Open Foundation and Fine-Tuned Chat Models
  37. Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoe… — Transformers: State-of-the-Art Natural Language Processing
  38. Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. arXiv preprint arXiv:2309.17453, 2023. find
  39. An Yang, Baosong Yang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Zhou, Chengpeng Li, Chengyuan Li, Dayiheng Liu, Fei Huang, Guanting Dong, Haoran Wei, Huan Lin, Jialong Tang, Jialin Wang, Jian Yang, Jianhong Tu, Jianwei Zhang, Jianxin Ma, Jianxin Yang, Jin Xu, Jingren Zhou, Jinze Bai, Jinzheng He, Juny… find
  40. Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. Big bird: Transformers for longer sequences. Advances in neural information processing systems, 33: 17283–17297, 2020.
  41. Zeyu Zhang, Xiaohe Bo, Chen Ma, Rui Li, Xu Chen, Quanyu Dai, Jieming Zhu, Zhenhua Dong, and Ji-Rong Wen. A survey on the memory mechanism of large language model based agents. CoRR, abs/2404.13501, 2024. doi: 10.48550/ARXIV.2404.13501. URL https://doi.org/10.48550/arXiv.2404.13501. ↗
  42. Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark W. Barrett, Zhangyang Wang, and Beidi Chen. H2O: heavy-hitter oracle for efficient generative inference of large language models. In Alice Oh, Tristan Naumann, Amir Globers… find
  43. Qianchao Zhu, Jiangfei Duan, Chang Chen, Siran Liu, Xiuhong Li, Guanyu Feng, Xin Lv, Huanqi Cao, Xiao Chuanfu, Xingcheng Zhang, Dahua Lin, and Chao Yang. Sampleattention: Near-lossless acceleration of long context LLM inference with adaptive structured sparse attention. CoRR, abs/2406.15486, 2024. d… — [doi:10.48550/arxiv.2406.15486](https://doi.org/10.48550/arxiv.2406.15486) find
  44. Zirui Liu, Jiayi Yuan, Hongye Jin, Shaochen Zhong, Zhaozhuo Xu, Vladimir Braverman, Beidi Chen, and Xia Hu. Kivi : Plug-and-play 2bit kv cache quantization with streaming asymmetric quantization. 2023. doi: 10.13140/RG.2.2.28167.37282. URL https://rgdoi.net/10. 13140/RG.2.2.28167.37282. ↗

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/zhang-2024-long-context-compression-activation-beacon/