Token-Shift 是 RWKV 模型中为解决并行训练效率而引入的关键机制。由于 RWKV 的结构天然具有时间递归特性,严格按时间步依赖前一个 token 的信息,这使得其在推理时效率极高,但也造成训练过程难以并行。为此,Token-Shift 技巧通过将序列输入向右平移一位(即“看到”前一个 token),并对当前 token 和上一个 token 的特征进行加权混合,模拟出时间依赖关系。这个过程通过广播和矩阵操作实现,使得模型可以像 Transformer 一样高效并行训练,又不会打破因果结构。Token-Shift 的引入使 RWKV 在保持低延迟推理优势的同时,兼顾了训练阶段的速度与效率,是其架构设计中的一项关键创新。 | Token-Shift 是 RWKV 模型中为解决并行训练效率而引入的关键机制。由于 RWKV 的结构天然具有时间递归特性,严格按时间步依赖前一个 token 的信息,这使得其在推理时效率极高,但也造成训练过程难以并行。为此,Token-Shift 技巧通过将序列输入向右平移一位(即“看到”前一个 token),并对当前 token 和上一个 token 的特征进行加权混合,模拟出时间依赖关系。这个过程通过广播和矩阵操作实现,使得模型可以像 Transformer 一样高效并行训练,又不会打破因果结构。Token-Shift 的引入使 RWKV 在保持低延迟推理优势的同时,兼顾了训练阶段的速度与效率,是其架构设计中的一项关键创新。 |