Résumé technique : Échantillonnage d'importance adaptatif pour l'RL quantifié (AIS)

1. Énoncé du problème

L'apprentissage par renforcement (RL) pour les grands modèles de langage (LLM) est actuellement limité par le coût computationnel de la génération de trajectoires par défilement (rollout). Pour y remédier, les pratiques industrielles récentes (par exemple, DeepSeek V3.2, veRL) utilisent une quantification de faible précision (par exemple, FP8) pour le moteur de défilement tout en maintenant un entraîneur de plus haute précision (BF16). Cela crée un décalage défilement-entraînement : la politique de défilement quantifiée ($\hat{\pi}_{\theta_{old}}$) génère des trajectoires qui diffèrent distributionnellement de la politique d'apprentissage pleine précision ($\pi_{\theta_{old}}$) utilisée pour calculer les gradients.

L'article identifie ce décalage comme non stationnaire et agissant comme une « épée à double tranchant » :

• Début de l'entraînement : Le bruit de quantification agit comme un bonus stochastique d'exploration, exposant le gradient à des trajectoires que l'entraîneur sous-échantillonnerait autrement.
• Fin de l'entraînement : À mesure que la politique se précise, ce même bruit devient une source déstabilisatrice de biais, provoquant un effondrement de l'entraînement sur des benchmarks de raisonnement (par exemple, une baisse de 6,63 % de la précision sur AIME25 pour Qwen3-8B).

L'échantillonnage d'importance (IS) standard tente de corriger ce biais mais échoue souvent en raison d'une variance prohibitivement élevée dans les séquences longues. À l'inverse, l'échantillonnage d'importance tronqué (TIS) stabilise l'entraînement mais utilise une force de correction fixe qui ne peut pas s'adapter à la nature évolutive du décalage, conduisant soit à une sur-correction (suppression de l'exploration utile), soit à une sous-correction (permettant l'instabilité).

2. Méthodologie : Échantillonnage d'importance adaptatif (AIS)

Les auteurs proposent l'échantillonnage d'importance adaptatif (AIS), un cadre qui ajuste dynamiquement la force de la correction d'échantillonnage d'importance sur une base par lot. Au lieu d'un seuil fixe, l'AIS calcule un coefficient de mélange unique, $\alpha(x) \in [0, 1]$, qui interpole entre le gradient non corrigé et le gradient entièrement pondéré par l'importance.

2.1 Le coefficient de mélange

Le cœur de l'AIS est un mécanisme de porte factorisé qui combine trois diagnostics en temps réel dérivés des statistiques déjà calculées lors de la mise à jour de la politique :

1. Fiabilité du poids ($\alpha_{ess}$) : Mesure le rapport de la taille d'échantillon effective (ESS) des poids d'importance tronqués. Si la distribution des poids est concentrée (faible ESS), la correction est jugée peu fiable et pondérée à la baisse.
   $$\alpha_{ess} = (1 + CV^2(\bar{w}))^{-1/2}$$
2. Sévérité de la divergence ($\alpha_{mis}$) : Mesure l'écart moyen absolu des probabilités logarithmiques entre les politiques de défilement et d'entraînement. La correction n'est activée que lorsque ce décalage dépasse un seuil de saturation $\delta$.
   $$\alpha_{mis} = \min\left(1, \frac{\bar{d}}{\delta}\right)$$
3. Amplification de la variance ($\alpha_{var}$) : Surveille si de grands poids d'importance coïncident avec des avantages de grande magnitude, ce qui gonflerait la variance du gradient au-delà de limites acceptables. Si le ratio d'amplification de la variance $\Delta\sigma$ dépasse une tolérance $\gamma$, une pénalité est appliquée.
   $$\alpha_{var} = \max\left(0, \frac{\Delta\sigma - \gamma}{\gamma}\right)$$

Le coefficient de mélange final est calculé comme suit :
$$\alpha = \text{clip}(\alpha_{ess} - \beta \cdot \alpha_{var}, 0, 1) \cdot \alpha_{mis}$$

2.2 Intégration avec GRPO

L'AIS est intégré à l'optimisation de politique relative de groupe (GRPO). La méthode introduit un poids de rectification $\tilde{w}_i$ qui ajuste le terme d'avantage dans l'objectif GRPO :
$$\tilde{w}_i = 1 + \alpha \cdot (\min(\rho_{rt}, C) - 1)$$
où $\rho_{rt}$ est le ratio défilement-entraînement et $C$ est le seuil de troncature. Cette formulation garantit que lorsque le décalage est négligeable ($\alpha \to 0$), l'objectif se réduit exactement au GRPO standard, maintenant la compatibilité ascendante.

