Mela：基于转化假说的测试时记忆巩固 Mela: Test-Time Memory Consolidation based on Transformation Hypothesis
用神经科学双频耦合机制构造层次记忆模块，把长上下文外推从崩溃曲线拉成平缓直线。
前置知识
Transformer 与自注意力
Transformer 通过 Q/K/V 投影计算 token 之间的加权聚合，其核心注意力机制具有 $O(n^2)$ 的时间与空间复杂度，$n$ 为序列长度，因此处理 32K 上下文需要的算力与显存是 4K 的 64 倍以上。
本文的出发点正是破解标准注意力的二次方瓶颈；理解 L-module/H-module 是如何替代或补充密集注意力，是看懂 Mela 设计取舍的前提。
Test-Time Training (TTT) 与神经记忆
TTT 把模型参数拆为内循环与外循环，内循环权重在测试时通过在线元学习持续更新；Titans 进一步把这种元学习模块解释为一种'神经长时记忆'，其状态 $M_t$ 沿序列通过梯度信号累积更新，例如 $M_t = \alpha M_{t-1} + \text{NewtonSchulz}_k(S_t)$。
Mela 的记忆模块正是 TTT 范式的直接继承者；HMM 的两个子模块本质上都是这种神经记忆，只是更新频率与容量不同。
跨频耦合与记忆转化假说
神经科学发现，记忆巩固涉及海马体（快更新、保留细节） 与皮层（慢更新、形成抽象语义） 之间的协调，并伴随不同频段神经振荡的同步现象。转化假说（Transformation Hypothesis）认为，记忆在被巩固的过程中会从情境丰富的情节记忆转变为去语境化的语义/图式表征。
这是 HMM 设计的理论来源：L-module 对应海马体，H-module 对应皮层，两者的频差则由 L 循环次数显式编码。
Hierarchical Reasoning Model (HRM) 与 TRM
HRM 用两个不同更新频率的循环模块 $f_L$、$f_H$ 来增强推理深度，并通过 1-step 梯度近似避免反向传播整段循环；TRM 则改用深度监督，将除最后一步外的循环步设为 no-grad，并在结构上做了简化。
Mela 的 HLR 借鉴 HRM 的双频循环，但放弃了固定点假设，并扩展到长程记忆建模；理解 HLR 与 HRM 的差异是把握本文技术贡献的关键。
Muon 优化器与 Newton-Schulz 正交化
Muon 在动量更新后对梯度做 Newton-Schulz 迭代（5 步即可近似半正交矩阵），目的是让权重更新方向更接近谱范数守恒，从而稳定深层训练。
Mela 把 Newton-Schulz 应用在记忆模块的梯度上，与传统 AdamW 形成对比；该选择直接影响 H-module 在长程循环中的数值稳定性。
研究动机
标准 Transformer 的 $O(n^2)$ 注意力复杂度使得上下文长度在训练阶段就被钉死：一旦推理时序列超过预训练窗口，困惑度会出现断崖式上升。在 4K 预训练窗口下，作者给出的基线数据极具说服力——Transformer++-400M 在 16K 时 PPL 已飙到 130.21，32K 时更是高达 303.56，而 1.2B 版本在 32K 时达到 597.37。已有的 TTT/Titans 路线把记忆编码进网络权重、用快速权重提供长程信息，但参数规模往往受限于单个记忆模块，更新频率单一，缺少对'情节细节'与'语义抽象'两类信息的显式分工。另一条路线（HRM/TRM）虽然引入双频循环，却以'假设循环已收敛到不动点'为前提来截断梯度，且主要用于小规模结构化推理任务，尚未被验证能处理 4K–32K 的开放域语言建模。
本文的目标是本文希望构造一种能在测试时持续进行记忆巩固的层次化神经记忆架构，使得一个 4K 预训练的语言模型能在 32K 推理时仍保持与 4K 相近的困惑度，且不依赖 RoPE 外推、YaRN 等长度扩展技术。具体来说，作者希望在 Transformer 解码器之外独立设置一个双模块记忆系统 HMM，使其中一个子模块以高频更新并保留情节细节，另一个子模块以低频更新并蒸馏语义 gist，再让二者的输出在每个解码层动态融合。
与已有工作不同的是，本文的独特切入角度是把神经科学里'海马–皮层系统巩固'与'跨频耦合'两个理论显式映射到模块设计：H-module（皮层）容量更大、循环更少，L-module（海马）容量较小、循环更多，二者通过 HLR 在不依赖固定点假设的前提下完成多步巩固。与 HRM 仅输出 $z_H$ 不同，Mela 强制在最终输出中保留 $h$ 与 $l$ 两路信息，并通过 MemStack 把循环过程中每一轮的中间状态也注入解码器早期层，相当于让模型不仅'看见巩固结果'，还能'看见巩固轨迹'。这套设计在公开论文中尚未见到先例。
核心方法
Mela 的整体思路可以拆成'先做记忆、再做语言'两个独立阶段。输入文本先经嵌入与 1-D 卷积分块（patch size $P=32$），得到 $N_p$ 个 patch 作为情节记忆单元；这些 patch 同时送入 H-module（高语义、低频）与 L-module（高细节、高频）两个神经记忆子模块。两者通过 HLR 协同工作：在 $H$ 个高层循环中，每个高层循环内部对 L-module 做 $L$ 次精细化，再用 H-module 蒸馏一次 $h$ 状态。所有循环结束后，$h$ 与 $l$ 状态经两层 MLP 融合，作为'记忆 token'拼接到每个 chunk 的文本 token 之前，输入到标准的 Transformer 解码器。
最核心的创新是把'记忆的层次化'从抽象比喻变成可计算的结构：用 Newton-Schulz 正交化后的梯度 $S_t$ 更新记忆参数 $M_t = \alpha M_{t-1} + \text{NewtonSchulz}_k(S_t)$（$k=5$），并让 H-module 与 L-module 的更新频率通过循环次数 $L$、$H$ 显式区分；最终输出不是单一 $h$ 状态，而是 $h$ 与 $l$ 的 MLP 融合，从而在一次前向中同时提供语义 gist 与情节细节。MemStack 进一步把每一轮循环的中间 $(l^{(i)}, h^{(i)})$ 状态沿解码器前 $H-1$ 层逐步叠加，让解码器获得'巩固过程本身'的信息，这是与 Titans/TTT 最本质的区别。
方法步骤详情
完整流程分为四步。第一步是输入切块：原始 token 嵌入为 $x \in \mathbb{R}^{N \times d}$，再用大小 $P$ 的 1-D 卷积切为 $N_p = N/P$ 个 patch，每个 patch 是一段情节记忆单元，$l$ 与 $h$ 状态均初始化为 $x_p$。第二步是双模块记忆：每个记忆块由 RMSNorm + 门控注意力 + 神经记忆模块组成；神经记忆以 chunk size $C=64$ 计算键值关联损失 $\ell(M_{t-1}, x_t) = \|N(k_t; M_{t-1}) - v_t\|_2^2$，并用 Newton-Schulz 正交化后的动量更新 $M_t$；H-module（8 层）负责抽象，参数多、循环少，L-module（4 或 8 层）负责细节，参数少、循环多。第三步是 HLR 协同：在 $H$ 个高层循环中每轮做 $L$ 次 L-module 细化再 1 次 H-module 蒸馏，前 $H-1$ 轮在 torch.no_grad 下进行，最后一轮开启梯度。第四步是解码与 MemStack：把 $h^{(i)}$、$l^{(i)}$ 融合为 $z^{(i)}$ 后，把 $z^{(H)}$ 作为 token 拼接到每 chunk 文本前送入解码器，并把 $z^{(1)}, \ldots, z^{(H-1)}$ 通过残差形式注入解码器前 $H-1$ 层，即 $y^{(\alpha)} = B^{(\alpha)}(y^{(\alpha-1)}) + z^{(\alpha)}$。
技术新颖性
技术新颖性体现在三个层次。其一是结构层：把神经科学的'快/慢两套巩固系统'显式实例化为可微分模块，且让 H-module 容量大于 L-module，这与传统 RNN/LSTM 共享参数的设计有本质区别。其二是算法层：HLR 借鉴 HRM 的双频循环，但放弃了固定点收敛假设，改用深度监督（仅最后一步反传），并允许最终输出融合 $h$ 与 $l$，因此既避免了 1-step 近似的精度损失，又规避了 TRM 提到的'扩层导致过拟合'的可扩展性陷阱。其三是信息流层：MemStack 把循环中间状态沿解码器深度方向堆叠，在不增加额外 token 的前提下让解码器'看'到记忆巩固的整条轨迹，这比 Titans 的单点注入更接近神经科学中'记忆痕迹分布在不同巩固阶段'的描述。
实验结果
作者在 FineWeb-Edu 的 5B token 子集、4K 上下文预训练下比较 Mela 与 Transformer++ 基线，最关键的一组数据来自 Table 5：在 1.2B 规模上，Mela 在 4K 时 PPL 为 9.20、8K 为 9.91、16K 为 11.64、32K 为 11.85；而 Transformer++ 在 8K 已经退化到 12.71，32K 达到 597.37，相差近 50 倍。400M 规模上 Mela 在 32K 时 PPL 仅 14.50，而 Transformer++ 已经是 303.56。Figure 4 直观展示了这条差异：Mela 三条曲线在 4K 之后几乎水平，Transformer++ 三条曲线则呈指数上扬。Table 2 的循环深度消融表明 H=4、L=4 是最佳配置，并且 H 循环带来的收益随上下文增大而扩大（如 H=2 与 H=4 的差距从 1K 的 0.06 扩到 16K 的 0.18），印证 H-module 是长程巩固的关键。Table 3 的容量消融呈现镜像结果：H-module 层数从 2 增到 8 在 8K 时获得 1.19 的 PPL 下降，而 L-module 同等扩容在 1K 时获得最大 0.49 的下降、16K 时只剩 0.19——这与神经科学中'情节编码决定短程、语义巩固决定长程'的判断高度一致。Table 4 进一步指出门控注意力贡献最大（去掉后 +1.24 PPL），Newton-Schulz 正交化去掉后 +0.21 PPL、MemStack 去掉后 +0.10 PPL，验证了'选择性写入''稳定更新''轨迹注入'三个机制缺一不可。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 语言建模（4K 预训练窗口内的标准 PPL） | Perplexity（4K 上下文） | Mela-400M=12.01, Mela-800M=10.04, Mela-1.2B=9.20 | Transformer++-400M=12.56, Transformer++-800M=11.35, Transformer++-1.2B=9.53 | 在预训练窗口内 Mela 三个规模分别比基线低 0.55、1.31、0.33 PPL |
| 长度外推（8K 上下文 PPL） | Perplexity（8K） | Mela-400M=12.64, Mela-800M=10.82, Mela-1.2B=9.91 | Transformer++-400M=28.26, Transformer++-800M=16.69, Transformer++-1.2B=12.71 | 8K 已是预训练窗口 2 倍，Mela 在三个规模上分别领先 15.62、5.87、2.80 PPL |
| 长度外推（16K 上下文 PPL） | Perplexity（16K） | Mela-400M=14.43, Mela-800M=12.74, Mela-1.2B=11.64 | Transformer++-400M=130.21, Transformer++-800M=129.10, Transformer++-1.2B=104.14 | 16K 时差距已拉大到约一个数量级 |
| 长度外推（32K 上下文 PPL） | Perplexity（32K） | Mela-400M=14.50, Mela-800M=13.41, Mela-1.2B=11.85 | Transformer++-400M=303.56, Transformer++-800M=497.85, Transformer++-1.2B=597.37 | 32K 时 Mela 相对基线下降约 21×–50×，且 PPL 几乎与 4K 一致 |
局限与改进
作者没有在论文中显式罗列 limitation，但可以观察到四点约束。其一，训练数据只用了 FineWeb-Edu 的 5B token，量级远小于主流 LLM 的数千亿 token 训练集，因此 1.2B 的 Mela-1.2B 在 4K 的 PPL 仍比同规模完整 LLaMA 高出不少，外推能力的强弱能否在更大数据规模下保持还未知。其二，消融实验是在 1B token 的小子集上做的，Table 2–4 的结论是否对 5B 或更大训练量仍然成立，需要补充实验。其三，论文完全没有报告推理时的显存与吞吐数据——Mela 需要在每个 chunk 维护额外的记忆 token（Table 1 显示 400M 时 memory tokens 占 27.45%、文本占 65.75%），实际部署成本可能高于同参数 Transformer。其四，HLR 的 $H$、$L$ 都被设为 4，缺少对更大循环步数的可扩展性验证，作者也没有讨论 $H \times L$ 步循环对长尾延迟的影响。
独立分析的弱点
独立审视论文时，有四个值得改进的方向。第一，记忆模块的更新完全依赖 Newton-Schulz 正交化后的关联损失，对训练数据的'惊喜度'分布高度敏感；如果长文本中存在大量重复或低信息密度片段，记忆权重可能反复在小特征值上做无意义更新，建议加入显式的'无操作门'或休眠机制。第二，MemStack 把中间状态注入前 $H-1$ 层时仅做残差相加，没有可学习的融合权重；Table 4 显示 Fusion 实现中 MLP 优于 Weighted Sum，说明对不同循环层注入强度应做自适应调节，而不是一视同仁地相加。第三，本文只在 FineWeb-Edu 上做语言建模，没有验证 Mela 在多轮对话、检索增强或长文档问答等需要精确回溯上下文的场景下是否依然鲁棒——H-module 是'语义抽象'，理论上会丢弃低频事实，这种设计在开放问答中可能成为负担。第四，HLR 训练时只在最后一步开启梯度，前 $H-1$ 步相当于多消耗显存但不更新参数，未来可以探索用 REINFORCE 类方法或隐式微分把梯度传遍整段循环。
未来方向
作者在结尾提到把'长程规划与更开放的能力'作为神经记忆的下一步，读者可以沿三条线延伸。其一是把 HMM 接到多模态/长视频上：L-module 的高频更新适合保留逐帧细节，H-module 的低频更新适合形成事件级抽象，正好对应 video QA 中的'何时发生'与'是什么'两类问题。其二是用 HMM 替代 RAG 中的检索压缩器：与其把检索结果原样塞进 prompt，不如让 HMM 在线巩固检索证据，使 LLM 看到的不是 top-k 原文而是经过跨频巩固的语义锚点。其三是与 MoE 结合，让 L-module/H-module 中的多个'海马–皮层对'分别负责不同的知识簇，进一步逼近大脑中'多重记忆系统'的分工。
复现评估
复现门槛中等偏低。作者在摘要中明确给出代码仓库 https://github.com/Musubi-ai/Mela，主干模块在论文 Figure 1/2/3 中给出了 PyTorch 风格伪代码（包括 hierarchical_latent_recursion 与 MemStack 的数学表达），超参数在 Table 1 全部公开。训练侧采用公开的 FineWeb-Edu 与 Llama 2 tokenizer，优化器为 AdamW（峰值 LR $10^{-3}$/$8 \times 10^{-4}$，β=(0.9, 0.95)）+ FSDP2，batch size 0.5M tokens，使用 4K 上下文。最大规模 1.2B 在 5B token 上训练在 8×H100 级别算力上数日即可完成，400M 版本在单节点多卡上即可复现 Table 2–4 的消融。潜在难点在于：(a) Newton-Schulz 5 步迭代与 Muon 风格的实现细节需要谨慎对齐；(b) HLR 的 no-grad / grad 切换位置容易写错；(c) MemStack 的中间融合权重初始化对最终收敛影响较大，需要参考作者发布的检查点。
论文图表