💻 computer science

One Vector Is All You Need for O(1) Self-Attention: The Ocean State

本文介绍了“海洋状态”（Ocean State），这是一种通过使用单个持久向量来取代标准自注意力机制的方法，旨在实现 O(1) 的计算与内存复杂度，同时展示了卓越的训练稳定性，且在 1000 万步上下文长度下也未出现遗忘迹象。

原作者： SHUYUAN YU

发布于 2026-09-04

📖 1 分钟阅读☕ 轻松阅读

原作者： SHUYUAN YU

原始论文采用 CC BY 4.0 许可（https://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

现代人工智能通常依赖于一种被称为“Transformer”的特定类型的计算机程序，它已成为处理文本生成、语言翻译和回答问题等任务的标准。这些程序通过观察一个词序列并决定下一个词是什么来工作。为了准确地做到这一点，程序必须记住它已经看到的词。在当前的设计中，当程序阅读长句子或整本书时，它会保留一份不断增长的已处理词汇列表。这个列表就像一个随每一个新词而扩张的记忆库。虽然这在处理短文本时效果很好，但对于非常长的文本来说，它会变成沉重的负担。计算机必须不断扫描这个不断增长的列表以寻找相关信息，这需要巨大的计算能力和内存。随着文本变长，阅读文本所需的时间和能量增长速度远快于文本本身的长度，最终使得高效处理超长文档变得不可能。

一位名叫 Yu Shuyuan 的研究人员针对这一记忆问题提出了另一种处理方法。该新方法建议不再保留一个不断增长的词汇列表，而是仅保留一个单一的摘要向量，作者称之为“海洋状态”（ocean state）。想象一下，这个状态就像一个单一且稠密的容器，承载着程序目前为止所读到的一切精华。当程序阅读一个新词时，它会更新这个单一的容器以包含新信息，用新鲜的摘要替换旧的摘要。然后，程序使用这个单一的容器来辅助预测下一个词。这种方法改变了问题的基本数学逻辑：处理过程的难度不再随文本长度增加，无论文本变得多长，其投入的精力都保持不变。程序可以阅读一百万个词，而每个词的处理量保持一致，并且无论文本长度如何，它使用的内存始终是固定且微小的。

这项工作的核心思想是一个简单的疑问：如果我们不在使用一个句子的最终摘要后将其丢弃，情况会怎样？在标准程序中，这个摘要被计算出来后就会被丢弃。而新方法保留了它，将其视为一种从一步流向下一步的持久记忆。程序读取当前的词和这个单一的摘要向量，将它们结合起来，并产生一个新的摘要。这个过程在文本中的每一个词上都会重复进行。至关重要的是，程序完全能够自主学习如何读取并更新这个单一的向量。没有人告诉它如何将信息压缩到一个点中；它只是在训练过程中通过设计本身学会了这样做。研究人员发现，这个单一向量的能力惊人。它可以在十百万步的序列中保存信息而不丢失任何内容。在测试中，该程序能够像检索前几步的信息一样准确地召回序列开头特定词汇的信息，没有表现出任何遗忘迹象。

研究人员使用一组大型文本数据，针对这一想法与标准方法进行了对比测试。他们训练了两个版本的模型来预测句子中的下一个词。结果显示，使用单一向量的新方法在表现上始终优于保留增长列表的标准方法。这种优势并非偶然，它出现在包括不同模型规模和不同文本长度在内的多种设置中。事实上，由于新方法非常稳定，它可以处理十百万步的序列而不会出现误差增加的情况，而标准方法甚至在处理不到这个长度的片段时，就会耗尽内存或耗费无法承受的时间。研究人员还测试了一项特定任务，即程序需要在长时间延迟后重复一个词。新方法重现该词时的损失仅为 0.0006，而标准方法则完全失败。

人们对这种压缩记忆的一个担忧是，它可能会随着时间的推移变得难以解读或变得“模糊”，从而失去进行良好预测所需的细节。然而，实验结果却恰恰相反。这个单一向量始终保持清晰锐利，即使在经过数百万步之后仍能检索特定信息。研究人员还探索了如何让这种方法在真实计算机上运行得更快。由于新方法是严格按顺序逐个处理单词，因此与可以同时处理多个单词的标准方法相比，其在单机上的训练速度可能较慢。为了解决这个问题，研究人员开发了一种将工作分配到程序不同层级的方法，从而允许同时处理多个步骤。这种工程上的改进显著减少了训练模型所需的时间，同时保持了同样的高性能。

