ReHyAt: 面向视频扩散Transformer的循环混合注意力机制 ReHyAt: Recurrent Hybrid Attention for Video Diffusion Transformers
ReHyAt通过混合注意力将视频生成复杂度从二次降为线性。
前置知识
Diffusion Transformer (DiT)
DiT是一种结合了扩散模型（Diffusion Model）和Transformer架构的生成模型。与早期的U-Net结构相比，DiT将视频视为时空图块（Spatiotemporal Patches）序列，利用Transformer强大的全局建模能力，能够更好地捕捉长时域的时间结构和高分辨率下的细节，已成为当前Sora、CogVideoX等先进视频生成系统的主流骨干网络。
本文的研究对象正是DiT架构。理解DiT是理解视频生成模型如何从U-Net演进到Transformer，以及为何现有模型面临算力瓶颈的基础。
Softmax Attention (Softmax注意力)
标准的自注意力机制，通过计算查询（Query）和键（Key）的点积并经过Softmax归一化来加权值（Value）。它使用指数核（$e^{qk^T}$），具有极强的表达能力，非常适合捕捉细粒度的依赖关系（如局部纹理、动作细节）。但其计算复杂度和内存占用随序列长度N呈二次方增长（$O(N^2)$），在处理长视频时成为性能瓶颈。
它是现有SOTA模型（如Wan2.1）的标准配置，提供了最高的生成保真度，但也带来了高昂的算力成本，是本文想要解决的核心矛盾之一。
Linear Attention (线性注意力)
线性注意力通过核技巧（Kernel Trick），将Softmax中的指数核替换为非负映射函数（如$1 + \text{elu}(x)$），从而利用矩阵乘法结合律将计算顺序从$(QK)V$改为$Q(KV)$。这使得计算复杂度降低到线性$O(N)$，并且可以改写为循环神经网络（RNN），实现常数级的内存占用。但其表达能力较弱，常导致生成质量下降或细节丢失。
它是本文提升效率的关键技术路径。文中通过改进核函数来弥补线性注意力在表达能力上的不足。
Knowledge Distillation (知识蒸馏)
一种模型压缩技术，通过让一个小的“学生模型”模仿一个大的“教师模型”的输出或中间特征，从而在保持性能的同时降低模型复杂度或计算量。在本文中，它特指利用强大的双向Softmax教师模型来指导轻量级的ReHyAt学生模型学习。
本文并没有从头训练一个线性模型（这通常代价极高），而是提出了高效的两阶段蒸馏方案，这是其能用极低算力（160 GPU hours）复现SOTA性能的核心原因。
研究动机
当前最先进的视频扩散模型（如Wan2.1, CogVideoX）普遍采用DiT架构，并使用双向全Softmax注意力机制。这种机制虽然保证了生成视频的高保真度，但其$O(N^2)$的计算复杂度和巨大的内存占用严重限制了模型的可扩展性。具体而言，当视频时长从5秒增加到10秒以上时，注意力模块的计算量和显存需求呈二次方爆炸式增长，导致在高端GPU上难以生成长视频，在手机等边缘设备上更是直接出现OOM（内存溢出）错误。此外，现有的线性注意力替代方案虽然降低了复杂度，但往往需要从零开始重新训练（如SANA-Video需数万GPU小时），且生成质量与Softmax模型存在明显差距，存在表达能力不足的问题。
本文的目标是本文旨在提出一种兼具Softmax注意力的高保真度和线性注意力的高效率的新型注意力机制，名为ReHyAt（Recurrent Hybrid Attention）。核心目标是将视频生成中注意力机制的计算复杂度从二次方$O(N^2)$降低到线性$O(N)$，并实现常数级的峰值内存占用，从而彻底解决长视频生成和移动端设备部署的算力与内存瓶颈。同时，为了降低研发成本，目标是以极低的训练代价（几百GPU小时量级）将现有的SOTA模型转化为这种高效形式。
与已有工作不同的是，已有的线性注意力或混合注意力方法存在两个关键盲点：一是许多方法（如Uniform Hybrid）虽然提升了质量，但仍保持二次方复杂度，无法改写为RNN，无法真正解决长序列的扩展性问题；二是它们通常采用“从头训练”或“重度微调”的模式，忽略了现有海量双向Softmax模型已经积累的庞大知识。本文的独特切入点在于“蒸馏”与“结构设计”的结合：通过设计一种能够自然支持RNN改写的分块混合注意力结构，并利用精心设计的轻量级多项式核映射来近似Softmax的指数核，从而在保留原模型架构的基础上，以极低的代价实现向高效模型的转化。
核心方法
ReHyAt的核心直觉是“术业有专攻”：对于视频中的局部、细粒度依赖（如相邻帧的纹理、物体边界），必须使用高精度的Softmax注意力；而对于长距离、全局的语义依赖（如背景变化、长程动作），则可以使用高效的线性注意力。基于此，ReHyAt将输入序列在时间维度上划分为重叠的块（Chunk）。在每个块内部，利用Softmax注意力以保证细节质量；在块与块之间（即全局范围），利用线性注意力以维持线性复杂度。技术路线上，作者将这种混合机制设计为因果（Causal）的，从而允许将其等价地改写为一种分块循环神经网络（Chunk-wise RNN），在推理时逐块生成，使得峰值内存保持恒定。
本文最核心的创新在于其“分块循环混合注意力”设计及其配套的“多项式特征映射”。与以往将所有Token同等对待或简单并行线性/Softmax不同，ReHyAt提出了非均匀的时域划分策略：将当前及附近的Token（当前Chunk）交给Softmax处理，将过去的所有Token交给线性处理。为了弥补线性注意力在表达Softmax指数核时的不足，作者没有使用固定的ELU函数，而是设计了可学习的多项式特征映射。该映射将输入经过轻量级网络后，进行不同阶数的幂次展开（如1次、2次），从而极大地扩展了特征的动态范围，使其能更准确地逼近Softmax的指数特性，保证了生成质量几乎不损失。
方法步骤详情
ReHyAt的方法流程主要分为两个阶段： 1. **注意力蒸馏（Attention Distillation）**：首先，将原始的SOTA双向DiT模型作为教师模型。对于学生模型中的每一个DiT块，仅保留其骨干网络参数并冻结，然后引入两个轻量级的可学习网络（phi_q和phi_k）来处理Query和Key，以计算线性注意力。训练时，不使用视频对，而是随机采样噪声和文本，让学生块的输出去拟合教师块的输出。通过这种方式，phi网络学会了如何用线性部分去捕捉教师模型中长距离的依赖关系。 2. **轻量级微调（Lightweight Fine-tuning）**：蒸馏完成后的模型虽然保持了场景结构，但在块与块的交界处可能存在不连贯（Episodic Incoherence）。因此，第二阶段使用少量的视频-文本对（如22K个样本），在全模型上进行短时间的微调（约1000步）。这一步利用标准的Flow Matching目标函数，对整个DiT模型（包括新增的phi参数和原有的FFN层）进行端到端优化，以平滑过渡并恢复生成细节。
技术新颖性
ReHyAt的技术新颖性主要体现在三个方面：第一，它是首个能够将现有的双向Softmax视频扩散模型转化为常数内存的循环（RNN）形式的混合注意力机制，解决了线性化过程中必须进行“暴力重训练”的痛点。第二，其重叠分块（Overlapping Chunks）设计巧妙地缓解了非连续分块带来的运动不连贯问题，而因果性设计保证了RNN改写的可行性。第三，提出的可学习多项式特征映射机制，相比传统的固定核函数，显著提升了线性注意力对指数核的逼近能力，在不牺牲效率的前提下大幅提升了生成质量。
实验结果
实验结果表明，ReHyAt在极低的训练成本下实现了与原始SOTA模型相媲美的性能。在VBench基准测试中，基于Wan2.1 1.3B蒸馏的ReHyAt模型（20x ReHyAt blocks）取得了83.79的总分，甚至略高于原始Wan2.1模型（83.31）。在更具挑战性的VBench-2.0测试中，ReHyAt也保持了高度的竞争力。更重要的是，人类主观评估（Human Preference Study）显示，在500对视频的对比中，受试者在总体一致性、场景、颜色等多个维度上对ReHyAt和原模型无显著偏好差异（总体比例为27.6% vs 29.0%，无偏好占43.5%），证明了其几乎无损的质量保持能力。在效率方面，ReHyAt相比Flash Attention最高可节省4倍的计算量（FLOPs）。在高通骁龙8 Gen4移动平台上，对于121帧（约7.5秒）的视频，ReHyAt的延迟仅为247ms，而Flash Attention为2964ms（提速约16倍）；内存读写更是降低了11倍以上，且能够在移动端轻松生成超过10秒的视频而不发生OOM。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 视频生成质量 (VBench) | Total Score | 83.79 (20x blocks) | 83.31 (Wan2.1 1.3B Original) | +0.48 (SOTA Competitive) |
| 移动端推理效率 | Latency (121 frames) | 247 ms | 2964 ms (Flash Attention) | 16x Speedup |
| 算力成本 (GPU Hours) | Training Cost | ~160 GPU hours | 12 days x 64 H100 (SANA-Video) | >100x Cheaper |
局限与改进
尽管ReHyAt在整体上表现出色，但作者承认在一些极端效率配置下（如最高效的变体），生成的视频仍存在少量的时间连贯性问题，特别是在快速运动场景下可能出现细微的瑕疵。此外，虽然线性注意力降低了计算复杂度，但RNN的循环推理机制虽然省内存，却在一定程度上牺牲了Transformer特有的并行生成能力，可能会增加单步推理的延迟。该方法目前主要在Wan2.1 1.3B模型上进行了验证，对于更大规模模型（如14B）的直接迁移效果和训练稳定性尚未完全在论文中展开讨论。
独立分析的弱点
尽管ReHyAt取得了成功，但仍存在一些可改进之处。首先，在分块处理机制中，虽然重叠设计缓解了边界效应，但在极高动态场景下，不同Chunk之间的特征传递效率仍可能低于全注意力机制，未来可探索自适应的Chunk大小或动态稀疏注意力策略。其次，当前的多项式特征映射虽然有效，但引入了额外的计算开销和参数量（每块约2.4M），针对移动端极致优化时，这部分可能需要进一步的量化或结构化剪枝。
未来方向
基于ReHyAt的成果，未来的研究方向十分广阔。一是探索其在更大规模模型（如CogVideoX或HunyuanVideo）上的应用，验证其通用性。二是结合模型量化（Quantization）技术，将ReHyAt转化为INT8或INT4精度，进一步降低移动端的部署门槛。三是利用RNN常数内存的优势，开发真正的“无限长度”流式视频生成（Streaming Video Generation）应用，实现视频的实时交互生成。
复现评估
该论文的复现性极佳，这主要得益于其极低的计算需求。作者声称仅需约160 GPU小时即可完成从Wan2.1 1.3B到ReHyAt的转化，这对于大多数实验室和开发者来说是完全可接受的成本。此外，文中详细描述了多项式核的参数配置、两阶段训练的超参（如学习率、Batch Size、迭代步数）以及在VBench上的评估流程，为后续研究提供了清晰的复现指南。
论文图表
