🤖 AI

LESA: Learnable Stage-Aware Predictors for Diffusion Model Acceleration

Le papier présente LESA, un cadre de prédicteurs apprenables et conscients des étapes basé sur un entraînement en deux phases et des réseaux KAN, qui accélère significativement les modèles de diffusion tout en préservant la qualité de génération pour les tâches d'image et de vidéo.

Auteurs originaux : Peiliang Cai, Jiacheng Liu, Haowen Xu, Xinyu Wang, Chang Zou, Linfeng Zhang

Publié 2026-03-17

📖 5 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Peiliang Cai, Jiacheng Liu, Haowen Xu, Xinyu Wang, Chang Zou, Linfeng Zhang

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 êtes un chef cuisinier de génie (le modèle de diffusion) qui doit préparer un plat complexe, étape par étape. Pour obtenir le résultat final parfait, vous devez suivre une recette très précise qui demande de goûter, d'ajuster et de remuer le mélange à chaque seconde pendant une heure. C'est ce que font les modèles d'IA actuels pour créer des images ou des vidéos : ils font des centaines de micro-ajustements pour transformer du "bruit" (une image floue et grise) en une photo magnifique.

Le problème ? C'est extrêmement lent. Votre chef cuisinier passe trop de temps à faire les mêmes gestes inutiles.