这项研究证明，利用固定的内存和计算能力来构建一个能够记住海量信息的系统是可能的。研究人员展示了程序如何在没有任何特殊指令的情况下，学会高效地组织这些信息。虽然目前的实验是在特定的数据集和相对较小的模型上进行的，且在现实世界规模下的表现尚未经过测试，但结果表明其底层原理是稳健的。该方法通过重新定义程序处理其内部状态的方式来起作用，将一个被丢弃的摘要转变为一种持久的记忆。这种改变消除了限制人工智能处理文本长度的结构性成本。研究结果表明，通过这种方法，上下文长度的限制可能不再是障碍，使得系统能够以一致的速度和准确度阅读并理解任何规模的文档。

技术摘要：一个向量即是全部——实现 $O(1)$ 自注意力的“海洋状态”（The Ocean State）

1. 问题陈述

标准 Transformer 存在一个结构性限制：自注意力会产生 $O(N^2)$ 的计算成本，并且需要随上下文长度线性增长的 $O(N)$ 内存缓存（键值对缓存/KV Cache）。虽然最近的研究（如 RetNet, RWKV, Mamba, Titans）试图通过引入线性注意力、选择性门控或外部存储模块等新机制来缓解这一问题，但它们从根本上改变了架构。此外，标准 Transformer 在每次前向传播后都会丢弃最终的隐藏状态，尽管该状态包含了整个上下文的吸收信息。

本文认为， $O(N^2)$ 的成本是当前实现方式的结构性特征，而非上下文建模概念本身的固有属性。它提出了一个问题：如果隐藏状态（它吸收了整个上下文）不被丢弃，而是被保留并作为下一个 token 的唯一注意力对象，情况会怎样？

2. 方法论：海洋状态（The Ocean State）

提出的解决方案——“海洋状态”——涉及对标准 Transformer 架构进行两种特定的重新定位，且不引入任何新组件（例如：没有门控，没有外部存储）。

2.1 核心机制

海洋状态（重新定位 1）：
- 不再在每次前向传播后丢弃隐藏状态，而是将当前 token 处理后的输出替换之前的隐藏状态，该状态被称为“海洋状态”。
- 注意力机制： 注意力严格在当前 token 与持久的海洋状态（一个稠密向量）之间进行计算。
- 后果： 这消除了对键值对（KV）缓存、位置编码和因果掩码的需求。状态通过“接力”（1:1 替换）而非累加进行更新，确保信息被重新排列而非被压缩或丢弃。
- 只读性： 海洋状态被当前 token 读取，但从不经过前馈网络（FFN）或层堆栈的处理；只有 token 本身会通过层堆栈。
海洋 BPTT（重新定位 2）：
- 由于海洋状态是一个递归变量，因此必须通过时间进行展开（随时间反向传播，BPTT）。
- 梯度流： 在反向传播期间，第 $t$ 步海洋状态的梯度是基于第 $t$ 步的 token 计算的。该梯度随后被添加到第 $t-1$ 步的隐藏状态梯度中。
- 读取路径算子： 一个特定的算子利用前向传播中保存的注意力权重，将 token 的梯度映射回海洋状态。一个 token 对海洋的“读取”程度越高，发送回更新海洋状态的梯度信号就越强。

2.2 算法实现

训练循环包含：

前向传播： Token $t$ 关注 $ocean_{t-1}$ 。Token 通过层堆栈（注意力 + FFN）。输出替换 $ocean_{t-1}$ 成为 $ocean_t$ 。
反向传播： 梯度从第 $t$ 步的损失流经层堆栈回到隐藏状态。至关重要的是，海洋状态的梯度会跨时间步累积，允许来自未来 token 的误差来更新过去上下文的表示。

2.3 工程变体

论文探索了三种训练方案，以平衡成本、时间与质量：

串行海洋 BPTT (Serial Ocean BPTT)： 全时间链展开。质量增益最高，计算成本最低，但实际运行时间（wall-clock time）为串行。
框架内海洋 (Ocean on Framework)： 在标准自动微分框架上运行完整链条。保持全额质量增益，并通过批处理提高吞效率。
分离海洋 (Ocean Detach)： 切断时间梯度链（无跨步 BPTT）。可以进行全并行训练（类似于标准 Transformer），但牺牲了源自时间梯度的性能增益。

