We propose the In-context Autoencoder (ICAE), leveraging the power of a large language models (LLM) to compress a long context into short compact memory slots that can be directly conditioned on by the LLM for various purposes. ICAE is first pretrained using both autoencoding and language modeling objectives on massive text data, enabling it to generate memory slots that accurately and comprehensively represent the original context; Then, it is fine-tuned on instruction data for producing desirable responses to various prompts. Experiments demonstrate that our lightweight ICAE, introducing about 1% additional parameters, effectively achieves $4\times$ context compression based on Llama, offering advantages in both improved latency and GPU memory cost during inference, and showing an interesting insight in memorization as well as potential for scalability. These promising results imply a novel perspective on the connection between working memory in cognitive science and representation learning in LLMs, revealing ICAE's significant implications in addressing the long context problem and suggesting further research in LLM context management. Our data, code and models are available at https://github.com/getao/icae.
翻译:我们提出上下文自编码器(ICAE),利用大语言模型(LLM)的能力将长上下文压缩为紧凑的存储槽,LLM可直接基于这些存储槽实现多种用途。ICAE首先通过自编码和语言建模目标在大规模文本数据上进行预训练,使其能够生成准确且完整表征原始上下文的存储槽;随后在指令数据上进行微调,以对各类提示产生理想响应。实验表明,轻量化的ICAE仅引入约1%的额外参数,即可基于Llama实现$4\times$的上下文压缩,在推理延迟和GPU内存成本方面均具有显著优势,同时展现出记忆特性与扩展潜力的有趣洞见。这些令人鼓舞的结果为认知科学中的工作记忆与LLM中的表征学习之间建立了新颖关联,揭示了ICAE在解决长上下文问题中的重要意义,并为LLM上下文管理研究开辟了新方向。相关数据、代码与模型已开源至https://github.com/getao/icae。