ReWY speeds up delta-rule attention gradients
Original title: ReWY: Faster Gradient Computation for Delta-Rule Attention Family
ReWY reformulates backward passes for delta-rule linear attention reducing seven local WY matrix products to three and achieving 1.15-1.81x speedup on GPUs.
Chunkwise implementations of delta-rule linear attention expose matrix parallelism but compute unnecessary auxiliary adjoints in backward passes. ReWY presents algebraic reformulation that reduces this work by differentiating chunk correction as triangular residual system. Complementary transformations expose further reuse via centered matrix contractions and scalar-gate reductions. Implementations on GPUs and Ascend NPUs show 1.15-1.81x geometric-mean backward speedups while preserving original architecture.
- Reduces seven local WY matrix products to three
- Eliminates inverse adjoint for chunk correction
- Achieves 1.15-1.81x speedups on NVIDIA GPUs
- Retains original FP32 working matrices
- Applies to DeltaNet, GDN, KDA and GDN2
From the original article
1 min read at arxiv.orgDelta-rule linear attention underlies language models such as Kimi Linear and Qwen3-Next. Chunkwise implementations expose matrix parallelism, but their backward passes compute unnecessary auxiliary adjoints. We present ReWY, an algebraic reformulation that reduces this work before optimizing accelerator execution.
By Yi Zhang · arxiv.org
A short excerpt; the article and its images belong to the publisher. Summary, key facts and assessment above are AI-written from the article and can be wrong.