💻 computer science

PISA: Piecewise Sparse Attention Is Wiser for Efficient Diffusion Transformers

Cet article introduit PISA, un mécanisme d'attention parcielle par morceaux (Piecewise Sparse Attention) sans entraînement qui atteint une complexité sous-quadratique dans les Transformers de diffusion en approximant les blocs non critiques via une expansion de Taylor par blocs plutôt qu'en les éliminant, offrant ainsi des accélérations significatives tout en maintenant une haute qualité de génération.

Auteurs originaux : Haopeng Li, Shitong Shao, Wenliang Zhong, Zikai Zhou, Lichen Bai, Hui Xiong, Zeke Xie

Publié 2026-02-04

📖 5 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Haopeng Li, Shitong Shao, Wenliang Zhong, Zikai Zhou, Lichen Bai, Hui Xiong, Zeke Xie

Article original sous licence CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Ceci est une explication générée par l'IA de l'article ci-dessous. Elle n'a pas été rédigée ni approuvée par les auteurs. Pour une précision technique, consultez l'article original. Lire la clause de non-responsabilité complète

Imaginez que vous essayez de peindre une fresque monumentale et incroyablement détaillée (comme une vidéo ou une image haute définition) en utilisant une équipe d'artistes. Dans le monde de l'IA, cette « équipe » est un Transformer de Diffusion, et le processus de « peinture » consiste à examiner chaque coup de pinceau (token) pour décider quel doit être le coup suivant.

Le problème ? La façon dont ces équipes d'IA travaillent actuellement est semblable à une règle stricte : « Regardez chaque un des autres coups de pinceau dans l'ensemble de la fresque pour décider de votre prochain mouvement. » À mesure que la fresque s'agrandit (résolution plus élevée ou vidéos plus longues), cela devient impossible. Les artistes sont submergés, le processus ralentit considérablement et l'ordinateur tombe en panne de ressources. C'est le « goulot d'étranglement de la complexité quadratique » mentionné dans l'article.

Pour résoudre cela, les méthodes précédentes ont tenté une stratégie de « Garder ou Jeter ». Elles disaient : « D'accord, ignorons 80 % des coups de pinceau et ne regardons que les 20 % les plus importants. »

La faille : C'est comme essayer de peindre un visage mais ignorer les yeux et la bouche parce qu'ils ne faisaient pas partie de votre « liste des 20 % les plus importants ». Le résultat est souvent flou, buggé ou manque de détails cruciaux.

La nouvelle idée : PISA (Attention parcellaire par morceaux)

Les auteurs de cet article proposent une méthode plus intelligente appelée PISA. Au lieu de simplement ignorer les parties « non importantes », PISA traite la fresque comme un puzzle par morceaux.

Voici comment cela fonctionne, en utilisant une analogie simple :

1. La stratégie « Exact vs Approximatif »

Imaginez que vous êtes un chef préparant une soupe géante pour une foule.

