Not All Memories are Created Equal: Learning to Forget by Expiring
Sainbayar Sukhbaatar Da Ju Spencer Poff Stephen Roller Arthur Szlam Jason Weston Angela Fan
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 (Rae et al., 2020). 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
中文速览
让Transformer学会"遗忘"是这项工作的核心出发点——现有的长序列注意力机制会把所有历史状态都保留下来，随着上下文变长，计算和存储开销急剧增加，而且大量无关信息混杂其中反而干扰模型判断。研究者提出了Expire-Span方法，为每个历史隐藏状态预测一个"过期时间"，过了这个时间点就用可微的软掩码把该记忆渐进地清零并永久删除，同时加入L1正则损失驱动模型主动缩短不重要记忆的寿命。这样一来，模型能自动识别并留住关键信息、丢弃无关内容，使注意力可以高效地覆盖数万个历史时间步，而无需像以往方法那样对所有记忆一视同仁地压缩或稀疏化。实验显示，Expire-Span在专门测试长程记忆的强化学习任务上表现出色，在字符级语言建模和逐帧运动物体预测等超长上下文基准上刷新了最优成绩，同时训练速度更快、显存占用更少，为大规模序列建模提供了一条兼顾效率与效果的新路径。
原文 arXiv:2105.06548;中英对照 + 大白话阅读 https://aha.fim.ai/paper/2105.06548v2