MeSH: Memory-as-State-Highways for Recursive Transformers
Chengting Yu, Xiaobo Shu, Yadao Wang, Yizhen Zhang, Haoyi Wu, Jiaang Li, Rujiao Long, Ziheng Chen, Yuchi Xu, Wenbo Su, Bo Zheng
ICLR 2026
cs.LG, cs.AI
2025-10-09
阿里提出 MeSH,给循环 Transformer 外挂记忆槽和逐步路由。Pythia-1.4B 上用少 33% 非嵌入参数,把 0-shot 平均准确率相对原版抬高 1.06 个百分点。
循环 Transformer 把同一组层反复跑,计算深度和参数深度就此分开。同等算力下,参数更少的循环模型经常打不过层参数各不相同的原版,困惑度更高、下游更差。循环看起来省参数,却把性能也省掉了。
阿里这篇用隐藏状态探针,把差距钉在两个结构问题。一是无差别计算:共享 core 不知道自己在第几圈,第一步几乎包办全部更新,后面几圈更新幅度掉到接近零,相邻圈的表示 CKA 相似度极高,等于卡在不动点。二是信息过载:长期要记住的输入语义、和每一步临时要用的工作记忆,挤在同一个 hidden state 里,奇异值谱迅速塌缩,有效秩变低。
骨架沿用 Prelude-Recurrent-Coda。Prelude 先把 token embedding 做成进入循环的初始状态,共享 core 循环 K 次,Coda 再接到最终表示。朴素循环每步只做 h ← fcore(h)。常见补丁是固定加法:residual 每步加回上一状态,anchor 每步加回循环入口状态,anchor 加回原始 embedding。这些 skip 能减轻遗忘,但 core 对圈次仍然盲目。
MeSH 把状态管理拆出去。模型维护 B 个与 hidden state 同形状的记忆槽,槽 0 初始化为 token embedding,其余为零。每一步有独立参数的 Write/Read 路由器,各是一层线性映射加 softmax,按 token 给出对 B 个槽的权重。Core 先算出本圈输出,按写权重软写入各槽,再按读权重从更新后的缓冲合成下一圈输入。Prelude 出口先进一次过渡读写,循环结束后再读一次交给 Coda。
缓冲长度用经验规则 B = Nloop + 3。以 4+8R2+4 为例,Nloop=2,B=5,对应每个主状态一个槽再加两个草稿槽。路由器参数量在 Pythia-1.4B 上只有 61470,相对非嵌入参数约 0.005%。
训练设定对齐 Pythia:在去重 Pile 的 2500 亿 token 上从零预训练一轮,序列长度 4096,batch 512。下游是 9 个任务的 0-shot/5-shot 平均准确率。
| 规模 | 变体 | 非嵌入参数 | 0-shot | 5-shot | Pile PPL |
| 1.4B | Vanilla 24 层 | 100% | 49.50 | 51.93 | 7.44 |
| 1.4B | 循环 base | -33.3% | 48.89 | 50.99 | 7.63 |
| 1.4B | +anchor | -33.3% | 49.39 | 51.27 | 7.51 |
| 1.4B | +MeSH | -33.3% | 50.56 | 52.79 | 7.39 |
1.4B 上 MeSH 相对 Vanilla 的 0-shot 平均准确率 +1.06,5-shot +0.86,同时 Pile 困惑度从 7.44 降到 7.39。更小规模上,160M 的 5-shot 刚过线(+0.06),410M 在 50% 压缩配置下 5-shot +0.73。2.8B 和 6.9B 用 6+10R2+6(非嵌入参数约 68.75%),没有 Vanilla 对照,MeSH 相对循环 base 的 0-shot 分别 +2.22 和 +2.16。
410M 上,把 residual、anchor、anchor 做固定或可学习线性组合,最好的手调组合 Pile PPL 是 9.17,MeSH 是 9.09。缓冲长度消融在 B=5 最好,再加槽略差。把 MeSH 接到不共享权重的 1.4B Vanilla(按 4+8+8+4 切开),Pile PPL 从 7.44 降到 7.26。OLMoE 2.6B 总参、512M 激活的循环版上,MeSH 优于循环 base,仍略低于 Vanilla。
探针侧,MeSH 三圈的更新幅度更均衡,圈间 CKA 下降,奇异值衰减变慢。
循环 Transformer 一直卡在「同等算力打不过原版」。MeSH 的判断是:问题不在共享权重本身,而在单通道状态过载、且无法按圈分工。把记忆外置、把圈次差异交给逐步路由,共享 core 可以继续纯共享。对在做 Universal Transformer、Huginn 这类循环深度的人,这是一个几乎零参数开销的插件。1.4B 上已经反超更大的非循环模型,2.8B/6.9B 相对循环基线也有稳定增益。
这仍是从零预训练的架构实验,不能直接接到现成 Llama 权重上就用。
作者自己写了两点。实验停在 Pythia-6.9B 和去重 Pile,更大模型和别的训练配方还没验证。非循环场景只做了 1.4B 的切块实验,没有系统扫层数和槽数。
6.9B 没有 Vanilla 对照,所以不能说大尺度也一定反超原版。路由按圈独立参数,循环次数是训练时定死的,没有 Adaptive Computation Time 那种按样本停圈。记忆槽与 hidden state 同形状,序列一长,缓冲的显存按槽数线性涨。OLMoE 上 MeSH 没超过 Vanilla,说明 MoE 骨干上的红利更薄。