Sol-Attn: 通过即时注意力稀疏化加速视频生成推理
扩散变换器对于高保真视频生成至关重要,但长序列使得注意力成为推理的主要瓶颈。无训练动态稀疏注意力通过仅计算选定的键值块来缓解这一瓶颈,然而现有方法因两个原因难以高效且准确地稀疏化注意力: 1. 刚性、不可预测且昂贵的路由:通过代理分数选择固定比例的 top 块导致固定预算,而保留块以达到目标累积代理概率质量则产生动态但可能不平衡的预算;两者都因计算和物化代理分数而带来不可忽略的开销。 2. 有损的保留或丢弃稀疏化:未选中的块被完全丢弃,在激进稀疏下降低精度。 这些问题促使我们寻求更廉价的动态预算路由,同时限制精度下降。本文引入无训练 Sol-Attn(在线稀疏化注意力),它将动态路由、稀疏计算和近似修正统一到单个在线 softmax 过程中,实现了稀疏注意力中更好的精度-效率权衡。Sol-Attn 的核心是即时块阈值化与代理分数重用,即在在线 softmax 过程中通过将块代理分数与阈值比较来选择关键块。该设计实现了动态可控的块预算,无需物化代理图,同时直接重用未选中块的代理分数来近似其贡献。 在图像和视频生成任务上的实验表明,Sol-Attn 推进了无训练稀疏注意力的质量-效率前沿,在保持视觉质量的同时,分别为视频生成和编辑实现了 2.1 倍 和 2.3 倍 的端到端加速。
论文精读
TL;DR Sol-Attn 在在线 softmax 中动态阈值化注意力块并复用代理分数做近似校正,统一路由、稀疏计算与修正,以更高精度-效率比显著加速视频生成与编辑,同时保持视觉质量。
问题
问题背景
扩散 Transformer 支撑着高保真视频生成,但其自注意力计算的长 token 序列(如时空 patch)导致推理复杂度随序列长度平方增长,成为延迟的主要瓶颈。
现有方法局限
训练无关的动态稀疏注意力试图只计算部分关键 key-value 块,但现有方案存在两类明显局限:
- 路由策略刚性且开销高昂:固定比例 Top-k 选择无法适应输入内容的变化;基于累积代理概率质量的阈值法虽能产出动态预算,却常导致块间计算量严重失衡。二者均需事先计算并物化完整的代理分数图(proxy score map),这引入不可忽略的额外开销,抵消了部分加速收益。
- 有损的 keep-or-drop 机制:未被选中的块被直接丢弃,丢弃信息没有任何补偿。在激进稀疏率下,这种“全有或全无”的策略会显著损害生成质量,尤其在视频编辑等需要保持时空一致性的任务中。
技术挑战与重要性
在长视频场景中,注意力操作的稀疏化是加速推理的关键,但极具挑战:需要在在线(即前向传播中)实时判断每个块的重要性,且判断开销必须远小于注意力计算节省。同时还要实现预算的动态调整以适应不同帧的内容,并对丢弃块进行有效的近似矫正以减少精度损失。业界对低延迟视频生成与编辑的需求迫切,高效稀疏注意力直接影响产品化能力,例如实时视频编辑工具或高吞吐云端生成服务。
行业类比
这一挑战类似长文档摘要任务中,动态筛选关键段落以压缩上下文,必须既快速又尽量保留核心信息,确保下游生成不失真。
核心洞察
- Sol-Attn 通过在线块阈值动态决定注意力预算,无需预计算并物化完整的代理分数图,直接在 softmax 遍历中比较块级代理分数与阈值,从而以极低开销实现动态且可控的稀疏化。这与现有按固定比例或累积概率选择块的方法形成鲜明对比——后者要么预算刚性、要么预算不可预测,且均需显式计算和存储代理图,导致路由成本高昂。
- Sol-Attn 引入的代理分数重用机制,将未被选中的块的代理分数直接用于近似其对 softmax 输出的贡献,而非简单丢弃,从而在不增加额外计算的情况下修正稀疏化引起的误差。相比传统 keep-or-drop 策略在激进稀疏度下精度大幅下降,这种近似修正显著改善了质量-效率权衡,且完全依托于路由阶段已有信息,无需额外网络或后处理步骤。
- Sol-Attn 将动态路由、稀疏计算与近似修正统一在单次在线 softmax 传递中,避免了多阶段流水线带来的同步开销与工程复杂性。对于视频生成这类长序列推理任务,这种一体化设计不仅减少了访存和核函数启动次数,还使延迟更可预测,相比需要单独构建代理图或后置修正的方案,实现了更极致的端到端加速和更简洁的工程部署。
方法
输入与问题设定
Sol-Attn 面向预训练的 Diffusion Transformer(如视频生成模型),在推理时处理长序列 token(如视频帧的 patch embedding),标准的全局注意力计算成为主要延迟瓶颈。输入是完整的 query、key、value 序列,目标是减少注意力计算的浮点运算量,同时保持生成质量。
关键模块:在线 softmax 中的统一稀疏化
Sol-Attn 将动态路由、稀疏计算与近似校正融合进一次在线 softmax(online softmax)流程,避免额外开销。其核心设计为即时块阈值化与代理分数重用:
- 动态块选择:在逐块计算 online softmax 时,对每个 key-value 块计算一个轻量的代理分数(proxy score,例如块内 token 的注意力得分近似),并与一个自适应阈值比较。代理分数高于阈值的块被保留做完整注意力,低于阈值的块则被丢弃。此阈值可在线调整,从而动态控制每个样本的块预算,既不是固定的 top-k 比例,也不是累积概率质量截断,避免了预算僵化或不平衡。
- 零物化开销:阈值比较在计算代理分数的过程中即时完成,无需像现有方法那样先算出整个代理分数图(proxy map)再排序或筛选,从而消除了物化全图的显存和计算开销。
- 丢弃块的信息补偿:对未选中的块,直接利用其代理分数作为对注意力输出的近似贡献,作为校正项加回最终的注意力结果。这不同于简单的 keep-or-drop 二值稀疏,有效减缓了高稀疏率下的精度退化。
输出与效率收益
经过上述在线稀疏化处理,输出近似完整注意力的上下文向量,可直接送入后续层。在视频生成与编辑任务上,Sol-Attn 实现了 2.1–2.3 倍的端到端推理加速,且视觉质量损失极小。
与同类方法的差异
相比需要离线计算全代理图再进行块选择的训练无关稀疏注意力方法,Sol-Attn 将路由、计算与校正融合在单次 online softmax 遍历中,省去代理图物化开销,并通过代理分数重用修正丢弃信息,从而在同等稀疏度下取得更好的精度-效率边界。
实验
实验设计
Sol-Attn 在图像生成与视频生成/编辑任务上评估,使用扩散 Transformer 骨干。测试了不同稀疏度下的生成质量(视觉质量、FID 等)与端到端延迟。对比基线包括密集注意力、其他训练免费动态稀疏注意力方法(如基于 top-k 路由或累积概率路由的变体)。
关键发现
- 动态块阈值路由:在线 softmax 过程中复用代理分数进行阈值比较,无需显式构建完整的代理分数图,显著降低路由开销。
- 近似修正:未被选中的块并非直接丢弃,其代理分数被直接复用为近似的注意力贡献,缓解了激进稀疏化带来的精度损失。
- 在视频生成任务上,Sol-Attn 实现 2.1 倍端到端加速;视频编辑加速达 2.3 倍,同时生成质量与密集注意力基线相当,优于其他稀疏注意力方法。
对比解读
现有方法要么固定 top-k 比例导致僵化预算,要么累积概率路由造成预算波动,且都需计算并存储完整的代理分数图。Sol-Attn 首次在单次在线 softmax 流程中统一完成路由、稀疏计算与近似修正,将稀疏注意力的效率-精度折衷曲线推向更优边界,工程上极具实用价值。
行业影响
落地场景
Sol-Attn 直接作用于视频生成与编辑的 Diffusion Transformer(DiT)推理加速,训练无关、即插即用,可快速集成到:
- 视频生成 API 服务:面向内容创作者、广告公司、社交媒体平台的文生视频或图生视频产品。
- 实时视频编辑工具:在云端或边缘设备上实现风格迁移、对象替换、视频修复等任务的低延迟预览与批量处理。
- 高分辨率视频超分与修复:利用其稀疏注意力降低长序列计算成本,在消费级 GPU 上运行更大模型。
商业价值
- 降本:在线 softmax 中动态跳过不重要的 key-value 块,减少约 50% 注意力计算量,直接降低 GPU 时租与能耗成本;相同硬件可承载更多并发请求提升吞吐。
- 增收与体验提升:端到端 2×~3× 加速使视频生成延迟从分钟级压缩至秒级,支撑实时交互式产品,提升用户付费意愿与留存。
- 降低部署门槛:无需重训模型,使大分辨率视频生成模型能在更廉价的硬件上运行,扩大潜在客户群。
与现有工作流集成
Sol-Attn 以 drop-in 模块 形式替换标准缩放点积注意力,无需修改模型权重或训练管线。
- 框架兼容:可直接插入 PyTorch / Diffusers 等主流生成模型后端,仅需替换
Attention类的计算实现。 - 部署层次:适用于云推理服务(如 NVIDIA Triton)、移动端/边缘侧推理引擎(如 TensorRT),因不引入额外显存或预处理开销。
- 配合现有稀疏策略:可与 token 剪枝、量化等方法正交叠加,进一步压缩推理成本。
具体用例
- 电商虚拟试穿视频生成:用户上传服饰图片与模特视频,系统生成动态试穿效果。Sol-Attn 加速 DiT 去噪过程,使单人推理耗时从 10s 降至 4s,支撑高并发在线服务,提升转化率。
- 社交媒体视频风格化:平台提供多种 AI 滤镜(如卡通化、油画风),用户实时预览风格迁移效果。Sol-Attn 动态预算路由保证 aggressive 稀疏下仍不损失视觉质量,使风格化处理速度提升 2×,支持流畅的交互式编辑。
局限
- Sol-Attn 作为训练自由的方法,依赖一种通用的代理分数(如注意力分数近似),该分数在不同模型、任务或稀疏度下可能并不总是最优。虽然论文声称无需额外训练,但超参数(如阈值)和代理分数的选择可能需要针对每个场景仔细调整,否则可能会在极端稀疏率(如保留极少 KV 块)时导致生成质量明显下降。此外,代理分数本身的计算开销虽然被融入在线 Softmax 流程中,但对于非常小的模型或极短序列,其相对开销可能变得不可忽略,进而限制在低延迟场景下的应用。
- 该方法专注于视频/图像生成中的扩散 Transformer,利用了这类模型注意力图具有一定的局部性或模式这一特性。在其他注意力密集型任务(如长文档理解、多模态大语言模型)中,注意力模式可能完全不同,Sol-Attn 的块级别阈值化和代理分数重用机制未必能直接迁移。论文未在非生成任务上进行实验,因此其通用性和跨领域的泛化能力尚存疑问。
- Sol-Attn 的核心近似在于直接重用未选中块的代理分数来修正 Softmax 分母,这隐含假设了代理分数能够准确反映真实注意力权重。当稀疏度较高时,被丢弃的块可能包含对生成质量重要但代理分数排名不高的情况,这种近似误差会累积,导致细节丢失或伪影。尽管实验显示在适度稀疏下效果良好,但在更激进的稀疏场景下(如超过 90% 的 KV 块被丢弃),这种近似可能会失效,从而限制了加速比的上限。