LooM:面向因果编码器-解码器语言模型的循环不变记忆
LooM: Loop-Invariant Memory for Causal Encoder-Decoder Language Models
Lele Yang · Muxi Diao · Zhonghao Yan · Jiankang Chen · Bin Wen · Tingting Gao · et al.
中文摘要
循环 Transformer 通过重复执行共享层来增加有效深度,但三个限制阻碍了其应用:内存随循环次数增长,计算只能按整块步长调整,循环次数在训练后难以更改。LooM(循环不变记忆)逐一解决这些问题。在因果编码器-解码器(CED)中,LooM 仅循环第一层之后的解码器层,这些层读取每个 token 写入一次的全局记忆,且每层仅在首次访问某个 token 时发布其键值(KV)条目,因此缓存不随循环次数增长。这些层独立循环,为每个 token 提供 3 到 21 次层执行的任意预算,一个在随机采样预算上训练的模型无需路由器即可覆盖所有预算。在足够循环次数下,LooM 在语言建模、零样本准确率和长上下文困惑度上优于在相同数据上训练的单次前向基线,且缓存大小相同。分析进一步表明,训练预算分布决定了额外深度带来何种收益。这些结果表明 CED 是循环的自然场所:记忆只写一次,深度成为部署时选择的预算。
关键要点
- 01循环 Transformer 存在内存随循环次数增长、计算调整粒度粗、训练后循环次数难改三个限制。
- 02LooM 仅循环解码器层,全局记忆每 token 写一次,KV 条目仅首次访问时发布,缓存不随循环次数增长。
- 03解码器层独立循环,支持每 token 3 到 21 次层执行,单模型在随机预算上训练即可覆盖全部预算。
- 04足够循环下,LooM 在语言建模、零样本准确率和长上下文困惑度上优于同数据单次前向基线,缓存大小相同。
- 05训练预算分布决定额外深度带来的收益类型,深度成为部署时可选的预算。
解读
尚无解读。
原始英文摘要
arXiv:2611.06835v1 Announce Type: new Abstract: Looped Transformers increase effective depth by re-executing shared layers, but three limitations hinder their adoption: memory grows with the loop count, compute can only be adjusted in whole-block steps, and the loop count is hard to change after training. We propose LooM (Loop-invariant Memory), which addresses each in turn. In a causal encoder-decoder (CED), LooM loops only the decoder layers after the first, which read a global memory written once per token, and each of them publishes the key-value (KV) entry for a token only on its first visit, so the cache does not grow with the loop count. These layers loop independently, providing every budget from 3 to 21 layer executions per token, and one model trained on randomly sampled budgets covers all of them without a router. With enough loops, LooM outperforms a single-pass baseline trained on the same data on language modeling, zero-shot accuracy and long-context perplexity, while using the same cache size. Our analysis further shows that the training budget distribution decides what extra depth buys. These results suggest that a CED is a natural place to loop: memory is written once, and depth becomes a budget chosen at deployment.