arXiv:2410.07145cs.CLcs.AI2024-10被引 8

Mamba模型虽快却难忘旧信息,因训练太短导致遗忘机制失效。

Stuffed Mamba: Oversized States Lead to the Inability to Forget

  • 用过长的上下文训练,让模型学不会主动遗忘旧信息。
  • 状态尺寸越大,需更长训练才能学会遗忘,最长记忆可达5位密码。
  • 提醒设计新RNN时要同步考虑状态大小与训练长度。

最近的循环架构(如Mamba和RWKV)展现出强大的语言能力。与Transformer不同,这些架构将所有上下文信息编码为固定大小的状态,从而实现高效推理。然而,这种设计可能导致信息干扰,即不同词元数据相互冲突,造成性能下降和超出一定上下文长度后的输出不连贯。为防止此问题,大多数RNN引入了“遗忘”机制。本文揭示,即使有内置遗忘机制,基于Mamba的模型仍难以有效遗忘早期词元。我们发现,这一问题源于在过短的上下文上训练,使模型无需学习遗忘即可表现良好。进一步表明,模型学会遗忘所需的最小训练长度与状态尺寸呈线性关系;而准确检索5位密码的最大上下文长度则随状态尺寸呈指数增长,说明模型在遗忘开始后仍保留部分信息。这些结果凸显当前RNN架构的关键局限,并为改进长上下文建模提供了重要启示。我们的工作建议未来RNN设计必须综合考虑状态尺寸、训练长度与遗忘机制之间的相互作用,以在长上下文任务中实现稳健性能。

原文摘要 · Abstract (English)

Recent advancements in recurrent architectures, such as Mamba and RWKV, have showcased strong language capabilities. Unlike transformer-based models, these architectures encode all contextual information into a fixed-size state, leading to great inference efficiency. However, this approach can cause information interference, where different token data conflicts, resulting in performance degradation and incoherent outputs beyond a certain context length. To prevent this, most RNNs incorporate mechanisms designed to "forget" earlier tokens. In this paper, we reveal that Mamba-based models struggle to effectively forget earlier tokens even with built-in forgetting mechanisms. We demonstrate that this issue stems from training on contexts that are too short for the state size, enabling the model to perform well without needing to learn how to forget. Then, we show that the minimum training length required for the model to learn forgetting scales linearly with the state size, and the maximum context length for accurate retrieval of a 5-digit passkey scales exponentially with the state size, indicating that the model retains some information beyond the point where forgetting begins. These findings highlight a critical limitation in current RNN architectures and provide valuable insights for improving long-context modeling. Our work suggests that future RNN designs must account for the interplay between state size, training length, and forgetting mechanisms to achieve robust performance in long-context tasks.

长序列建模遗忘机制RNN改进Mamba

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。