通过激活稀疏化加速扩散模型推理,保持生成质量的同时提升速度。
RT-Lynx: Putting GEMM Sparsity in the Right Place for Diffusion Models

- 将稀疏化从权重转移到激活,更适应扩散模型特性
- 在多个模型上实现平均1.55倍加速,生成质量无损失
- 专为激活稀疏设计的优化内核,适合部署于高性能计算场景
扩散Transformer(DiT)在图像生成中表现优异,但推理开销大。尽管已有研究通过量化和蒸馏降低开销,半结构化稀疏性(可近似减半计算量)仍被忽视。主要原因在于多数方法聚焦于权重稀疏化,剪掉50%权重会破坏模型容量并降低生成质量。我们发现DiT的激活值本身具有内在稀疏性,且对N:M半结构化稀疏化更为鲁棒。基于此,我们提出范式转变:从权重稀疏化转向激活稀疏化。提出RT-Lynx方法,在激活上应用N:M稀疏化,并引入误差补偿机制缓解精度损失。进一步开发了针对该场景的高度优化CUDA内核,使线性层平均提速1.55倍。大量实验表明,该方法在多个扩散模型上保持原始生成质量的同时显著加速推理。
原文摘要 · Abstract (English)
Diffusion Transformers (DiT) achieve strong performance in image generation but incur substantial inference costs. While prior work has reduced this cost via quantization and distillation, semi-structured sparsity, which can nearly halve FLOPs, remains underexplored. A key reason is that most existing approaches focus on weight sparsification, and pruning 50% of the weights can remove critical model capacity and degrade generation quality. Our study, however, shows that DiT activations are intrinsically sparse and significantly more robust to N:M semi-structured sparsification than weights. Motivated by this observation, we advocate a paradigm shift from weight sparsification to activation sparsification. We propose RT-Lynx, which applies N:M sparsification to activations and incorporates error-compensation techniques to mitigate accuracy loss. We further implement highly optimized CUDA kernels tailored to this setting, achieving up to a 1.55x speedup on average in linear layers. Extensive experiments across multiple diffusion models demonstrate that our method preserves the generation quality of the original models while substantially accelerating inference.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。