Résumé

Dans les modèles Diffusion Transformer (DiT), notamment pour la génération vidéo, la latence de l’attention constitue un goulot d’étranglement majeur en raison de la longueur des séquences et de la complexité quadratique associée. Nous constatons que les poids d’attention peuvent être séparés en deux parties : une petite fraction de poids importants à rang élevé, et le reste des poids à très faible rang. Cette observation suggère naturellement d’appliquer une accélération creuse à la première partie et une accélération de faible rang à la seconde. À partir de cette constatation, nous proposons SLA (Sparse-Linear Attention), une méthode d’attention apprenable qui fusionne l’attention creuse et l’attention linéaire afin d’accélérer les modèles de diffusion. SLA classe les poids d’attention en trois catégories : critiques, marginales et négligeables. Elle applique une attention de complexité O(N²) aux poids critiques, une attention de complexité O(N) aux poids marginaux, et ignore les poids négligeables. SLA combine ces calculs dans un unique noyau GPU et prend en charge à la fois les passes avant et arrière. En n’effectuant qu’un petit nombre d’étapes de fine-tuning avec SLA, les modèles DiT atteignent une réduction de 20 fois du calcul d’attention, entraînant une accélération significative sans perte de qualité de génération. Les expériences montrent que SLA réduit le calcul d’attention de 95 % sans dégrader la qualité de génération en boucle complète, surpassant ainsi les méthodes de référence. En outre, nous avons implémenté un noyau GPU efficace pour SLA, qui permet une accélération de 13,7 fois du calcul d’attention et une accélération de 2,2 fois en boucle complète pour la génération vidéo sur Wan2.1-1.3B.