论文

RoLA:旋转定位低秩线性注意力,实现高效扩散 Transformer

RoLA: Rotary-Positioned Low-Rank Linear Attention for Efficient Diffusion Transformers

模型架构Transformer

摘要

Diffusion Transformer (DiTs) 实现了强大的视频生成质量,但其密集的时空自注意力随序列长度呈二次方缩放,并很快成为主要的推理瓶颈。稀疏低秩混合通过将局部稀疏分支与全局压缩分支相结合来减轻这种成本。在配备 3D 旋转位置嵌入 (RoPE) 的视频 DiT 中,全局分支面临结构兼容性问题:当在非线性特征图之前应用 RoPE 时,旋转和非线性通常不会交换,因此很难在保留相对旋转几何的同时保持独立于查询的线性摘要。现有的工作通常通过用坐标条件代理或可学习的绝对位置模块替换真正的跨词元全局聚合来回避这个问题。这些妥协可能是有效的,但它们近似于绝对坐标的相对衰减并引入了额外的位置参数。我们提出 \textbf{RoLA},一个旋转定位的低秩线性注意力分支,它保持真正的跨词元聚合,同时保持与可重用的线性摘要兼容。该设计将 RoPE \emph{outside} 应用于非线性低秩特征图,并重用与低秩瓶颈匹配的预训练旋转调度的截断子集。这产生了一个线性时间低秩全局分支,通过设计具有相对位置行为,并且没有额外的位置参数;完整的稀疏-低秩模块仍然包括固定稀疏稀疏分支。对开源视频 DiT 的实验表明,所得方法在稀疏度为 90% 时在生成质量方面保持竞争力,同时在 Wan2.1-14B(720p、81 帧,在 NVIDIA H100 GPU 上测量)上实现 2.63$\times$ 端到端推理加速。