Résumé Technique : STaR-Quant

Énoncé du Problème

Les modèles de langage à diffusion (DLLM - Diffusion Large Language Models) sont apparus comme une alternative prometteuse aux LLM autorégressifs, utilisant un débruitage itératif masqué avec un contexte bidirectionnel pour générer du texte. Cependant, leur déploiement est entravé par des surcharges substantielles en mémoire et en calcul découlant de la taille importante des modèles et du processus de débruitage itératif à plusieurs étapes. Bien que la quantification post-entraînement (PTQ - Post-Training Quantization) ait réussi à compresser les LLM autorégressifs, l'application directe des méthodes PTQ existantes aux DLLM entraîne une dégradation significative des performances, particulièrement sous une quantification agressive des poids et des activations en bas nombre de bits.

L'article identifie deux défis spécifiques propres aux DLLM qui causent cet échec :

1. Disparité d'Activation Dépendante de l'État : Au sein d'une seule étape de débruitage, la séquence contient à la fois des jetons masqués (positions incertaines à prédire) et des jetons non masqués (contexte visible). Ces deux états présentent des distributions d'activation et des motifs d'anomalies (outliers) distincts. L'application d'une transformation d'activation ou d'un espace de quantification unifié aux deux états ne parvient pas à lisser efficacement les distributions pour l'un ou l'autre.
2. Accumulation d'Erreurs Temporelles : La génération des DLLM procède par étapes itératives où la sortie d'une étape sert d'entrée pour la suivante. Les erreurs de quantification, particulièrement dans la multiplication matricielle entre la carte d'attention softmax et les états de valeur (value states), ne sont pas isolées à une seule étape mais s'accumulent tout au long de la trajectoire de débruitage, entraînant un écart significatif par rapport au chemin en pleine précision.

Méthodologie : STaR-Quant

Pour répondre à ces défis, les auteurs proposent STaR-Quant, un cadre de PTQ conçu spécifiquement pour les DLLM, basé sur la cohérence état-temps. Le cadre opère sur deux dimensions :

1. Transformation d'Activation Guidée par l'État (SGAT - State-Guided Activation Transformation)

SGAT traite la disparité d'activation entre les jetons masqués et non masqués.

• Mécanisme : Au lieu d'une transformation unique pour tous les jetons, SGAT décompose la dimension cachée en sous-espaces partagés et spécifiques à l'état.
  • Sous-espace Partagé ($U_{sh}$) : Capture les composantes d'activation communes.
  • Sous-espaces Spécifiques à l'État ($U_m, U_u$) : Absorbent les motifs d'anomalies spécifiques aux jetons masqués et non masqués, respectivement.
• Implémentation : Des portes de sous-espace binaires dirigent les jetons masqués et non masqués vers leurs espaces de transformation respectifs ($P_m$ et $P_u$). Crucialement, le côté des poids maintient une transformation statique unifiée ($P_e$) appliquée une seule fois aux poids. Cela évite de dupliquer les poids transformés lors de l'inférence tout en permettant au côté activation de s'adapter à l'état du jeton.
• Optimisation : Les paramètres de transformation sont appris en minimisant l'erreur de quantification par couche entre la sortie en pleine précision et la projection guidée par l'état quantifiée.

2. Compensation d'Attention Temporelle (TAC - Temporal Attention Compensation)

TAC traite l'accumulation des erreurs de quantification à travers les étapes de débruitage itératives.

• Cible : La méthode cible la sortie de la multiplication softmax-valeur dans le mécanisme d'attention, une opération connue pour être sujette aux erreurs sous une quantification à bas nombre de bits.
• Mécanisme : TAC insère une application affine légère, de type bloc-diagonal, avant la projection de sortie de l'attention. Elle ne modifie pas le calcul de l'attention elle-même, mais corrige la représentation d'attention quantifiée ($\hat{Z}_t$) pour l'aligner avec son homologue en pleine précision ($Z^{FP}_t$).
• Estimation : Les paramètres de compensation (matrices par blocs $M_b$ et biais $d_b$) sont estimés en forme fermée en faisant correspondre les statistiques du premier et du second ordre (moyenne et covariance) des représentations en pleine précision et quantifiées. Un terme d'amortissement et un lissage des valeurs singulières sont appliqués pour assurer la stabilité compte tenu des données de calibration limitées.