L'ancienne méthode (Attention dense) : Vous goûtez chaque cuillerée de soupe pour obtenir la saveur parfaite. C'est précis, mais cela prend un temps infini.
La méthode du « Jet » (Attention parcellaire standard) : Vous ne goûtez que les premiers 20 % de la soupe et ignorez le reste. La soupe a un goût étrange parce que vous avez manqué le sel au fond.
La méthode PISA :
- Blocs critiques (La partie « Exacte ») : Vous identifiez les ingrédients les plus importants (comme les épices principales ou la viande). Vous goûtez ceux-là exactement et avec soin.
- Blocs non critiques (La partie « Approximative ») : Pour le reste de la soupe (l'eau, le bouillon, les légumes), vous n'avez pas besoin de goûter chaque goutte. Au lieu de cela, vous utilisez un raccourci mathématique (un développement de Taylor) pour estimer la saveur en fonction du goût moyen de cette section.

2. Pourquoi est-ce plus « sage » ?

L'article a découvert quelque chose de fascinant : les parties « non importantes » de l'attention de l'IA (les blocs non critiques) sont en réalité très prévisibles. Elles suivent un schéma fluide et calme.

Parce qu'elles sont prévisibles, vous n'avez pas besoin de les jeter. Vous pouvez les approximer très rapidement sans perdre la « saveur » de l'image finale.
C'est comme estimer la température d'une grande pièce en mesurant la température moyenne des coins, plutôt que de mesurer chaque centimètre cube d'air.

3. Le résultat : La vitesse sans le sacrifice

En combinant des calculs exacts pour les parties importantes et des approximations intelligentes pour le reste, PISA atteint deux objectifs :

Vitesse : Il est 2 à 2,5 fois plus rapide que les meilleures méthodes actuelles. Dans les tests de l'article, générer une vidéo qui prenait auparavant 26 minutes n'en prend plus qu'environ 11.
Qualité : Les images et les vidéos sont tout aussi bonnes que la version lente et « parfaite ». En fait, sur certains tests, PISA a produit de meilleurs résultats que d'autres méthodes « rapides » qui se contentaient de supprimer des informations.

La « Magie » en coulisses

L'article mentionne quelques astuces techniques qui rendent cela possible sans avoir besoin de réentraîner l'IA :

Aucun réentraînement nécessaire : Parce que PISA imite très étroitement la façon de penser « parfaite » originale, vous pouvez l'intégrer dans des modèles d'IA existants (comme Wan2.1 ou FLUX.1) et cela fonctionne simplement. Vous n'avez pas besoin d'enseigner un nouveau langage à l'IA.
Le « Noyau Hybride » (Hybrid Kernel) : Les auteurs ont construit un programme informatique personnalisé (un noyau) qui bascule instantanément entre « goûter exactement » et « estimer », afin que l'ordinateur ne soit pas confus ou ralenti.

Résumé

Considérez PISA comme un gestionnaire intelligent pour un studio d'art par IA.

Les anciens gestionnaires : « Ignorez 80 % du travail ! » (Résultat : Une œuvre médiocre).
Le gestionnaire PISA : « Concentrez 100 % de l'effort sur les détails qui comptent, et utilisez une supposition rapide et intelligente pour le bruit de fond. » (Résultat : Une œuvre de haute qualité et rapide).

L'article prouve que cette « supposition » n'est pas un raccourci paresseux ; c'est une façon mathématiquement solide de gagner du temps tout en préservant l'intégrité du chef-d'œuvre.

Résumé Technique : PISA - Attention Éparse par Morceaux (Piecewise Sparse Attention) pour les Transformers de Diffusion Efficaces

1. Énoncé du Problème

Les Transformers de Diffusion (DiTs) sont devenus le standard pour la génération d'images et de vidéos de haute fidélité. Cependant, leur extensibilité est limitée par la complexité quadratique ( $O(L^2)$ ) du mécanisme d'auto-attention à mesure que les longueurs de séquence d'entrée ( $L$ ) augmentent pour des résolutions plus élevées ou des vidéos plus longues.

Les solutions existantes présentent des limitations significatives :

Attention Éparse (« Garder ou Supprimer ») : Les méthodes qui rejettent des blocs de clés-valeurs (KV) non critiques accélèrent le calcul mais souffrent d'une dégradation des performances à des ratios de parcimonie élevés en raison de la perte de contexte. Elles peinent également avec les séquences plus courtes où le surcoût de sélection l'emporte sur les bénéfices.
Attention Linéaire/Hybride : Les approches qui modifient la distribution de l'attention (par exemple, via des mappages de caractéristiques ou des combinaations additives de branches éparses et linéaires) perturbent souvent la normalisation intrinsèque des modèles pré-entraînés. Cela nécessite un réentraînement ou un ajustement fin coûteux, empêissant la réutilisation directe des poids pré-entraînés.
Troncation Brutale : Le fait de supprimer directement des blocs entraîne des écarts marqués par rapport à la distribution d'attention originale, provoquant une perte de qualité dans le contenu généré.

Le défi central est de concevoir un mécanisme capable d'accélérer l'inférence avec une complexité sous-quadratique tout en maintenant la précision de l'attention complète et en restant compatible avec les poids pré-entraînés sans réentraînement.

2. Méthodologie : Attention Éparse par Morceaux (PISA)

PISA introduit un nouveau paradigme « exact ou approximatif », dépassant la stratégie traditionnelle de « garder ou supprimer ». Il traite l'attention comme un processus par morceaux qui couvre toute l'étendue de l'attention avec une complexité sous-quadratique.

Concepts Clés

Les auteurs ont analysé les DiTs pré-entraînés et ont découvert deux propriétés critiques :

Stabilité Distributionnelle : Les scores d'attention pré-softmax des blocs non critiques (non essentiels) présentent une distribution symétrique en forme de cloche centrée autour de zéro ou de valeurs négatives. Cela les rend très propices à une approximation précise via un développement de Taylor, plutôt qu'à un calcul exact ou une suppression.
Cohérence de Normalisation : Pour maintenir la compatibilité sans réentraînement, les approximations doivent être intégrées de manière interne dans le numérateur et le dénominateur du softmax, plutôt que de simplement ajouter les sorties de différentes branches (ce qui violerait les règles de somme pondérée).

Algorithme Central

PISA partitionne les blocs KV en deux ensembles :

Blocs Critiques ( $S_i$ ) : Calculés de manière exacte pour préserver les informations de haute saillance.
Blocs Non Critiques ( $U_i$ ) : Approximés efficacement en utilisant un développement de Taylor par bloc.

La sortie $o_t$ est dérivée en normalisant l'agrégation pondérée des valeurs :
$o_t = \frac{N_t}{D_t}$

Où le dénominateur ( $D_t$ ) et le numérateur ( $N_t$ ) sont composés de :

Termes Épars Exacts : Calculés pour les blocs sélectionnés $j \in S_i$ .
Approximation du Zéroième Ordre par Bloc : Pour les blocs non sélectionnés, la contribution est approximée en utilisant le centroïde du bloc $\bar{k}_j$ et la somme des valeurs.
Correction Globale du Premier Ordre : Au lieu de calculer un terme de premier ordre distinct pour chaque bloc non sélectionné (ce qui serait limité par la mémoire), PISA emploie une approximation d'ordre hybride. Il utilise une statistique globale $\bar{H}$ (la moyenne des matrices de premier ordre par bloc sur l'ensemble des blocs) pour approximer le gradient de premier ordre pour l'ensemble du groupe de queue.

