[论文解读] Not All Memories are Created Equal: Learning to Forget by Expiring

Sainbayar Sukhbaatar, Da Young Ju+5 more|arXiv (Cornell University)|May 13, 2021

Topic Modeling参考文献 47被引用 5

一句话总结

本文提出 Expire-Span，一种针对 Transformer 的可微分遗忘机制，可学习淘汰无关记忆，实现在数万个时间步上的高效注意力计算。通过动态淘汰低重要性内容，该方法在长上下文任务（如字符级语言建模和逐帧目标追踪）中达到最先进性能，同时训练速度比基线快3倍，内存使用更少。

ABSTRACT

Attention mechanisms have shown promising results in sequence modeling tasks that require long-term memory. Recent work investigated mechanisms to reduce the computational cost of preserving and storing memories. However, not all content in the past is equally important to remember. We propose Expire-Span, a method that learns to retain the most important information and expire the irrelevant information. This forgetting of memories enables Transformers to scale to attend over tens of thousands of previous timesteps efficiently, as not all states from previous timesteps are preserved. We demonstrate that Expire-Span can help models identify and retain critical information and show it can achieve strong performance on reinforcement learning tasks specifically designed to challenge this functionality. Next, we show that Expire-Span can scale to memories that are tens of thousands in size, setting a new state of the art on incredibly long context tasks such as character-level language modeling and a frame-by-frame moving objects task. Finally, we analyze the efficiency of Expire-Span compared to existing approaches and demonstrate that it trains faster and uses less memory.

研究动机与目标

• 解决长上下文 Transformer 因存储所有历史信息而导致的效率低下与可扩展性限制问题。
• 使模型能够选择性遗忘无关记忆，模仿人类记忆中优先保留关键信息的能力。
• 通过学习动态记忆淘汰而非固定压缩或注意力范围，提升长上下文任务的性能。
• 在不牺牲性能的前提下，降低长序列建模中的计算与内存开销。

提出的方法

• Expire-Span 通过可学习的过期预测器扩展自注意力机制，为每个隐藏状态输出一个保留跨度。
• 每个记忆在过期跨度结束后被清除，且该过程可微分，支持端到端反向传播。
• 过期机制按层独立应用，使不同层可专门处理不同时间尺度的记忆保留。
• 通过早期丢弃无关信息，实现在数万个时间步上的高效注意力计算。
• 模型使用交叉熵损失进行端到端训练，且过期机制集成于注意力计算中。
• 该方法应用于仅解码器的 Transformer，实现可扩展的长上下文建模，无需固定压缩或稀疏注意力模式。

实验结果

研究问题

• RQ1 可学习的遗忘机制是否能提升长上下文序列建模任务的性能？
• RQ2 与固定跨度或压缩记忆方法相比，动态记忆淘汰如何影响训练效率与内存使用？
• RQ3 模型能否学习保留关键信息（如命名实体）而遗忘通用或无关内容？
• RQ4 长期记忆容量在多大程度上影响长上下文基准上的预测准确率？
• RQ5 遗忘记忆的能力是否能实现更长上下文长度的扩展，而不会导致计算量爆炸？

主要发现

• Expire-Span 在 Enwik8 字符级语言建模基准上达到最先进性能，优于 Compressive Transformer 和 Adaptive-Span 基线模型。
• 该模型训练速度接近基线的 3 倍，同时保持更高准确率。
• 在最大跨度为 64k 的情况下，Expire-Span 顺利训练并避免了 GPU 内存溢出，而 Adaptive-Span 在 32k 时已失败。
• 模型学习到长期保留稀有或关键实体（如 'Egypt' 和 'Humpty Dumpty'），而快速淘汰通用术语（如 'somewhere'）。
• 当推理跨度从 16k 降低至 1k 时，对 'Guinea coast' 等长距离依赖预测的准确率显著下降，该词在 3,584 个时间步前被提及。
• 分析显示，结构元素（如章节标题和换行符）被保留更长时间，表明各层在记忆使用上存在特定分工。