如您有工作需要分享,欢迎联系:aigc_to_future
作者:Youping Gu等
解读:AI生成未来
文章链接:https://arxiv.org/pdf/2508.10774
项目链接:http://ziplab.co/BLADE-Homepage/
亮点直击
BLADE,一种新颖且无需数据的联合训练框架,通过将自适应稀疏注意力机制直接融入稀疏感知的步数蒸馏过程,克服了以往顺序或训练无关集成方法的局限性。 自适应块稀疏注意力(ASA),一种动态、内容感知且硬件友好的注意力机制,能够实时生成稀疏掩码以聚焦于关键特征,其性能优于现有的静态稀疏注意力方法。 在多样化模型上实现了显著的端到端推理加速:在 Wan2.1-1.3B 上达到 14.10× 加速,在较短序列的 CogVideoX-5B 上实现稳健的 8.89× 加速。质量持续提升,VBench-2.0 分数在 Wan2.1-1.3B()和 CogVideoX-5B()上均有所提高。
总结速览
解决的问题
推理效率低:扩散Transformer在高质量视频生成中迭代去噪过程缓慢,且长序列的二次注意力计算成本过高,导致推理速度瓶颈。 独立加速策略的局限性: 步数蒸馏(step distillation)和稀疏注意力(sparse attention)单独使用时效果有限。 训练无关的稀疏注意力与蒸馏结合效果不佳,而分阶段训练(先蒸馏后稀疏化)需要昂贵的高质量视频数据。 稀疏注意力的适应性不足: 静态稀疏模式(如固定局部窗口)无法适应视频动态时空特征,导致质量下降。 动态稀疏方法(如VSA、SparseAttention)受限于训练场景或序列长度,灵活性不足。
提出的方案
BLADE框架: 联合训练:将稀疏注意力与步数蒸馏协同优化,而非分阶段处理。 数据无关(data-free):避免依赖高质量视频数据集,降低训练成本。 关键技术组件: 自适应块稀疏注意力(ASA, Adaptive Block-Sparse Attention):动态生成内容感知的稀疏掩码,聚焦关键时空特征。 稀疏感知的步数蒸馏:基于轨迹分布匹配(TDM),将稀疏性直接融入蒸馏过程,加速收敛。
应用的技术
动态稀疏注意力(ASA): 支持训练无关(inference-only)和训练感知(ASA-GT)两种模式,适应不同场景。 通过内容自适应掩码动态选择重要token,平衡计算效率与生成质量。 轨迹分布匹配(TDM): 在蒸馏过程中显式建模稀疏性,使学生模型学习教师模型的紧凑轨迹。 块稀疏计算优化: 减少注意力计算的二次复杂度,提升长序列处理效率。
达到的效果
显著加速: 在Wan2.1-1.3B模型上实现14.10×端到端加速(50步基线对比)。 在CogVideoX-5B(短序列)上实现8.89×加速。 质量提升: VBench-2.0评分提升:CogVideoX-5B从0.534→0.569,Wan2.1-1.3B从0.563→0.570。 人类评估中生成质量优于基线。 灵活性与通用性: 支持不同模型规模(1.3B~5B参数量)和视频序列长度。 动态稀疏机制适应多样内容,减少固定模式导致的质量损失。
方法
整体架构
BLADE 是一个通过将动态稀疏性协同整合到强大的步数蒸馏过程中来加速视频扩散模型的整体框架。如下图1所示,本文的架构基于师生范式。教师模型 是一个预训练的高质量但计算成本高的多步扩散模型。学生模型 最初与教师模型共享相同的基于Transformer(DiT)的架构和权重。本文关键创新在于将学生模型中的标准自注意力层替换为提出的自适应块稀疏注意力(ASA, Adaptive Block-Sparse Attention) 机制,以实现少步生成。训练过程遵循 轨迹分布匹配(TDM, Trajectory Distribution Matching) 范式。在每次迭代中,稀疏学生模型 生成一个中间轨迹,随后通过无需数据的分数蒸馏损失引导该轨迹与教师模型的轨迹分布匹配。这确保了学生模型在ASA施加的计算约束下学习生成高质量输出。
预备知识:轨迹分布匹配(TDM)
轨迹分布匹配(TDM) 是一种先进的蒸馏框架,旨在创建高效的少步扩散模型。其核心思想是在分布级别上对齐学生模型与教师模型的整个生成轨迹,而非要求精确的实例级匹配。这一目标通过无需数据的分数蒸馏过程实现,该过程依赖三个关键组件:
预训练的教师模型 ,提供真实数据分数 。 学生生成器 ,学习以少量步骤生成高保真样本。 伪分数模型 ,通过近似学生的难解样本分数提供伪分数 。
训练过程涉及两个相互交织的目标:一个针对伪分数模型,另一个针对学生生成器。
训练伪分数模型()
分数蒸馏过程需要学生模型的分数函数 ,但该函数难解。TDM通过引入伪分数模型 来解决这一问题,该神经网络在训练过程中同步学习以近似学生的分数。为确保近似准确,伪分数模型 通过以下去噪目标进行训练:
其中,干净目标首先由学生模型通过对输入去噪获得。随后通过扰动该目标生成噪声样本,模型学习从该噪声输入中预测干净样本。
训练学生生成器()在同时获得教师分数和学生自身分数估计的情况下,可训练学生生成器。其目标是最小化学生轨迹分布与教师轨迹分布之间的KL散度。该对齐操作在扩散过程的个阶段执行,确保学生高效学习跟随教师的生成路径。核心蒸馏损失为:
在实际应用中,通过匹配分数来实现该KL散度的最小化。通过用伪分数模型的输出替代学生模型难解的真实分数来计算该目标的梯度,从而得到以下梯度近似:
自适应块稀疏注意力(ASA)
本文工作的核心设计是自适应块稀疏注意力(ASA)机制,该机制通过动态剪枝注意力矩阵来聚焦关键时空交互。这种内容感知方法克服了先前工作中静态掩码的局限性,其流程包含预处理步骤和动态掩码生成阶段。
预处理:保持局部性的token重排
输入矩阵、和(表示展平的视频token序列)首先被分块重组。关键预处理步骤是通过Gilbert空间填充曲线 对token重新排序,以修复标准光栅扫描token化破坏的空间局部性。这确保分块后的token语义更连贯,包含空间连续信息,从而显著提升后续基于阈值的剪枝效果。
步骤1:高效块重要性估计
理论上,可先计算完整的稠密注意力矩阵,将其划分为的块,并对每个块进行最大池化,得到降采样的重要性矩阵(每个元素表示对应块的重要性)。随后对每行应用阈值生成稀疏掩码,使每个查询块仅关注最关键的键-值块。然而,计算完整矩阵使得该方法无法实际加速。
为解决此问题,本文提出一种高效的在线近似方法:从和的每个块中采样个代表token()构成小矩阵和,并基于这些采样token计算低分辨率注意力图,进而导出块重要性矩阵。该方法将掩码生成复杂度从降至约(为序列长度),使在线掩码生成可行。理论分析表明,在适当选择和时,该近似能稳定生成接近完整注意力图的掩码,同时大幅降低计算成本(详见附录证明)。
步骤2.1:稀疏掩码构建
获得后,基于阈值策略生成最终稀疏注意力掩码:对每行按注意力分数降序排序,并选择最少的键块使其累积分数超过预设阈值(如90%)。这种动态剪枝保留了关键注意力路径,跳过信息量较少的块,实现了精度与效率的灵活权衡。
步骤2.2:计算实现基于上述掩码生成技术,我们定义两种机制变体:
标准ASA(训练无关)
基础形式中,生成的二进制稀疏掩码直接与块稀疏注意力核集成。该变体可直接应用于预训练模型而无需微调,通过聚焦细粒度关键信息实现直接推理加速。带全局token的ASA(用于蒸馏)
为缓解高稀疏率下的全局信息丢失问题,引入增强变体:对键()和值()矩阵添加"全局token",通过对窗口大小进行均值池化生成,将序列长度压缩至原/的。增强矩阵构造为(同理)。注意力计算时,查询与原区域的交互由稀疏掩码控制以保留细节;对全局token区域则在softmax前分数添加固定偏置,补偿均值池化的稀释效应,使每个全局token能代表其个组成token的完整重要性。该设计确保所有查询保持全局上下文感知,避免过度剪枝导致的信息损失。
全文将基础版本称为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
致谢
如果您觉得这篇文章对你有帮助或启发,请不吝点赞、在看、转发,让更多人受益。同时,欢迎给个星标⭐,以便第一时间收到我的最新推送。每一个互动都是对我最大的鼓励。让我们携手并进,共同探索未知,见证一个充满希望和伟大的未来!
技术交流
加入「AI生成未来社区」群聊,一起交流讨论,涉及 图像生成、视频生成、3D生成、具身智能等多个不同方向,备注不同方向邀请入群!可添加小助手备注方向加群!
没有评论:
发表评论