RT-Lynx : Utiliser la sparsité de GEMM de la bonne manière pour les modèles de diffusion

Résumé

Les Transformers de Diffusion (DiT) atteignent des performances élevées en génération d'images, mais entraînent des coûts d'inférence substantiels. Bien que des travaux antérieurs aient réduit ce coût via la quantification et la distillation, la parcimonie semi-structurée, qui peut quasiment diviser par deux le nombre de FLOPs, reste peu explorée. Une raison clé est que la plupart des approches existantes se concentrent sur l'élagage des poids, et qu'élaguer 50 % des poids peut supprimer une capacité critique du modèle et dégrader la qualité de génération. Notre étude montre cependant que les activations des DiT sont intrinsèquement parcimonieuses et significativement plus robustes à la parcimonie semi-structurée N:M que les poids. Motivés par cette observation, nous préconisons un changement de paradigme, passant de l'élagage des poids à l'élagage des activations. Nous proposons RT-Lynx, qui applique la parcimonie N:M aux activations et intègre des techniques de compensation d'erreur pour atténuer la perte de précision. Nous implémentons en outre des noyaux CUDA hautement optimisés, adaptés à ce contexte, permettant d'atteindre une accélération moyenne allant jusqu'à 1,55x dans les couches linéaires. Des expériences approfondies sur plusieurs modèles de diffusion démontrent que notre méthode préserve la qualité de génération des modèles originaux tout en accélérant substantiellement l'inférence.

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.