3. 核心贡献

架构重定位： 证明了通过重用被丢弃的隐藏状态作为持久的“海洋”向量，可以实现 $O(1)$ 注意力和 $O(1)$ 内存，从而消除了对 KV 缓存和位置编码的需求。
涌现能力： 展示了 Transformer 可以在没有任何显式指令的情况下，仅凭端到端训练，学会从单个稠密向量中读取和提取信息。
极端上下文的可扩展性： 证明了该架构可以处理高达 1000 万步的上下文长度而无性能退化，这对于受限于内存和计算约束的标准 Transformer 而言是不可能完成的任务。
梯度展开： 引入了“海洋 BPTT”，将标准 BPTT 应用于整个层堆栈输出（即递归状态），并指出这是性能提升的主要来源。

4. 实验结果

实验在 Shakespeare 语料库（字符级和单词级）上进行，使用了自定义的 NumPy 实现（float64）并使用 PyTorch 进行了验证。

4.1 性能与训练动态

优越性： 在所有测试配置（宽度 $d \in \{32, 64, 128, 256\}$ ，深度 $L \in \{2, 4, 6\}$ 以及各种头数）下，“Ocean Full”模式始终优于标准 Transformer。
稳定的增益： 性能差距（ $\Delta = Loss_{std} - Loss_{full}$ ）为正且稳定（约减少 0.15 的 loss），在整个网格测试中均如此。
增益来源： 这种增益主要是结构性的。当时间链被分离时（Ocean Detach），性能显著下降，这表明时间梯度展开是关键因素，而非仅仅是注意力机制本身。

4.2 长上下文与内存

无遗忘： 在一个 1000 万步的预填充（prefill）测试中，损失保持平稳（100 万步时为 1.532，1000 万步时为 1.531）。海洋状态没有表现出容量饱和或遗忘的迹象。
回声任务（Echo Task）： 模型在具有 200 万步延迟的回声任务中实现了接近零的损失（0.0006），而标准 Transformer 则完全失败。
内存效率： 海洋状态仅需 $O(1)$ 内存（一个维度为 $d$ 的向量），与上下文长度无关。相比之下，标准 Transformer 需要 $O(N)$ 内存用于 KV 缓存，这在 100 万个 token 的规模下变得不可行（对于 13B 模型，大约需要 1 TB 内存）。

4.3 推理与生成

无需外推： 海洋状态生成的文本具有连贯性，不需要位置编码外推（如 RoPE, ALiBi），因为上下文长度并不是注意力机制的一个输入变量。
实际运行时间： 虽然递归的串行性质使得其每步训练比并行的标准 Transformer 慢，但论文指出，通过层级流水线（Ocean A-2）可以实现 1.9 倍的实际运行时间加速，而嵌套并行（Ocean A-3）在理论上可提供 $10^3$ 倍的增益。作者强调，训练涉及一个关于成本、时间与质量的显式权衡三角形，实际应用中需要选择特定的方案（例如，为了工业规模选择 detach，为了研究质量选择 full chain）。

5. 意义与主张

论文提出了三个主要观点，并将其定义为可证伪的：

训练一致性： 海洋状态的训练曲线形状与标准 Transformer 相同，证实了这是一种有效的重定位，而非一种全新的学习范式。
普遍增益： 对于任何在相同条件下的 Transformer 架构，海洋状态都能产生正向的性能增益（ $\Delta > 0$ ）。
无限容量： 在足够的维度（ $d$ ）下，海洋状态表现出没有容量限制、没有遗忘以及在任何可观测规模下都不会出现检索失败。

作者总结道，“海洋状态”解决了任意上下文长度下的 $O(1)$ 推理与训练原则，并表示：“我们相信该原则已被解决。”通过依赖 Transformer 在单个向量内组织信息的涌现能力，该方法消除了二次方注意力和线性内存增长的结构性瓶颈，可能使“小型且个性化”的 Transformer 以及拥有连续认知状态的人形机器人成为现实。这项工作强调，这些结果源于对现有组件的“假设性”实验，尊重了原始架构的涌现属性，同时也承认了针对特定配置（例如，为什么 detach 在某些情况下表现更好）仍存在开放性问题。

技术摘要：一个向量即是全部——实现 O(1)O(1)O(1) 自注意力的“海洋状态”（The Ocean State）