ReWY重构DeltaNet反向传播：WY矩阵乘法减至3次，GPU/NPU加速1.8倍

Delta-rule线性注意力（Delta-rule linear attention）是Kimi Linear、Qwen3-Next等语言模型的底层机制。这类模型在训练时，分块（chunkwise）实现虽能释放矩阵并行性，但反向传播却要计算大量不必要的辅助伴随（adjoint），成为效率瓶颈。近期一篇工作提出ReWY，通过代数重构而非硬件优化来削减计算量，在NVIDIA GPU和昇腾NPU上均取得显著加速。

ReWY: Faster Gradient Computation for Delta-Rule Attention Family

反向传播的冗余从何而来

Delta-rule线性注意力的分块实现将序列切分为多个块，块内做校正（chunk correction），块间传递隐状态。这一过程天然适合矩阵并行，但反向传播时，为计算梯度需要引入额外的伴随矩阵。问题在于，这些伴随矩阵中有部分并非必需——它们只是中间推导的副产品，却让计算图变得臃肿。

据论文描述，分块校正的微分可以表示为一个三角残差系统（triangular residual system）。把校正当作整体来微分，就能消去逆伴随（inverse adjoint），同时把WY矩阵相关的七个局部矩阵乘法压缩到三个。这是ReWY的核心代数观察：先做公式层面的化简，再谈加速器执行。

伴随（adjoint）与反向传播：在自动微分中，伴随量指目标函数对中间变量的梯度。对于含矩阵逆或递推结构的计算图，朴素反向传播需要为每个中间矩阵单独存储并传播梯度，产生大量辅助伴随（auxiliary adjoint）。这些伴随并非原始模型参数梯度所必需，而是计算图展开的副产物。它们不仅占用显存，还引入额外的矩阵乘法。Delta-rule线性注意力的分块反向传播中，逆伴随（inverse adjoint）就是典型例子：它来自对块内校正矩阵求逆这一操作的微分，但实际上可以通过对整体残差系统求微分来避免显式构造。理解伴随的冗余来源，是理解ReWY为何能通过代数化简加速的关键。

三项变换如何复用计算

除了消去逆伴随，ReWY还利用两个互补的变换进一步暴露可复用性。其一是用中心化矩阵收缩（centered matrix contractions）替代逐对向量门控处理；其二是标量门控归约（scalar-gate reductions）不再构造键的Gram矩阵，而是复用键梯度乘积。

这些恒等式在精确算术下保持递推关系和梯度不变，不改变模型架构，也不改变训练目标。换言之，ReWY是一种“纯计算图”优化，对DeltaNet、GDN、KDA和GDN2四种算子均适用，模型行为不受影响。

WY矩阵与Delta-rule的代数结构：Delta-rule线性注意力可视为一种带校正的快速权重更新规则，其分块形式常借助WY矩阵表示——WY矩阵由键、值、门控向量通过一系列外积和逐元素操作构造，用于将块内递推写成矩阵乘法形式。WY矩阵相关的计算图包含多个局部矩阵乘法（如键与值的交互项），朴素实现中数量较多。ReWY观察到这些局部乘法中有可合并或可消去的部分，通过中心化矩阵收缩（centered matrix contractions）和标量门控归约（scalar-gate reductions）等恒等式，将七个局部矩阵乘法压缩到三个，并避免构造键的Gram矩阵。这种化简不改变WY矩阵的数学值，只改变计算路径，因此对DeltaNet、GDN、KDA等变体通用。

实测：GPU与NPU上的加速幅度

研究者在GPU和昇腾NPU上实现了化简后的计算图。在匹配的重计算与全保存策略下，与冻结的FLA参考实现做六种形状的精度合格对比，NVIDIA GPU上完整低层反向传播的几何平均加速为1.15–1.81倍。具体到H100和H200，四种算子的加速在1.18–1.70倍之间，且保留原始FP32工作矩阵。

昇腾950DT上的独立公共autograd评测（包含梯度转换）在可用匹配策略下取得1.09–2.89倍加速。配对消融实验进一步确认，在优化互补收缩之后ReWY仍有额外收益。完整梯度检查验证了所记录的算子用例，补充的全模型反向传播测量则评估了集成到社区FLA模型以及GDN2层移植的效果。

重计算与全保存策略：在训练大型序列模型时，反向传播需要前向传播的中间激活。为节省显存，常用重计算（recomputation/checkpointing）策略：不保存全部中间结果，而是在反向传播时重新计算部分前向。全保存（full storage）则保留所有激活以换取速度。两种策略下反向传播的计算图不同，ReWY的化简收益也会变化。论文在匹配策略下与FLA参考实现对比，确保加速来自代数化简而非策略差异。几何平均加速比（geometric mean）是对多种形状的加速比取几何平均，能更稳健地反映整体提升，避免个别极端值拉偏算术平均。

为什么值得关注

DeltaNet类模型的训练成本中，反向传播占大头。ReWY的价值在于它不依赖新的硬件特性或近似计算，而是从数学结构上砍掉冗余矩阵乘法。对于已经在使用Kimi Linear或Qwen3-Next架构的团队，这类改动可以较低风险地集成，代码与复现材料已在GitHub开源（flashserve/ReWY）。

当然，1.15–1.81倍的加速是低层反向传播的几何平均，全模型端到端收益会因模型规模和训练配置而异，论文的补充测量也仅覆盖部分集成场景。但方向明确：Delta-rule注意力的效率优化，正从“堆硬件”转向“改公式”。