Trellis让注意力模型动态压缩记忆，解决长序列计算负担问题。
Trellis: Learning to Compress Key-Value Memory in Attention Models
- 用固定大小内存替代传统缓存，通过递归压缩机制存储关键信息。
- 在长序列任务中表现优于基线，序列越长优势越明显。
- 适合需要长期记忆的场景，如长文本生成与时间序列建模。
Transformer 模型虽强大，但存在二次计算复杂度及不断增长的键值（KV）缓存问题。本文提出 Trellis，一种具有有界记忆的新架构，在测试时动态学习压缩其键值记忆。Trellis 将标准 KV 缓存替换为固定大小的记忆体，并训练一种两阶段递归压缩机制，将新键值存入记忆体。为此，它采用带遗忘门的在线梯度下降方法，使压缩记忆能在测试时递归更新，同时学习保留来自输入标记的重要上下文信息。在语言建模、常识推理、需记忆的任务和时间序列上的大量实验表明，该架构显著优于强基线。尤其值得注意的是，其性能提升随序列长度增加而增大，凸显其在长上下文应用中的潜力。
原文摘要 · Abstract (English)
Transformers, while powerful, suffer from quadratic computational complexity and the ever-growing Key-Value (KV) cache of the attention mechanism. This paper introduces Trellis, a novel Transformer architecture with bounded memory that learns how to compress its key-value memory dynamically at test time. Trellis replaces the standard KV cache with a fixed-size memory and train a two-pass recurrent compression mechanism to store new keys and values into memory. To achieve this, it leverages an online gradient descent procedure with a forget gate, enabling the compressed memory to be updated recursively while learning to retain important contextual information from incoming tokens at test time. Extensive experiments on language modeling, common-sense reasoning, recall-intensive tasks, and time series show that the proposed architecture outperforms strong baselines. Notably, its performance gains increase as the sequence length grows, highlighting its potential for long-context applications.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。