作者:Youping Gu等
解读:AI生成未来
文章链接:https://arxiv.org/pdf/2508.10774 项目链接:http://ziplab.co/BLADE-Homepage/
亮点直击
BLADE 是一个通过将动态稀疏性协同整合到强大的步数蒸馏过程中来加速视频扩散模型的整体框架。如下图1所示,本文的架构基于师生范式。教师模型 是一个预训练的高质量但计算成本高的多步扩散模型。学生模型 最初与教师模型共享相同的基于Transformer(DiT)的架构和权重。本文关键创新在于将学生模型中的标准自注意力层替换为提出的自适应块稀疏注意力(ASA, Adaptive Block-Sparse Attention) 机制,以实现少步生成。训练过程遵循 轨迹分布匹配(TDM, Trajectory Distribution Matching) 范式。在每次迭代中,稀疏学生模型 生成一个中间轨迹,随后通过无需数据的分数蒸馏损失引导该轨迹与教师模型的轨迹分布匹配。这确保了学生模型在ASA施加的计算约束下学习生成高质量输出。
轨迹分布匹配(TDM) 是一种先进的蒸馏框架,旨在创建高效的少步扩散模型。其核心思想是在分布级别上对齐学生模型与教师模型的整个生成轨迹,而非要求精确的实例级匹配。这一目标通过无需数据的分数蒸馏过程实现,该过程依赖三个关键组件:
训练过程涉及两个相互交织的目标:一个针对伪分数模型,另一个针对学生生成器。
训练伪分数模型() 分数蒸馏过程需要学生模型的分数函数 ,但该函数难解。TDM通过引入伪分数模型 来解决这一问题,该神经网络在训练过程中同步学习以近似学生的分数。为确保近似准确,伪分数模型 通过以下去噪目标进行训练:
其中,干净目标首先由学生模型通过对输入去噪获得。随后通过扰动该目标生成噪声样本,模型学习从该噪声输入中预测干净样本。
训练学生生成器()在同时获得教师分数和学生自身分数估计的情况下,可训练学生生成器。其目标是最小化学生轨迹分布与教师轨迹分布之间的KL散度。该对齐操作在扩散过程的个阶段执行,确保学生高效学习跟随教师的生成路径。核心蒸馏损失为:
在实际应用中,通过匹配分数来实现该KL散度的最小化。通过用伪分数模型的输出替代学生模型难解的真实分数来计算该目标的梯度,从而得到以下梯度近似:
本文工作的核心设计是自适应块稀疏注意力(ASA)机制,该机制通过动态剪枝注意力矩阵来聚焦关键时空交互。这种内容感知方法克服了先前工作中静态掩码的局限性,其流程包含预处理步骤和动态掩码生成阶段。
预处理:保持局部性的token重排 输入矩阵、和(表示展平的视频token序列)首先被分块重组。关键预处理步骤是通过Gilbert空间填充曲线 对token重新排序,以修复标准光栅扫描token化破坏的空间局部性。这确保分块后的token语义更连贯,包含空间连续信息,从而显著提升后续基于阈值的剪枝效果。
步骤1:高效块重要性估计 理论上,可先计算完整的稠密注意力矩阵,将其划分为的块,并对每个块进行最大池化,得到降采样的重要性矩阵(每个元素表示对应块的重要性)。随后对每行应用阈值生成稀疏掩码,使每个查询块仅关注最关键的键-值块。然而,计算完整矩阵使得该方法无法实际加速。
为解决此问题,本文提出一种高效的在线近似方法:从和的每个块中采样个代表token()构成小矩阵和,并基于这些采样token计算低分辨率注意力图,进而导出块重要性矩阵。该方法将掩码生成复杂度从降至约(为序列长度),使在线掩码生成可行。理论分析表明,在适当选择和时,该近似能稳定生成接近完整注意力图的掩码,同时大幅降低计算成本(详见附录证明)。
步骤2.1:稀疏掩码构建 获得后,基于阈值策略生成最终稀疏注意力掩码:对每行按注意力分数降序排序,并选择最少的键块使其累积分数超过预设阈值(如90%)。这种动态剪枝保留了关键注意力路径,跳过信息量较少的块,实现了精度与效率的灵活权衡。
步骤2.2:计算实现基于上述掩码生成技术,我们定义两种机制变体:
全文将基础版本称为ASA,增强版本称为ASA-GT。
Video-BLADE框架的核心是稀疏感知蒸馏原则。与以往将稀疏化作为训练后压缩步骤的方法不同,将ASA直接嵌入TDM训练循环:每次迭代时,学生模型通过ASA机制生成轨迹,分布匹配损失则在动态稀疏约束下更新学生权重以提升输出质量。 这种协同设计形成强正则化,迫使模型学习更具语义鲁棒性的表示,往往能产生更优的感知质量。
模型。在两个文本到视频扩散模型上评估BLADE:CogVideoX-5B和Wan2.1-1.3B。这些模型代表了不同的架构和规模,使我们能够测试方法的通用性。
数据集。本文的训练过程由包含10,000个文本提示的数据集指导。这些提示从JourneyDB基准测试中采样,并使用Qwen2.5-3B-Instruct模型进一步优化质量和多样性。
指标。使用一组标准指标评估性能:VBench-1.0、VBench-2.0、SSIM & PSNR以及人工评估。
实现细节。除非另有说明,本文使用块大小,每块采样个token用于注意力探测。蒸馏通常运行250-500次迭代。CogVideoX-5B和Wan2.1-1.3B的实验在8块A800(80GB)GPU集群上进行。
对比方法。ASA GT、ASA、STA和RaA分别表示使用我们的自适应注意力、其免训练变体、滑动瓦片注意力和径向注意力。FA2指FlashAttention-2。
实验表明,Video-BLADE在不牺牲生成质量(甚至通常提升质量)的情况下实现了显著加速。
质量分析。下表1展示了CogVideoX-5B和Wan2.1-1.3B在VBench-2.0基准测试中的结果,对比了多种方法,包括我们提出的ASA GT、稀疏基线STA、FA2以及50步密集基线。
对于CogVideoX-5B,ASA GT在所有主要质量维度上均表现出一致且全面的提升。其VBench-2.0总分最高(),优于50步基线和FA2,并在创造力、可控性和物理合理性等关键指标上领先。值得注意的是,ASA GT仅通过8个解码步骤和17k令牌的短序列实现了这一性能,速度提升达,同时生成质量更高。这些结果表明,即使在序列长度极度受限的情况下,ASA GT仍能实现稳健的生成质量。
对于Wan2.1-1.3B,ASA GT同样展现出明显优势。其VBench-2.0得分为,人工保真度最高(),物理合理性表现强劲,且仅需原始推理时间的(加速)。与稀疏性相似的STA相比,ASA GT在几乎所有指标上均显著更优。尽管FA2总分略高于ASA GT,但其可控性较弱且计算成本更高。附录中提供了多样本视频的视觉对比图集。
从结果中发现一个有趣的现象:尽管BLADE具有高稀疏性和少量推理步骤,但其质量仍能超越50步密集基线。我们将此归因于联合训练框架引入的正则化效应。50步教师的长期迭代轨迹有时会累积数值误差或过拟合噪声及不连贯细节。相比之下,稀疏感知蒸馏迫使学生模型学习更直接、稳定的生成路径(这一原理与DMD2等先前工作的发现一致),使其捕捉最关键的语义信息,同时隐式过滤教师过程中的“弯路”和噪声。自适应稀疏性通过仅聚焦最显著特征进一步强化了这一效果。附录中的注意力图分析提供了视觉佐证。因此,最终模型不仅是更快的近似,还能成为更鲁棒、连贯的生成器。我们在更强调语义保真度的VBench-2.0上评估模型,这与我们方法的优势高度契合。
这些发现验证了ASA GT能很好地泛化到不同模型规模和视频长度,并通过稀疏感知蒸馏与全局上下文集成,在效率和感知质量之间实现了出色平衡。
效率分析。在核函数层面,本文的ASA实现相比8步FA2基线使用的标准密集注意力实现了加速( vs. ),这得益于的有效稀疏率。这种底层增益直接转化为显著的端到端加速:基于ASA的模型在秒内完成生成,而其密集对应版本需秒——实现的端到端加速。
值得注意的是,虽然核函数加速超过,但端到端增益呈次线性。这表明在蒸馏模型中注意力不再是主要瓶颈,其他操作(如VAE编码器/解码器和Transformer中的非注意力层)开始主导运行时。这一转变验证了我们的定向核函数优化在最小化现代扩散流程中注意力开销方面的有效性。
为单独评估ASA机制的性能,在Wan2.1-1.3B上以免训练推理设置将其与其他稀疏注意力方法对比。下表3显示,在的相近稀疏度下,ASA在PSNR和SSIM上均显著优于STA与SVG,确立了其作为动态注意力机制的优越性。下图3展示了不同方法的视频采样结果,附录中提供了包括人工评估在内的进一步消融研究。
本文提出BLADE框架,通过协同设计动态内容感知的自适应块稀疏注意力(ASA)与无数据的轨迹分布匹配(TDM)蒸馏流程,显著提升视频扩散模型的推理效率且不牺牲生成质量。实验表明,通过训练中引入稀疏感知,其视觉质量与内在保真度甚至可超越原始多步教师模型和密集蒸馏学生模型。
本文在多种视频模型上的实验验证了该方法在核函数效率、端到端推理速度和生成质量(基于VBench-2.0与人工评估)上的显著提升。
局限性与未来方向。当前实验仅针对中等长度视频序列,扩展ASA至生成长达数分钟(数十万tokens)的视频是重要下一步。此外,现有ASA核函数基于Triton实现,未能完全释放理论加速潜力,未来计划开发更优化的CUDA实现。稀疏感知训练作为正则化方法的潜力也有望拓展至其他生成领域。
[1] Video-BLADE: Block-Sparse Attention Meets Step Distillation for Efficient Video Generation
如果您觉得这篇文章对你有帮助或启发,请不吝点赞、在看、转发,让更多人受益。