3. Contributions clés

• Caractérisation de la non-stationnarité : L'article établit que le décalage défilement-entraînement n'est pas uniformément nuisible ; il fournit un bonus d'exploration au début de l'entraînement mais devient une source de biais plus tard. Cela remet en question la vision du décalage de quantification uniquement comme un adversaire.
• Cadre de correction adaptatif : L'AIS est la première méthode à moduler dynamiquement la force de la correction d'échantillonnage d'importance basée sur des diagnostics en temps réel (fiabilité, divergence et variance), plutôt que d'utiliser un seuil statique.
• Application aux modèles de diffusion : Ce travail présente la première étude de défilement de faible précision pour le RL sur des LLM basés sur la diffusion (spécifiquement LLaDA-8B-Instruct), démontrant que le problème du décalage et la solution AIS se généralisent au-delà des architectures autorégressives.
• Compromis efficacité-précision : La méthode atteint un équilibre où elle récupère la précision de l'entraînement BF16 pleine précision tout en conservant les avantages de débit et de mémoire des défilements FP8.

4. Résultats expérimentaux

Les auteurs ont évalué l'AIS sur Qwen3-8B, Qwen3.5-9B (autorégressif) et LLaDA-8B-Instruct (diffusion) à travers des benchmarks de raisonnement mathématique (GSM8K, MATH500, AIME25, AMC, Olympiades) et de planification (Sudoku, Countdown).

• Récupération de la précision :
  • Modèles autorégressifs : Sur Qwen3-8B, un défilement FP8 naïf a fait chuter la précision AIME25 de 36,70 % (BF16) à 26,81 %. L'AIS l'a récupérée à 43,33 %, surpassant la référence BF16. Des améliorations similaires ont été observées sur Qwen3.5-9B, où l'AIS a égalé ou dépassé le BF16 sur 11 des 12 paires de benchmarks.
  • Modèles de diffusion : Sur LLaDA-8B, le défilement FP8 a causé une dégradation significative (par exemple, -19,00 % sur Countdown à 512 tokens). L'AIS a récupéré les performances pour égaler ou dépasser les niveaux BF16 à travers diverses longueurs de séquence et régimes d'affinage (FT complet et LoRA).
• Capacités générales : L'AIS a préservé les performances sur les benchmarks généraux (MMLU, PIQA, ARC, etc.), indiquant aucun surapprentissage au signal de récompense RL.
• Gains d'efficacité :
  • Accélération : Le défilement FP8 a fourni une accélération de 1,5× à 2,76× du temps de défilement de bout en bout, selon la taille du lot et l'architecture du modèle (gains plus élevés observés dans les modèles de diffusion en raison de l'utilisation des cœurs tensoriels).
  • Mémoire : Le défilement FP8 a réduit de moitié l'empreinte du modèle de défilement (par exemple, ~7,95 Go contre ~15,9 Go pour LLaDA-8B), permettant des tailles de lot plus grandes qui étaient auparavant hors mémoire (OOM) sous BF16.

5. Signification et revendications

L'article revendique que l'AIS est la première étude de défilement de faible précision pour le RL sur des LLM de diffusion. Sa signification principale réside dans la résolution du « décalage défilement-entraînement » sans sacrifier les gains d'efficacité de la quantification.

Les auteurs affirment que :

1. Gestion plutôt qu'élimination : Le décalage doit être géré, non éradiqué. L'AIS supprime avec succès le biais déstabilisateur de l'entraînement de fin de phase tout en préservant les bénéfices exploratoires du bruit de quantification aux étapes précoces.
2. Évolutivité : La méthode est computationnellement efficace, ajoutant une surcharge négligeable car les diagnostics sont des réductions scalaires sur des statistiques existantes.
3. Robustesse : L'AIS se généralise à travers différentes familles de modèles (autorégressif vs diffusion) et paradigmes d'entraînement (FT complet vs LoRA).

L'article conclut que l'AIS permet un entraînement RL haute capacité et basse précision qui égale ou dépasse les références pleine précision, en faisant une solution viable pour mettre à l'échelle le RL dans des environnements aux ressources limitées.