KroQuant : Quantification post-entraînement efficace pour les Diffusion Transformers via structure de produit de Kronecker

Contexte

Les Diffusion Transformers (DiT) constituent aujourd'hui l'architecture de référence pour la génération d'images haute fidélité, mais leur complexité computationnelle et leur empreinte mémoire élevée freinent considérablement leur déploiement sur des appareils edge ou mobiles. Bien que la quantification post-entraînement (PTQ) soit une voie prometteuse pour réduire ces coûts, l'application de formats à très faible précision, tels que W4A4 (4 bits pour les poids et 4 bits pour les activations), pose des défis techniques majeurs. Le problème central réside dans les valeurs aberrantes présentes dans les activations entrant dans les couches linéaires. Ces valeurs extrêmes dépassent la capacité de représentation des formats 4 bits, provoquant une amplification drastique des erreurs de quantification et entraînant la génération d'images présentant des artefacts visuels sévères et une perte de détails critiques.

Les solutions théoriques consistent souvent à appliquer des transformations linéaires réversibles aux activations et leurs inverses aux poids pour lisser les distributions. Cependant, la présence de couches de normalisation dans les architectures DiT impose que ces transformations soient exécutées en ligne à chaque étape de débruitage. Cette contrainte fait du coût computationnel lors de l'inférence un facteur limitant strict. Les approches actuelles sont confrontées à un dilemme : SmoothQuant, bien que rapide, sacrifie la précision en altérant les distributions d'amplitude ; les transformations de Hadamard fixes offrent une meilleure précision mais engagent des coûts de calcul élevés ; enfin, les transformations de rang complet, bien que précises, introduisent une surcharge de multiplication matricielle incompatible avec le temps réel. Cette tension entre efficacité et précision constitue le contexte principal de cette recherche.

Analyse approfondie

KroQuant propose une innovation majeure en introduisant une transformation réversible basée sur une structure de produit de Kronecker, spécifiquement optimisée pour les caractéristiques locales par bloc des activations. La méthode divise les activations en blocs indépendants de 32 éléments, appliquant une transformation apprise à chacun. Cette conception structurelle offre un avantage significatif en termes de stockage, utilisant moins de paramètres que les méthodes de mise à l'échelle par canal traditionnelles comme SmoothQuant, réduisant ainsi la charge globale du modèle. En exploitant une structure locale par bloc, la transformation est convertie en opérations GEMM de petites tuiles sur les cœurs de tenseur, plutôt qu'en multiplications matricielles denses. Cela permet à KroQuant de tirer pleinement parti de l'accélération matérielle des GPU modernes.

Pour atténuer les erreurs résiduelles, notamment celles liées au biais de quantification des poids, KroQuant intègre une calibration de poids LoRaQ hors ligne. Cette technique affine les poids lors d'une phase préliminaire pour absorber les décalages de distribution causés par les transformations d'activation, garantissant que le modèle quantifié final préserve au maximum la puissance d'expression du modèle original. Les expérimentations ont été menées sur des modèles DiT majeurs, incluant PixArt-Σ, SANA et FLUX.1-schnell, en utilisant strictement le format W4A4 (MXFP4e2). Les résultats démontrent que KroQuant génère une qualité d'image nettement supérieure à celle de SVDQuant et LoRaQ, se rapprochant significativement de la qualité de référence en précision flottante (FP). Des études d'ablation confirment que la structure locale par bloc réduit drastiquement la complexité computationnelle avec une perte de précision négligeable par rapport aux transformations de rang complet.

Impact sur l'industrie

L'introduction de KroQuant ouvre de nouvelles voies techniques pour le déploiement efficace des modèles de diffusion. À mesure que les applications AIGC s'étendent des serveurs cloud vers les appareils mobiles et embarqués, la taille du modèle et la vitesse d'inférence deviennent des facteurs déterminants pour la viabilité des produits. KroQuant démontre qu'une quantification à très faible précision peut être réalisée sans sacrifier la précision visuelle, ouvrant la voie à l'adoption massive des architectures DiT sur du matériel grand public. Pour la communauté open source, cette méthode fournit un paradigme PTQ reproductible et efficace, encourageant une focalisation accrue sur l'optimisation structurelle dans la quantification à faible précision.

Sur le plan industriel, l'accélération de l'inférence de 14 % et les économies de stockage substantielles offertes par KroQuant se traduisent directement par une réduction des coûts des services cloud ou une prolongation de l'autonomie des appareils terminaux. La stratégie combinant calibration hors ligne et transformation d'activation offre également une référence pour la recherche future, suggérant que la découplisation de ces processus permet d'optimiser séparément les erreurs de quantification. Cette capacité à maintenir une haute fidélité visuelle tout en réduisant les exigences computationnelles rend KroQuant particulièrement précieux pour les industries nécessitant des capacités de génération en temps réel, telles que les médias interactifs et les pipelines de création de contenu automatisés.

De plus, le succès de KroQuant souligne l'importance d'une conception d'algorithmes consciente du matériel. En alignant la structure mathématique de la transformation de quantification avec les capacités physiques des cœurs de tenseur, la méthode atteint des gains de performance que les optimisations purement algorithmiques ne peuvent égaler. Cette synergie entre innovation logicielle et utilisation matérielle établit une nouvelle norme pour le déploiement efficace des modèles d'IA. Elle incite les fabricants de matériel à optimiser davantage les architectures de cœurs de tenseur pour les opérations de petites tuiles, potentiellement entraînant des avancées dans la conception des puces AI et l'efficacité algorithmique.

Perspectives

Les principes sous-jacents à KroQuant sont susceptibles d'influencer le développement des prochaines générations de modèles d'IA générative efficaces. L'approche consistant à utiliser des transformations structurées apprises pour gérer les valeurs aberrantes d'activation peut être étendue à d'autres types d'architectures Transformer, y compris celles utilisées dans le traitement du langage naturel et les tâches multimodales. À mesure que la demande pour des modèles plus complexes et volumineux augmente, le besoin de techniques de quantification efficaces qui ne compromettent pas la qualité de sortie deviendra encore plus critique. La capacité de KroQuant à atteindre une haute précision avec une surcharge computationnelle faible le positionne comme un composant potentiel standard dans les futures boîtes à outils d'IA.

La recherche future pourrait explorer l'intégration de KroQuant avec d'autres techniques de compression, telles que l'élagage ou la distillation de connaissances, pour atteindre une efficacité encore supérieure. De plus, la stratégie de calibration hors ligne pourrait être affinée pour s'adapter dynamiquement à différentes distributions de données, renforçant ainsi la robustesse dans divers scénarios d'application. Le gain de vitesse de 14 % observé sur les GPU MI350 suggère que des gains similaires pourraient être obtenus sur d'autres architectures GPU modernes, à condition que les opérations GEMM de petites tuiles soient bien optimisées. En fin de compte, KroQuant représente une étape significative vers la démocratisation de la génération d'images IA de haute qualité, équilibrant précision, vitesse et utilisation des ressources pour un écosystème IA plus inclusif et durable.