Proteus: Incremental Memory Activation for Long-Context Sequence Modeling¶
会议: NeurIPS 2026
arXiv: 2608.16844
OpenReview: https://openreview.net/forum?id=BJtTMkeGhz
领域: LLM 效率
关键词: 循环神经网络, 联想记忆, 渐进记忆激活, 长上下文建模, 动态容量调度
一句话总结¶
Proteus 针对固定容量循环记忆在序列早期过度拟合、后期严重干扰与覆盖的痛点，提出在上下文推进过程中逐步激活记忆参数与前馈权重的渐进激活范式，以零额外参数与计算开销显著增强长上下文检索与泛化性能。
研究背景与动机¶
基于自注意力机制的 Transformer 架构虽然在各类自然语言任务中取得了巨大成功，但其注意力矩阵与无损 KV 缓存带来的序列长度二次方计算与线性显存开销，严重制约了超长上下文的处理能力。为此，现代线性循环神经网络与基于联想记忆的状态模型（如 RetNet、DeltaNet、Titans 以及 TTT）重新受到广泛关注。这类模型通过将历史上下文在线压缩进常数尺寸的循环隐藏状态中，实现了恒定的单步推理开销与线性的序列级计算复杂度。然而，将无限上下文压缩进有限维度的记忆状态必然带来信息瓶颈，现有循环架构在超出训练长度时往往表现出剧烈的性能衰退，难以胜任长程检索与依赖精确召回的任务。
现存记忆增强模型几乎无一例外地在整条序列处理期间暴露静态且恒定的记忆容量。深入分析在线记忆更新机制可以发现，这一静态设定存在严重的结构性失衡：在序列初始阶段，记忆状态几乎为空，前序 token 在没有任何容量竞争与压缩压力的情况下，无节制地占据了模型状态的大量自由度，从而倾向于“死记硬背”早期局部表面特征，导致记忆空间被严重“污染”；而当序列推进到中后期，后续到达的丰富信息只能挤入剩余的狭窄空间，向高度饱和的状态写入新内容不可避免地覆盖既有记忆并引发剧烈表征干扰。
这种早期压缩不足与晚期容量饱和的双重困境，正是导致模型在长上下文检索中发生“迷失”（lost in the middle）与外推失败的本质根源。若要打破这一僵局，就必须改变“何时激活容量”的时间分配策略。核心 idea：打破静态记忆暴露机制，引入渐进记忆激活范式（Proteus），在序列早期施加硬性容量瓶颈倒逼模型对历史信息进行高维提炼与紧凑摘要，随着上下文长度增长分阶段解锁全新的未污染记忆块，为后续 token 提供充足的新生自由度以消除覆盖干扰。
方法详解¶
整体框架¶
Proteus 将序列建模统一纳入在线联想记忆优化视角，把隐状态更新视作根据键值对映射在线拟合内部损失函数的过程。输入序列经过线性变换生成查询 \(q_t\)、键 \(k_t\) 与值 \(v_t\)。传统架构将全量记忆参数 \(\mathcal{M}\) 始终暴露给每一步的读写更新，而 Proteus 则引入一个依赖输入绝对位置的门控激活算子 \(\mathcal{G}_t\)，使得任意时刻只有记忆状态的前缀子集被激活，未激活参数被严格锁定且不参与前向检索与反向梯度更新。
整个处理管线由位置调度器与分块门控算子协调运转：在序列步进至位置 \(t\) 时，系统根据全局调度规则计算当前激活的块数量 \(k(t)\)，构造二值掩码向量 \(g_t\)。记忆读取仅作用于门控后的激活子空间；在线写入时，惊喜信号（内部目标关于激活态的梯度）仅更新激活块，未激活块保持原值封存，直至上下文长度推进到对应扩容区间才被平滑解锁。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
In["当前输入 Token<br/>x_t (映射生成 q_t, k_t, v_t)"] --> Sched["均匀步进扩容调度<br/>根据位置 t 计算激活块数 k(t) 与掩码 g_t"]
Sched --> Gate["分块前缀门控机制<br/>锁定未激活块，暴露前缀激活块"]
Gate --> Read["门控记忆读取<br/>y_t = Read(g_t ⊙ M_{t-1}, q_t)"]
Gate --> Write["门控在线写入<br/>M_t = M_{t-1} + g_t ⊙ δM_t"]
Read --> Out["输出表征与下游预测 y_t"]
Write --> State["更新后的状态 M_t (供下一位置使用)"]关键设计¶
1. 均匀步进扩容调度：随位置单调递增有效容量以平衡压缩瓶颈与容量储备 针对静态容量下早期欠压缩与晚期强干扰的核心矛盾，该设计旨在建立一种显式的容量扩张时间表。设预训练的最大上下文窗口为 \(N\)，将总记忆参数划分为 \(E\) 个大小均等的连续块（Block），单个块的参数维度为 \(d = \dim(\mathcal{M}) / E\)。调度器定义激活步长为： $\(\Delta = \max\left(1, \left\lfloor \frac{N}{E} \right\rfloor\right)\)$ 在序列第 \(t\) 步（\(1 \le t \le N\)），激活的块数量 \(k(t)\) 严格满足： $\(k(t) = \min\left(E, 1 + \left\lfloor \frac{t - 1}{\Delta} \right\rfloor\right)\)$ 由此，在区间 \([1, \Delta]\) 内，模型仅能使用前 \(1/E\) 的记忆容量，严酷的容量限制迫使模型学习对早序 token 提取高度压缩的抽象摘要；当序列进入 \([(i-1)\Delta + 1, i\Delta]\) 时，第 \(i\) 个记忆块被激活，释放出全新的、未被先前写入破坏的自由度来接纳最新输入。当序列长度超过 \(N\) 时，\(k(t) = E\)，全量容量完全解锁，在长文本外推阶段不再引入额外分块干扰。
2. 分块前缀门控机制：读写双向严格隔离的无开销参数屏蔽 为了在现存硬件友好的矩阵记忆或神经网络记忆上以零计算代价实现容量受限，Proteus 构建了一个离散前缀指示掩码 \(g_t \in \{0, 1\}^{\dim(\mathcal{M})}\)： $\(g_t[j] = \begin{cases} 1, & 1 \le j \le k(t) \cdot d \\ 0, & \text{其他} \end{cases}\)$ 在记忆读取阶段，读取函数仅对激活子空间计算：\(y_t = \text{Read}(\mathcal{G}_t(\mathcal{M}_{t-1}), q_t) = \text{Read}(g_t \odot \mathcal{M}_{t-1}, q_t)\)。在记忆更新阶段，针对内部惊喜目标 \(\tilde{\mathcal{L}}\) 的在线更新被限制在掩码指示的子区域内： $\(\mathcal{M}_t = \mathcal{M}_{t-1} - \theta_t \left( g_t \odot \nabla \tilde{\mathcal{L}}\left(g_t \odot \mathcal{M}_{t-1}; k_t, v_t\right) \right)\)$ 未激活区域满足 \((\mathbf{1} - g_t) \odot \mathcal{M}_t = (\mathbf{1} - g_t) \odot \mathcal{M}_{t-1}\)，这意味着未激活块在被唤醒前其数值完全保持初始状态，既不产生读取串扰，也不承担任何前序梯度的耗散。该机制不改变底层 Hebbian 规则、Delta 规则或带有动量衰减的 Titans 规则的数学本质，具有完全的架构正交性与即插即用特性。
3. 嵌套学习参数扩容延伸：将记忆激活原则推广至前馈层参数轨迹 基于嵌套学习（Nested Learning）理论，神经网络中反向传播更新的 MLP 前馈层参数本质上也是一种以样本误差为内部目标的联想记忆系统。为将渐进激活思想延拓至无显式循环记忆的混合模型（如 Hope-Attention），Proteus 将该调度机制应用于模型权重的在线更新过程中。在优化器（如 AdamW）更新阶段，引入样本步级别的调度掩码 \(g_{i+1}\) 作用于梯度与各阶动量： $\(m_{i+1} = \beta_1 m_i + (1 - \beta_1) \left( g_{i+1} \odot \nabla_{\theta_i} L(\theta_i, x_{i+1}) \right)\)$ $\(v_{i+1} = \beta_2 v_i + (1 - \beta_2) \left( g_{i+1} \odot \nabla_{\theta_i} L(\theta_i, x_{i+1}) \right)^2\)$ 只有被选中的前缀参数切片在当前训练步吸收梯度，使得模型在优化初始阶段避免过度参数化带来的低秩表面拟合，保留深层容量以持续吸收长序列复杂依赖。
损失函数 / 训练策略¶
模型采用标准的自回归下文预测交叉熵损失进行端到端预训练。所有模型在 FineWeb 数据集上统一以 8K 上下文窗口进行训练，分为 760M 参数（训练 50B tokens）与 1.3B 参数（训练 100B tokens）两级规模。优化器采用 AdamW，最大学习率设为 \(4 \times 10^{-4}\)，配合余弦退火学习率衰减策略，全局 batch size 为 0.5M tokens，权重衰减系数为 0.1。分块超参数在默认状态下设为 \(E = 16\)。测试阶段直接继承训练时拟定的 \(N=8\text{K}\) 调度表，超出 8K 的位置全部默认激活 \(E\) 个分块。
实验关键数据¶
主实验¶
在标准语言建模（Wikitext 与 LAMBADA 困惑度）及 8 项常识推理基准（PIQA、HellaSwag、WinoGrande、ARC-e、ARC-c、SIQA、BoolQ 以及 LAMBADA 准确率）上，Proteus 在所有基线骨干模型上均取得了一致且显著的提升。
| 模型骨干 | 规模 / Tokens | Wiki ppl ↓ | LMB ppl ↓ | 常识推理平均准确率 (Avg.) ↑ |
|---|---|---|---|---|
| Transformer++ | 760M / 50B | 22.08 | 22.41 | 51.56 |
| Hope-Attention | 760M / 50B | 20.67 | 21.36 | 53.15 |
| Hope-Attention + Proteus | 760M / 50B | 19.87 | 19.72 | 53.99 |
| SWLA (\(c=2\)) | 760M / 50B | 22.76 | 22.85 | 50.12 |
| SWLA + Proteus | 760M / 50B | 21.82 | 20.88 | 50.90 |
| Comba | 760M / 50B | 21.94 | 21.77 | 51.43 |
| Comba + Proteus | 760M / 50B | 21.32 | 21.20 | 52.15 |
| Titans | 760M / 50B | 20.92 | 21.28 | 52.65 |
| Titans + Proteus | 760M / 50B | 20.44 | 20.93 | 53.36 |
| Transformer++ | 1.3B / 100B | 17.04 | 17.42 | 55.81 |
| Hope-Attention | 1.3B / 100B | 15.66 | 13.39 | 56.93 |
| Hope-Attention + Proteus | 1.3B / 100B | 15.70 | 13.01 | 57.30 |
| SWLA (\(c=2\)) | 1.3B / 100B | 18.11 | 16.95 | 53.67 |
| SWLA + Proteus | 1.3B / 100B | 18.07 | 17.18 | 54.20 |
| Comba | 1.3B / 100B | 16.98 | 14.17 | 55.42 |
| Comba + Proteus | 1.3B / 100B | 16.95 | 14.01 | 56.00 |
| Titans | 1.3B / 100B | 15.36 | 13.18 | 56.95 |
| Titans + Proteus | 1.3B / 100B | 14.94 | 13.03 | 58.00 |
在长上下文大海捞针（Needle-In-A-Haystack, NIAH）评测中，针对数值检索（S-NIAH-2）与 UUID 复杂精确检索（S-NIAH-3）以及多键检索（MK-NIAH-1）在 16K 超长窗口（超出 8K 训练长度）下的性能对比如下：
| 评估任务 (16K 上下文) | Titans 基线 | Titans + Proteus | Comba 基线 | Comba + Proteus | Hope-Attn 基线 | Hope-Attn + Proteus |
|---|---|---|---|---|---|---|
| S-NIAH-2 (数字检索) | 69.4 | 74.2 (+4.8) | 13.4 | 21.2 (+7.8) | 96.8 | 97.2 (+0.4) |
| S-NIAH-3 (UUID 检索) | 21.4 | 29.8 (+8.4) | 7.4 | 10.8 (+3.4) | 50.0 | 54.4 (+4.4) |
| MK-NIAH-1 (多键检索) | 11.8 | 16.8 (+5.0) | 10.4 | 14.2 (+3.8) | 62.2 | 64.8 (+2.6) |
| LongBench (综合平均分) | 13.80 | 14.15 (+0.35) | 13.05 | 13.23 (+0.18) | 15.72 | 16.65 (+0.93) |
消融实验¶
分块数量 \(E\) 是 Proteus 控制有效容量调度精细度的核心超参数。当 \(E=1\) 时即退化为全序列暴露静态容量的基线模型。作者在保持其他所有设定严格一致的情况下，对 \(E \in \{1, 2, 4, 8, 16, 32\}\) 进行了详细消融实验，验证集困惑度变化如下：
| 分块数量配置 \(E\) | 早期每个区间占用容量比例 | 验证集困惑度 (Perplexity) 表现 | 核心观察说明 |
|---|---|---|---|
| \(E = 1\) (静态基线) | 100.0% (无门控) | 最高 (基线水平) | 早期容量无瓶颈，早期 token 过度拟合且后期覆盖严重 |
| \(E = 2\) | 50.0% | 显著优于 \(E=1\) | 粗粒度两阶段瓶颈即带来明显收益 |
| \(E = 4\) | 25.0% | 进一步下降 | 压缩压力适度增强，有效提升泛化性 |
| \(E = 8\) | 12.5% | 持续大幅下降 | 早期抽象提取与中后期容量供给趋向平衡 |
| \(E = 16\) (默认最优) | 6.25% | 最低 (最优性能拐点) | 形成最佳的前期强摘要与后期新生空间解耦 |
| \(E = 32\) | 3.125% | 出现反弹退化 | 初始有效容量仅剩 \(1/32\)，瓶颈过窄导致早期有效信息发生丢弃 |
关键发现¶
- 收益随序列长度单调递增：在短序列（如 4K 简单 passkey 检索）上，基线模型本已接近饱和，Proteus 表现持平；但在 16K 超长上下文及高难度 UUID 任务中，Proteus 的提升极其剧烈（如 Titans 16K S-NIAH-3 准确率从 21.4 暴涨至 29.8）。
- 非单调分块敏感性印证核心假设：消融实验清晰证明，分块数过少（\(E=1, 2\)）无法提供足够的压缩诱导；而分块过多（\(E=32\)）会导致最初分配的有效维度过小，损害了极早期语义的基本表达，验证了容量增长调度的权衡边界。
- 全序列位置困惑度普惠降低：按 token 位置分析困惑度表明，Proteus 在序列从头至尾每一个 token 处均取得了比基线更低的困惑度，并非通过牺牲早期质量来换取后期收益，其优势在 8K 训练边界处达到最大，并能平滑延续至 32K 外推区域。
亮点与洞察¶
- 概念视角的根本颠覆：传统研究往往聚焦于“如何更新记忆”（改用 Delta 规则、增加动量或引入非线性），Proteus 则指出“何时暴露记忆”同等关键。静态记忆暴露是长期被忽视的架构缺陷。
- 极简工程实现与零显存开销：通过预定义的前缀切片二值掩码实现渐进激活，不引入额外的可学习路由网络或前向分支，在训练与推理中不增加任何参数量和显存峰值。
- 嵌套学习视角的跨架构统合：成功将隐状态的容量调度推广至前馈层参数的优化轨迹，打通了循环状态模型与 Transformer 混合架构之间的范式阻隔。
局限与展望¶
- 手动预定调度表的局限：当前采用的是与绝对位置绑定的确定性线性分块，缺乏对不同文本信息密度与难度的自适应感知，未来可探索基于内容驱动的数据依赖动态激活路由。
- 训练外推边界的固定性：在超过预训练窗口 \(N\) 之后，模型记忆已完全解锁，无法继续以新生分块来缓解 32K 以上超长区域的干扰，需结合动态增长记忆（如 Memory Caching）进行长效拓展。
相关工作与启发¶
- vs Titans / Comba / TTT: 这类工作致力于设计更复杂的在线优化目标、自适应遗忘门或闭环控制，但状态矩阵均是全时全量暴露；Proteus 作为正交机制，可直接叠加于其上并进一步释放其长程记忆潜力。
- vs 动态计算路由 (MoD / MoR): Mixture-of-Depths 等方法在 token 维度上动态分配层数计算，而 Proteus 专注于在位置时序上对记忆容量与参数更新实施时间调度。
- vs 动态扩展记忆 (Memory Caching): Memory Caching 是在外部随长度增加显式分配新的物理存储，而 Proteus 是在固定常数大小的物理状态内部调度逻辑激活比例，二者互为补充。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次系统论证静态记忆容量的有害性，并提出时间维度上的渐进激活范式。
- 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 4 大前沿骨干模型、多尺度训练、语言建模、常识推理、NIAH 及 LongBench 完整评测。
- 写作质量: ⭐⭐⭐⭐⭐ 联想记忆数学形式化严谨，问题动机阐述极其深刻通透。
- 价值: ⭐⭐⭐⭐⭐ 零额外成本普适增强线性序列模型与长上下文建模，工程落地价值极高。