yifanzhang-pro/recurrent-looped-tranformer

解决的问题

Recurrent Looped Transformer (RLT) 解决了标准 Transformer 在长度泛化方面的限制，即在训练期间未见过的更长序列上执行任务的能力。它特别针对算法任务（如奇偶校验和状态追踪），这些任务需要计算路径随序列长度增长以保持准确性。

运作原理

RLT 修改了 Transformer 架构以引入跨 token 的循环：

• 循环反馈 (Recurrent Feedback)： 前一个 token 的最终解码器状态通过门控合并 (gated merge) 反馈到当前 token 的解码器输入中。这创造了一条随序列长度线性增长的计算路径。
• 混合记忆 (Hybrid Memory)： 它使用并行因果编码器进行全局键值 (KV) 记忆，并使用带有滑动窗口注意力 (SWA) 的解码器进行局部记忆。
• 状态连续性 (State Continuity)： 循环状态在提示词与响应之间的边界处无缝延续。
• 反馈变体 (Feedback Variants)： 该项目探索了不同的反馈间隔 ($B$)：RLT-1 每个 token 更新一次，RLT-2 每 $B$ 个 token 的区块更新一次，而 RLT-0 则没有反馈。

适用对象

致力于序列建模、长度泛化以及 LLM 算法推理的研究人员与开发者，他们希望突破标准 Transformer 固定深度的限制。

亮点

• 强大的长度泛化： 在显著超过训练数据长度的情况下（例如奇偶校验为 256 位元），在奇偶校验和 S5 交换任务上达到 100% 的准确度。
• 循环路径： 计算路径对于 $t$ 个 token 遍历 $tL_D$ 个解码器区块，同时保持固定的每个 token 成本。
• 灵活的深度分配： 实验展示了如何在编码器和解码器之间分配层数，从而影响不同算法任务的性能。
• 效率权衡： 与 RLT-1 相比，RLT-2（分块反馈）在训练中提供了加速（高达 2.27 倍），同时在某些任务上保持了高准确度。