La formulation pour le numérateur $N_t$ devient :
$N_t = \sum_{j \in S_i} \sum_{n} e^{q_t k_{j,n}^\top} v_{j,n} + \sum_{j \in U_i} \alpha_{t,j} \left(\sum_{n} v_{j,n}\right) + q_t \bar{H} \sum_{j \in U_i} \alpha_{t,j}$
Où $\alpha_{t,j} = \exp(q_t \bar{k}_j^\top)$ .

Routage Sensible à la Covariance

Pour minimiser l'erreur d'approximation, PISA utilise une Stratégie Top-k Sensible à la Covariance. Le score de sélection d'un bloc incorpore non seulement le score d'attention, mais aussi la norme de la matrice de covariance du bloc ( $M_j = \|H_j - \bar{H}\|_2$ ). Les blocs présentant une forte pertinence sémantique ou une forte hétérogénéité sont sélectionnés pour un calcul exact, tandis que les blocs plus lisses sont délégués au développement de Taylor.

Implémentation Matérielle

PISA est implémenté via un noyau fusionné qui entrelace le calcul exact et le balayage approximatif au sein du processus de softmax en ligne (style FlashAttention). Le pipeline comprend :

Pré-calcul : Calcul des centroïdes de blocs, des sommes de valeurs et de la statistique globale $\bar{H}$ .
Phase 1 (Exacte) : Chargement des blocs sélectionnés dans la SRAM pour l'attention exacte.
Phase 2 (Approximation) : Balayage des vecteurs agrégés des blocs pour les blocs non sélectionnés en utilisant l'expansion du zéroième ordre.
Phase 3 (Correction) : Injection du terme de correction de premier ordre global.