Voici comment LESA (le sujet de l'article) résout ce problème, expliqué simplement :

1. Le problème : "On ne peut pas toujours deviner"

Les méthodes actuelles pour accélérer ce processus fonctionnent comme un assistant un peu paresseux :

Méthode A (Réutilisation) : L'assistant dit : "Le mélange a l'air stable, je vais juste copier ce que j'ai fait il y a 2 secondes." Parfois, ça marche, mais souvent, le plat change de texture et le résultat est raté.
Méthode B (Prédiction simple) : L'assistant dit : "Le mélange change doucement, je vais deviner la prochaine étape avec une règle mathématique simple." Ça va bien au début, mais dès que le plat commence à cuire (les étapes finales), la règle simple ne suffit plus et le plat brûle.

Le problème, c'est que la "cuisson" d'une image n'est pas linéaire. Elle a trois phases très différentes :

Le début (Haute température) : Tout change vite et de manière chaotique. Il faut être très attentif.
Le milieu (Cuisson stable) : Les choses évoluent doucement et régulièrement. On peut aller plus vite.
La fin (Finition) : On ajoute les détails fins (les yeux, les textures). Il faut de nouveau être très précis.

Les anciennes méthodes utilisaient la même règle pour les trois phases, ce qui était inefficace.

2. La solution LESA : "L'Assistant qui a trois cerveaux"

LESA (Learnable Stage-Aware Predictors) est comme un chef qui a embauché trois assistants spécialisés, chacun avec un cerveau différent, pour gérer une phase précise de la recette.

L'Assistant "Chaos" (Début) : Il est expert dans le mouvement rapide. Il sait que tout change vite au début, donc il ne se repose pas sur des règles simples. Il apprend à anticiper les changements brusques.
L'Assistant "Calme" (Milieu) : Il est expert dans la régularité. Il sait que le milieu est stable, donc il peut prédire l'avenir très facilement et faire des bonds de géant dans le temps.
L'Assistant "Détail" (Fin) : Il est un chirurgien. Il sait que la fin demande de la précision, donc il ralentit pour ne pas gâcher les détails.

3. Le secret : Le "Cerveau KAN" (Le cerveau qui apprend)

Ce qui rend LESA génial, c'est le cerveau qu'il utilise pour faire ces prédictions. Au lieu d'utiliser un cerveau rigide (comme les méthodes actuelles), LESA utilise un réseau KAN.

Imaginez un cerveau qui peut changer de forme selon ce qu'il doit faire.

Si le problème est simple, il devient une règle droite.
Si le problème est complexe et courbe, il se courbe pour s'adapter parfaitement.

Contrairement aux assistants classiques qui ont une "recette fixe" (une fonction mathématique rigide), le cerveau KAN apprend la forme exacte de la recette en regardant des milliers d'exemples. Il devient un expert en "prédiction de l'avenir" pour chaque phase spécifique.

4. L'entraînement : "Répéter avant de jouer"

Pour que ces assistants soient parfaits, LESA suit un entraînement en deux temps :

Apprentissage guidé : On leur montre la "vraie" réponse (la photo finale) et on leur dit : "Regardez, c'est ce qu'il faut faire." Ils apprennent la théorie.
Entraînement en boucle : On les met dans une situation réelle où ils doivent prédire l'étape suivante, puis utiliser leur propre prédiction pour deviner la suivante, et ainsi de suite. C'est comme un répétition de théâtre où ils apprennent à ne pas faire d'erreurs même s'ils commencent avec une petite imperfection.

Le Résultat ?

Grâce à cette méthode, LESA peut sauter des étapes inutiles sans gâcher le plat.

Sur des images, il est 6 fois plus rapide que la normale, avec une qualité quasi identique.
Sur des vidéos, il est 5 fois plus rapide, en gardant le mouvement fluide et les couleurs justes.

En résumé : LESA ne force pas le modèle à aller plus vite en courant (ce qui fait trébucher). Il lui donne un GPS intelligent qui sait exactement quand on peut accélérer (au milieu de la route) et quand il faut ralentir pour prendre un virage serré (au début et à la fin), le tout piloté par un cerveau capable d'apprendre la forme exacte de la route.

1. Problématique

Les modèles de diffusion, en particulier les Diffusion Transformers (DiTs), ont révolutionné la génération d'images et de vidéos. Cependant, leur déploiement pratique est entravé par des coûts computationnels élevés dus à la nécessité d'itérer de nombreuses étapes de débruitage.

Les méthodes d'accélération actuelles reposent principalement sur le mise en cache et la réutilisation des caractéristiques (Feature Caching). On distingue deux approches :

Cache-Then-Reuse : Réutilise directement les caractéristiques de l'étape précédente.
Cache-Then-Forecast : Prédit les futures caractéristiques en supposant une évolution temporelle lisse (souvent via des méthodes sans entraînement comme l'expansion de Taylor).

Le problème central identifié par les auteurs est que ces hypothèses sont incorrectes. L'évolution des caractéristiques dans les modèles de diffusion n'est pas uniforme : elle varie considérablement selon les étapes (dynamique dépendante de l'étape).

Phase à haut bruit : Changements brusques et instables.
Phase intermédiaire : Évolution stable et continue.
Phase à faible bruit : Raffinement fin des détails.

Les méthodes existantes, qui appliquent une stratégie de prédiction fixe à toutes les étapes, échouent à s'adapter à ces dynamiques complexes, entraînant une dégradation de la qualité de l'image générée et une incohérence temporelle.

2. Méthodologie : Le Framework LESA

Pour résoudre ce problème, les auteurs proposent LESA (Learnable Stage-Aware), un cadre basé sur l'entraînement qui modélise explicitement les dynamiques temporelles du processus de diffusion.

A. Architecture Multi-Experts et Consciente des Étapes

LESA divise le processus de débruitage en trois étapes distinctes basées sur les niveaux de bruit, attribuant un prédicteur spécialisé ("expert") à chaque phase :

Expert Haut Bruit : Gère la phase initiale de changements drastiques.
Expert Bruit Moyen : Gère la phase de débruitage stable et continue.
Expert Faible Bruit : Se concentre sur le raffinement final des détails.

Chaque expert utilise une fenêtre temporelle adaptée à sa phase (plus courte pour les changements rapides, plus longue pour les phases stables).

B. Utilisation des Réseaux KAN (Kolmogorov-Arnold Networks)

Au cœur de chaque prédicteur se trouve un réseau KAN. Contrairement aux réseaux de neurones classiques (MLP) qui utilisent des fonctions d'activation fixes, les KAN apprennent des fonctions univariées (splines) à chaque nœud.

Fonctionnement : Le prédicteur prend les caractéristiques des $K$ étapes précédentes et l'information temporelle (décalage temporel) en entrée.
Modulation : Un module KAN calcule un facteur de modulation scalaire ( $\alpha$ ) qui ajuste la force de la mise à jour résiduelle. Cela permet au modèle d'apprendre des relations temporelles non linéaires complexes et d'adapter la prédiction à la dynamique spécifique de chaque étape.

C. Stratégie d'Entraînement en Deux Étapes

Pour garantir la robustesse et éviter la dérive de prédiction (prediction drift) lors de l'inférence accélérée, LESA utilise une stratégie d'entraînement spécifique :

Entraînement guidé par la vérité terrain (GT-Guided Training) : Le prédicteur apprend à reproduire les caractéristiques réelles du modèle de diffusion en utilisant les sorties du modèle comme référence (Ground Truth).
Entraînement autorégressif en boucle fermée (Closed-Loop AR Training) : Le prédicteur est entraîné en utilisant ses propres prédictions passées comme entrée pour les étapes suivantes. Cela simule les conditions réelles d'inférence et apprend au modèle à corriger les erreurs cumulées et à rester robuste face à la dérive.

3. Contributions Clés

Cadre de prédiction temporelle apprenable : Introduction d'un prédicteur basé sur les KAN qui apprend directement les dynamiques temporelles à partir des activations du modèle, surpassant les méthodes de prédiction sans entraînement.
Architecture Multi-Experts Consciente des Étapes : Une conception qui segmente le processus de débruitage et attribue des experts spécialisés, garantissant une modélisation précise des dynamiques uniques à chaque niveau de bruit.
Stratégie d'Entraînement Robuste : Une méthode en deux étapes (GT-guidé puis autorégressif) qui assure la stabilité de la génération même à des taux d'accélération élevés.

4. Résultats Expérimentaux

Les auteurs ont évalué LESA sur plusieurs modèles de pointe pour la génération d'images et de vidéos (FLUX.1-dev, Qwen-Image, HunyuanVideo).

Performances principales :

FLUX.1-dev : Accélération de 5,00× avec une dégradation de qualité minimale (seulement 1,0 % de baisse sur les métriques).
Qwen-Image : Accélération de 6,25× avec une amélioration de 20,2 % de la qualité par rapport à l'état de l'art précédent (TaylorSeer).
HunyuanVideo (Vidéo) : Accélération de 5,00× avec une amélioration de 24,7 % du PSNR par rapport à TaylorSeer.

Comparaison qualitative :
Contrairement aux méthodes concurrentes qui souffrent de distorsions structurelles, de perte de détails ou d'incohérences temporelles à haut taux d'accélération, LESA préserve la structure spatiale, la cohérence des couleurs et la fidélité des détails (ex: anatomie des baleines, texte sur des images, dynamique de l'encre dans l'eau).

5. Signification et Impact

Ce travail est significatif car il remet en question l'hypothèse commune selon laquelle les caractéristiques des modèles de diffusion évoluent de manière lisse et continue. En démontrant que la dynamique est dépendante de l'étape, LESA ouvre une nouvelle voie pour l'accélération des modèles de diffusion.

Efficacité : Permet un déploiement plus rapide et moins coûteux des modèles DiTs massifs.
Qualité : Démontre qu'il est possible d'obtenir des accélérations extrêmes (jusqu'à 6x) sans sacrifier la fidélité visuelle, voire en l'améliorant par rapport aux méthodes de prédiction heuristiques.
Généralisation : La méthode s'avère efficace aussi bien pour la génération d'images que pour la vidéo, prouvant la robustesse de l'approche basée sur les KAN et l'architecture multi-experts.

En résumé, LESA représente une avancée majeure en passant d'une approche de cache statique ou heuristique à une approche dynamique, apprenable et adaptée au contexte, permettant d'exploiter pleinement le potentiel des modèles de diffusion modernes.