AB-Attention accélère la génération vidéo jusqu'à 2,23× en sautant 92% de l'attention

L'équipe InfRec, avec des chercheurs de CUHK Shenzhen, des universités du Shandong, de Tokyo et de Fudan, de Berkeley, de HKUST et de NUS, a publié AB-Attention. Cette méthode accélère la génération vidéo dans les transformeurs de diffusion (Wan2.1 et HunyuanVideo) sans réentraîner les poids.

L'opération la plus coûteuse de ces modèles est l'attention : chaque fragment de vidéo est comparé à tous les autres, dans l'espace et dans le temps, si bien que le coût croît de façon quadratique avec la longueur de la séquence. AB-Attention saute 89–92% de ces calculs. La vidéo obtenue reste plus proche de celle produite avec l'attention complète « honnête » que celle des concurrents.

Le PSNR (métrique de similarité avec la vidéo de référence, exprimée en décibels : plus il est élevé, plus la vidéo est proche) est meilleur de 1,6 à 4,4 dB, et l'ensemble de la génération sur une seule carte NVIDIA H20 est 1,84 à 2,23 fois plus rapide.

Où les anciens schémas d'attention parcimonieuse atteignent leurs limites

L'attention parcimonieuse ne calcule pas toutes les paires « requête — clé », mais seulement une partie. Les premières méthodes fixaient un motif à l'avance, par exemple des fenêtres 3D locales. Elles ignorent ce qui compte pour un prompt donné, une tête d'attention donnée ou une étape de débruitage donnée, et doivent donc conserver beaucoup de calculs superflus. Les méthodes dynamiques sélectionnent les blocs de clés et de valeurs d'après les activations courantes, mais les auteurs leur relèvent trois faiblesses :

• Les estimations d'importance sont grossières ou coûteuses. Moyenner un bloc en un seul vecteur mélange des tokens au comportement différent, et le clustering global (SVG2) exige un k-means sur toute la séquence ainsi qu'une permutation des tokens. • La sélection top-k coupe la queue de distribution. Un bloc écarté disparaît à la fois du numérateur de l'attention et de la normalisation du softmax. De meilleures estimations n'y changeraient rien : même une sélection top-k parfaite, avec 10,8% des clés, ne conserve qu'environ 69% de la masse d'attention. • L'unité de sélection est la tuile de calcul. Un bloc est conservé ou écarté en entier, même si seules quelques-unes de ses clés sont utiles.

Les trois techniques d'AB-Attention

La première, IBCS (Intra-Block Clustering Selector), estime l'importance des blocs de clés. Les requêtes sont découpées en blocs de 128 tokens. Dans chaque bloc, elles sont regroupées en 8 clusters, chacun doté d'un centroïde et d'un poids égal à la part des requêtes qu'il contient. Un bloc de clés reçoit pour score l'attention pondérée des centroïdes sur ses clés moyennées. Les clusters servent uniquement à l'estimation : le calcul lui-même suit l'ordre d'origine des tokens. Ainsi, un bloc regardé par un seul cluster n'est pas perdu, contrairement à ce qui se passe quand on moyenne tout le bloc. L'écart de rappel de sélection entre la moyenne et un top-k idéal est réduit de moitié.

La deuxième, SIP (Sampling with Inclusion-Probability Weighting), remplace la sélection brutale par un échantillonnage. Un bloc entre dans l'échantillon avec une probabilité proportionnelle à son score. Les blocs les plus importants sont toujours retenus, et les autres gardent une chance non nulle. Chaque bloc retenu est multiplié par l'inverse de sa probabilité d'inclusion : un bloc avec une probabilité de 25% compte pour quatre. C'est l'estimateur classique de Horvitz-Thompson, qui fournit des estimations non biaisées du numérateur comme de la normalisation du softmax. L'erreur d'attention est inférieure de 41% à celle d'un top-k idéal avec le même nombre de blocs. Aucune passe supplémentaire de correction n'est nécessaire : le calcul reste une attention softmax ordinaire avec décalage des logits.

La troisième technique consiste en une sélection fine sur un matériel « épais ». Les blocs de clés et de valeurs ne font que 4 tokens, mais le noyau assemble les blocs choisis, par indices, en tuiles pleine taille pour les cœurs tensoriels. Les tokens ne sont pas permutés et aucune copie superflue n'est créée. Réduire la taille des blocs de 128 à 4 tokens fait passer de 59% à 69% la part de masse d'attention conservée par une sélection idéale, sans que le nombre de tuiles augmente à budget fixe. Le plan d'échantillonnage est recalculé toutes les quatre étapes de débruitage, tandis que l'attention elle-même est calculée à chaque étape.

Protocole de mesure et résultats

Les auteurs ont comparé trois modèles en 720p : deux versions de Wan2.1 (81 images) et HunyuanVideo (129 images), sur 600 prompts du benchmark Penguin Video. La référence était la vidéo générée avec l'attention complète, à prompt et graine de génération identiques. Outre le PSNR, ils ont calculé le SSIM et le LPIPS, ainsi que cinq évaluations portant sur la vidéo elle-même : esthétique, cohérence de l'arrière-plan, qualité d'image, fluidité du mouvement et constance du sujet. Les concurrents sont RadialAttention (masque statique à décroissance d'énergie), SVG2 et DFSAttn, tous exécutés avec leur code officiel et le backend FlashInfer avec FlashAttention-3. Le premier bloc du transformeur et les 10 premières des 50 étapes utilisent l'attention complète pour toutes les méthodes.

• Parcimonie : 89,16% sur Wan2.1 et 92,35% sur HunyuanVideo. • Les concurrents calculent dans le même temps 1,8 à 3,6 fois plus d'attention. • AB-Attention est la meilleure en PSNR, SSIM et LPIPS sur les trois modèles. • Accélération de bout en bout : 1,84–2,23× sur une seule H20 en BF16.

Ce que cela change en pratique

La méthode ne demande aucun réentraînement : on peut donc la brancher sur des poids existants. Pour les équipes qui font tourner Wan2.1 ou HunyuanVideo, cela représente presque deux fois moins de temps par vidéo, avec une qualité proche de la référence. Le gain devrait augmenter sur les vidéos longues et en haute résolution, où l'attention pèse davantage dans la latence. Il existe toutefois des limites : les mesures ont été réalisées sur H20 et sur trois modèles, et le noyau comme le plan d'échantillonnage sont conçus pour les cœurs tensoriels NVIDIA. Le code est ouvert sur GitHub, dans le dépôt InfRec-Cardinal/AB-Attention : les résultats peuvent donc être revérifiés sur son propre matériel.