Workflow
视觉生成模型
icon
搜索文档
用好视觉Attention局部性,清华、字节提出Token Reorder,无损实现5倍稀疏、4比特量化
机器之心· 2025-06-30 11:18
视觉生成模型的注意力优化 - 视觉生成模型的输入序列长度可达10K-100K,Attention操作占据模型开销的60-80%,成为主要性能瓶颈 [1] - 现有稀疏注意力与低比特量化方法在低稠密度(<50%)和低比特(INT8/INT4)时面临显著性能损失 [1] - 视觉生成任务的注意力模式具有"多样且分散"的特点,这是优化面临的关键挑战 [3][11] 技术方案与创新点 - 提出Token重排方案将多样分散的注意力模式统一为硬件友好的块状模式 [14][19] - 针对每个注意力头选择最优的Token重排方式,共6种可能的置换方式 [24] - 静态稀疏方案在50%稠密度下取得1.73x加速,且不引入额外开销 [55][56] - 量化方案可将PV计算从FP8降至INT4而无精度损失 [46] 性能表现 - 在50%稀疏比下,PAROAttention的CLIPSIM达0.203,与FP16全精度相当 [50] - 20%稀疏比时仍能保持良好生成质量,优于基线50%稀疏比的效果 [45] - 最激进优化方案(50%+INT4)实现近10倍Attention延迟优化 [49] - 量化方案在INT4下PSNR达24.16,优于SageAttentionV2的24.46 [50][56] 系统优化 - 采用算子融合技术将Token重排的额外开销降至可忽略程度 [47] - 通过预取策略将稀疏掩码的显存开销从GB级降至MB级 [47] - 方案完全兼容FlashAttention,无需特殊CUDA Kernel优化 [43] 应用前景 - 方法可启发训练方法优化、图像参数化方式和3D空间位置编码设计 [58] - 技术路线不仅限于推理优化,可推动具有更好归纳偏置的视觉基座模型构建 [58]