Résumé Technique : Allocation de Pas Hétérogène (HSA) pour une Génération Vidéo Efficace

1. Énoncé du Problème

Les Transformers de Diffusion (DiT) ont établi des performances de pointe dans la génération vidéo, mais souffrent de coûts computationnels immenses. Les protocoles d'inférence standards appliquent un nombre uniforme de pas de débruitage ($T$) à chaque token spatio-temporel de la séquence. Cette approche monolithique ignore l'asymétrie inhérente des données visuelles : les régions d'arrière-plan, les patches statiques et les structures grossières sont perceptuellement plus simples et nécessitent moins de pas de débruitage que les objets d'avant-plan riches en détails ou les régions à fort mouvement. De plus, la vision humaine est moins sensible à la dégradation de la qualité dans les composants hautement dynamiques par rapport aux composants statiques. Par conséquent, appliquer un budget de débruitage complet uniformément à tous les tokens est un gaspillage computationnel et ne s'aligne pas avec les priorités perceptuelles.

Les méthodes d'accélération existantes, telles que la mise en cache des caractéristiques au niveau des pas (par exemple, TeaCache), prennent généralement des décisions globales : à n'importe quelle itération donnée, soit tous les tokens sont calculés, soit tous les tokens réutilisent des sorties mises en cache. Ces méthodes n'exploitent pas l'hétérogénéité spatiale ou au niveau des tokens, traitant la séquence comme un bloc monolithique plutôt que comme une collection de tokens aux exigences dynamiques variables.

2. Méthodologie : Allocation de Pas Hétérogène (HSA)

Les auteurs proposent l'Allocation de Pas Hétérogène (HSA), un algorithme d'inférence sans entraînement et agnostique au modèle qui assigne des trajectoires de débruitage spécifiques aux tokens au sein d'un DiT pré-entraîné. L'idée centrale est de partitionner la séquence de tokens en groupes basés sur leurs dynamiques de vitesse, allouant moins de pas aux tokens à évolution lente tout en maintenant un calendrier complet pour un groupe de référence.

2.1 Partitionnement et Planification des Tokens

L'ensemble de $N$ tokens est partitionné en $K$ groupes disjoints ($G_1, \dots, G_K$). Chaque groupe $G_k$ se voit attribuer un budget de pas $s_k$, qui doit être un diviseur du nombre total de pas $T$.

• Groupe de Référence ($G_1$) : Se voit attribuer le budget complet $s_1 = T$. Ces tokens sont actifs à chaque itération.
• Groupes Réduits : Se voient attribuer des budgets plus faibles (par exemple, $T/2, T/4$). Ces tokens ne sont actifs que lorsque l'itération globale $i$ satisfait $i \pmod{T/s_k} = 0$.
• Sélection Dynamique : L'article évalue quatre stratégies pour assigner les tokens aux groupes. La stratégie de Sélection Dynamique des Tokens, qui assigne les tokens en fonction du changement relatif L1 par token des prédictions de vitesse (identifiant les régions à évolution lente comme les arrière-plans), produit systématiquement les meilleures performances.

2.2 Synchronisation du Cache KV

Un défi majeur de la planification hétérogène est que les tokens avec des nombres de pas différents existent à différents niveaux de bruit ($\sigma$) à n'importe quelle itération donnée en temps réel. L'attention par auto-attention exige que tous les tokens s'attendent les uns aux autres, créant un décalage de synchronisation. HSA résout cela via un mécanisme de synchronisation du cache KV :

1. Tokens Actifs : Seul le sous-ensemble de tokens actifs ($A_i$) effectue un passage avant à travers les projections QKV.
2. Mise à Jour du Cache : Les projections de Clé ($K$) et de Valeur ($V$) fraîches pour les tokens actifs sont écrites dans un cache par couche, écrasant les entrées obsolètes.
3. Attention Contexte Complet : Les tokens actifs calculent des requêtes par rapport aux matrices $K, V$ mises en cache complètes (contenant à la fois les entrées fraîches pour les tokens actifs et les entrées obsolètes mais valides pour les tokens sautés).
4. Réduction de Complexité : Cela réduit la complexité de l'auto-attention de $O(N^2)$ à $O(|A_i| \cdot N)$, où $|A_i|$ est le nombre de tokens actifs à l'étape $i$. Les tokens sautés n'encourent aucun calcul de requête, attention croisée ou opérations feed-forward pendant les itérations inactives.

2.3 Mise à Jour Euler Mise en Cache

Pour faire avancer les états latents des tokens sautés sans évaluations de modèle supplémentaires, HSA emploie une mise à jour Euler mise en cache.