Contributions Clés

1. Identification du Problème : Les auteurs identifient formellement la disparité d'activation dépendante de l'état et l'accumulation d'erreurs temporelles comme les principaux goulots d'étranglement pour la quantification des DLLM à bas nombre de bits.
2. Transformation d'Activation Guidée par l'État (SGAT) : Une nouvelle méthode qui projette les activations masquées et non masquées dans des espaces de transformation partagés et spécifiques à l'état, tout en maintenant une transformation unifiée statique du côté des poids, lissant efficacement les anomalies spécifiques à l'état sans duplication de poids.
3. Compensation d'Attention Temporelle (TAC) : Une méthode de compensation affine par blocs en forme fermée qui aligne les représentations d'attention quantifiées avec leurs homologues en pleine précision avant la projection de sortie, atténuant l'accumulation d'erreurs le long de la trajectoire de débruitage.
4. Validation Empirique : Des expériences approfondies démontrent que STaR-Quant surpasse systématiquement les bases de référence PTQ fortes (incluant RTN, AWQ, QuaRot et DLLMQuant) sur plusieurs DLLM et tâches.

Résultats Expérimentaux

Les auteurs ont évalué STaR-Quant sur trois DLLM représentatifs : LLaDA-8B, LLaDA-1.5-8B et Dream-7B, en utilisant une quantification des poids et des activations en 4 bits (W4A4).

• Performance : STaR-Quant a obtenu la meilleure performance moyenne sur 9 benchmarks (couvrant la connaissance générale, le raisonnement et la génération de code).
  • Sur LLaDA-8B, il a amélioré le score moyen de 2,78 points par rapport à la base la plus forte (DLLMQuant++).
  • Sur Dream-7B, il a amélioré le score moyen de 1,69 points.
  • Des gains notables ont été observés dans les tâches de raisonnement (GSM8K) et de génération de code (HumanEval), qui sont sensibles aux erreurs de représentation intermédiaire.
• Efficacité :
  • Accélération : STaR-Quant a atteint des accélérations d'inférence de 1,65× à 1,69× par rapport au déploiement FP16.
  • Mémoire : Il a réduit l'utilisation de la mémoire GPU de pointe de 3,05× à 3,14× (par exemple, réduisant la mémoire de LLaDA d'environ 15,9 Go à environ 5,2 Go).
• Études d'Ablation : Le retrait de TAC a provoqué une chute significative des performances (indiquant l'importance de la correction d'erreur temporelle), tandis que le retrait de SGAT a également dégradé les résultats, confirmant la nature complémentaire de la cohérence par état et par temps.

Signification et Revendications

L'article affirme que STaR-Quant fournit un changement nécessaire de stratégie de quantification pour les modèles de langage basés sur la diffusion. En modélisant explicitement les dimensions de l'état (masqué vs non masqué) et du temps (accumulation itérative), le cadre résout l'incompatibilité entre les hypothèses de la PTQ conventionnelle et le processus d'inférence de la diffusion.

Les auteurs positionnent STaR-Quant comme une solution pratique pour le déploiement des DLLM sur des appareils aux ressources limitées, offrant un équilibre entre des taux de compression élevés (W4A4), des économies de mémoire significatives et des accélérations, tout en maintenant des niveaux de précision proches des modèles en pleine précision. Le travail suggère que la "cohérence état-temps" est un principe efficace pour la quantification efficace des DLLM, bien que les auteurs notent qu'une validation supplémentaire sur des architectures diverses et des nombres de bits plus agressifs (par exemple, 2-bit ou 3-bit) reste une direction ouverte.