Cette conception garantit que la méthode est efficace en termes de mémoire et évite les goulots d'étranglement associés au chargement de matrices distinctes pour chaque bloc.

3. Principales Contributions

Nouvelle Attention Éparse par Morceaux : Un mécanisme sans réentraînement qui permet de couvrir toute l'étendue de l'attention avec une complexité sous-quadratique en unifiant le calcul épar l'exact et l'approximation par bloc.
Schéma d'Approximation d'Ordre Hybride : Une stratégie qui combine l'expansion du zéroième ordre par bloc avec une correction globale du premier ordre. Cela augmente considérablement la précision avec un coût de calcul négligeable par rapport à l'attention éparse standard.
Routage Sensible à la Covariance : Une dérivation issue de l'analyse d'erreur qui minimise la divergence d'approximation en sélectionnant les blocs sur la base des scores d'attention et des normes de covariance des blocs.
Performance de l'État de l'Art : Démontre une efficacité et une qualité supérieures à travers diverses tâches (génération d'images et de vidéos) sans nécessiter de réentraînement.

4. Résultats Expérimentaux

Les auteurs ont évalué PISA sur des modèles de génération de vidéo (Wan2.1-1.3B/14B, Hunyuan-Video) et des modèles de génération d'images (SD 3.5, FLUX.1).

Génération de Vidéo :
- Wan2.1-14B : A obtenu une accélération de 1,91× par rapport à l'attention dense (1468s $\to$ 818s) à 87,5 % de parcimonie, tout en maintenant ou en améliorant légèrement les métriques de qualité (VBench, PSNR, SSIM) par rapport à l'attention dense.
- Hunyuan-Video : A obtenu une accélération de 2,57× (1651s $\to$ 641s) avec une préservation constante de la qualité.
- Comparaison : A surpassé les méthodes éparses existantes comme SpargeAttn et SVG2 tant en vitesse qu'en qualité, particulièrement dans la préservation des détails fins et de la cohérence temporelle.
Génération d'Images :
- FLUX.1-dev : A atteint une accélération de 1,2× (8,32s $\to$ 6,87s) à 85 % de parcimonie. Notamment, PISA a obtenu des scores plus élevés sur les métriques de préférence humaine (ImageReward, HPSv2) et les métriques de similitude (PSNR, SSIM) par rapport à l'attention dense et à SpargeAttn.
- SD 3.5 : A démontré des accélérations significatives (ex: 1,4× sur SD 3.5-M) tout en maintenant des scores FID comparables ou meilleurs que l'attention dense.
Efficacité du Noyau :
- PISA a systématiquement surpassé FlashAttention-2/3 et SpargeAttn en latence à travers diverses longueurs de séquence (4K à 32K) et densités.
- Contrairement à SpargeAtt even, qui se dégrade sur les séquences plus courtes, PISA a maintenu ses avantages de vitesse même à 4K jetons.

5. Signification et Revendications

L'article affirme que PISA représente un changement de paradigme, passant du « garder ou supprimer » à l'« exact ou approximatif » pour l'attention éparse. Sa principale signification réside dans :

Combler l'écart Vitesse-Qualité : Il comble efficacement l'écart entre la vitesse d'inférence et la qualité de génération, en atteignant une approximation quasi sans perte avec un surcoût négligeable.
Compatibilité Sans Réentraînement : En intégrant les approximations directement dans le cadre de normalisation du softmax, PISA permet la réutilisation directe des poids pré-entraînés, éliminant le besoin de réentraînement ou d'ajustement fin coûteux requis par les méthodes d'attention linéaire ou hybride.
Évolutivité : Il fournit une solution unifiée pour la génération d'images à haute résolution et de vidéos de longue durée, répondant au goulot d'étranglement de la complexité quadratique sans sacrifier la fidélité du contenu généré.

Les auteurs soulignent que leur méthode est la première à intégrer nativement le développement de Taylor par bloc et l'attention éparse dans un cadre unifié qui préserve la distribution intrinsèque des modèles pré-entraînés.