[论文解读] In-context Autoencoder for Context Compression in a Large Language Model

Tao Ge, Jing Hu+4 more|arXiv (Cornell University)|Jul 13, 2023

Speech Recognition and Synthesis被引用 7

一句话总结

本文提出了一种轻量级方法——上下文自编码器（ICAE），利用LoRA微调的大语言模型将长上下文压缩为紧凑的记忆槽，实现最高达4倍的上下文压缩，从而在保持性能的同时显著降低延迟和显存占用，为大模型记忆机制与上下文管理提供了新见解。

ABSTRACT

We propose the In-context Autoencoder (ICAE), leveraging the power of a large language model (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 imes$ 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.

研究动机与目标

• 通过将输入序列压缩为更短的、信息丰富的记忆表示，解决大语言模型中的长上下文问题。
• 通过用紧凑的记忆槽替代长上下文，实现更低延迟和GPU显存成本的高效推理。
• 探究自监督预训练如何增强大语言模型的记忆能力，类比人类工作记忆机制。
• 开发一种参数高效、即插即用的解决方案，兼容现有大语言模型，无需架构修改。
• 探究记忆槽在下游任务（如思维链推理和任务特定压缩）中的潜力。

提出的方法

• ICAE使用LoRA微调的大语言模型作为可学习编码器，将长输入上下文压缩为少量记忆槽。
• 解码器即为目标大语言模型本身，其基于生成的记忆槽进行响应生成或原始上下文重建。
• 模型在自编码和语言建模目标上联合预训练，确保记忆槽兼具重建能力与语义意义。
• 在指令数据上进行微调可提升记忆槽质量，适用于多样化提示场景，并增强与提示的交互能力。
• 记忆槽通过特殊标记格式和可学习嵌入进行学习，使模型能够生成紧凑且高保真的表示。
• 该方法参数效率高，仅通过LoRA增加原始模型参数的1%以下。

实验结果

研究问题

• RQ1 大语言模型能否在不修改架构的前提下，有效将长输入上下文压缩为一组少量记忆槽？
• RQ2 此类记忆槽在多大程度上能保留原始上下文的语义和结构信息？
• RQ3 自监督预训练在多大程度上影响模型记忆和重建长上下文的能力？
• RQ4 在真实场景中，上下文压缩比与推理性能之间存在何种权衡？
• RQ5 记忆槽能否在多样化下游任务中有效应用，包括思维链推理和任务特定提示？

主要发现

• ICAE在Llama模型上实现了最高达4倍的上下文压缩，显著降低GPU显存占用和推理延迟。
• 使用2048个记忆槽替代4096个上下文标记，性能下降可忽略，节省约20GB GPU显存。
• 自编码与语言建模联合预训练显著增强了模型记忆和重建长上下文的能力。
• 该方法与现有大语言模型具有强兼容性，仅通过LoRA微调增加1%以下的参数。
• 记忆槽在思维链推理中表现有效，且可进一步微调以实现任务特定压缩。
• 研究揭示了大语言模型记忆与人类工作记忆之间存在强烈类比，提示其在表征学习中可能存在共享机制。