• Chaque token stocke la vitesse ($\hat{v}_n$) prédite à son dernier pas actif.
• À chaque itération globale, tous les tokens (actifs et sautés) sont avancés par la même étape Euler incrémentielle $(\sigma_{i+1} - \sigma_i)$.
• Les tokens actifs utilisent leur vitesse fraîchement calculée ; les tokens sautés réutilisent leur vitesse mise en cache.
• Cette mise à jour est formulée comme une seule opération tensorielle sur tous les $N$ tokens, évitant les branchements et maintenant l'efficacité du GPU.

2.4 Fenêtre de Mise en Cache Sensible à la Phase

Reconnaissant que les étapes précoces et tardives du débruitage sont critiques respectivement pour la structure globale et les détails fins, HSA restreint la mise en cache à la portion médiane de la trajectoire. Les $m$ premiers et derniers pas exécutent des passages avant complets et non mis en cache pour tous les tokens, garantissant que l'intégrité structurelle et le raffinement des détails ne sont pas compromis par des approximations obsolètes.

3. Contributions Clés

1. Algorithme HSA : Un paradigme d'inférence entièrement sans entraînement, plug-and-play, qui assigne des trajectoires de débruitage hétérogènes aux groupes de tokens spatio-temporels basés sur les dynamiques de vitesse.
2. Synchronisation du Cache KV : Un mécanisme qui préserve le contexte d'attention complet à $N$ tokens pour les tokens actifs tout en contournant entièrement les tokens inactifs, réduisant le coût d'attention par itération à $O(|A_i| \cdot N)$.
3. Mise à Jour Euler Mise en Cache : Une simple opération tensorielle qui fait avancer les tokens sautés en utilisant des vitesses mises en cache, maintenant la qualité de génération sans évaluations de modèle supplémentaires.
4. Évaluation Complète : Une analyse de quatre stratégies de regroupement de tokens et de quatre préréglages de budget, démontrant une robustesse à travers différents budgets d'inférence.

4. Résultats Expérimentaux

Les auteurs ont évalué HSA sur Wan-2.1-1.3B et Wan-2.2-A14B (Texte-vers-Vidéo et Image-vers-Vidéo) et LTX-2 (Audio-Vidéo).

• Compromis Qualité-Temps d'Exécution : HSA atteint une frontière de Pareto supérieure par rapport au Flow Matching (FM) uniforme et aux méthodes de mise en cache de pointe comme TeaCache.
• Accélération Agressive : Les avantages sont les plus prononcés dans les régimes d'accélération agressive (50 % et 25 % du temps d'exécution). À 25 % du temps d'exécution, HSA surpasse significativement les références sur le benchmark VBench (Score Total : 79,87 % contre 75,68 % pour FM et 78,33 % pour TeaCache).
• Robustesse Dimensionnelle : Alors que les méthodes de référence souffrent d'un « effondrement dimensionnel catastrophique » (dégradation sévère dans des dimensions spécifiques de VBench comme la fluidité du mouvement ou la cohérence temporelle) sous des budgets serrés, HSA suit de manière robuste la référence à budget complet sur les 16 dimensions d'évaluation.
• Image-vers-Vidéo (I2V) : HSA préserve avec succès l'alignement de la condition d'image et la fidélité visuelle dans les tâches I2V, tandis que les références FM uniformes à pas réduits ($T=10$) présentent un effondrement structurel sévère.
• Pas de Profilage Hors Ligne : Contrairement à TeaCache, qui nécessite un profilage hors ligne coûteux pour estimer les fluctuations de la sortie du modèle, HSA fonctionne sans aucun réglage ou fine-tuning hors ligne.

5. Importance et Revendications

L'article revendique que HSA adresse une inefficacité fondamentale dans l'inférence DiT en alignant les ressources computationnelles avec la complexité perceptuelle et dynamique des données visuelles. En traitant les tokens individuellement plutôt que globalement, HSA réalise des accélérations significatives sans la dégradation de qualité généralement associée à l'accélération agressive.

Les auteurs soulignent que HSA est un cadre flexible qui peut être instancié avec diverses stratégies de regroupement de tokens. Bien que la sélection dynamique produise actuellement les résultats les plus solides, le cadre anticipe des améliorations futures avec des stratégies de regroupement plus sophistiquées. Ce travail pose les bases d'une nouvelle classe d'algorithmes d'inférence qui exploitent l'hétérogénéité inhérente des données spatio-temporelles pour démocratiser la génération vidéo haute fidélité.