Résumé Technique : Comment Vous Commencez Détermine Comment Vous Raisonnez : Piloter l'Exploration en RLVR via des Priors Ajustés par Préfixe

1. Énoncé du Problème

L'Apprentissage par Renforcement avec Récompenses Vérifiables (RLVR) est devenu un paradigme puissant pour améliorer les capacités de raisonnement des Modèles de Langage de Grande Taille (LLM) dans des tâches telles que la résolution de problèmes mathématiques et la preuve de théorèmes. Cependant, les auteurs identifient une limitation critique dans les approches RLVR actuelles : l'effondrement de l'entropie.

Bien que le RLVR améliore souvent la précision sur un seul déroulement (Pass@1), il échoue fréquemment à élargir la couverture des trajectoires de raisonnement réussies. Cela se manifeste par une réduction de la diversité des sorties, où le modèle converge vers un ensemble restreint de chemins de raisonnement. Les techniques d'exploration existantes, telles que la régularisation de l'entropie ou les récompenses intrinsèques, traitent souvent la distribution du modèle de base comme fixe et tentent de modifier l'objectif RL ou le processus d'échantillonnage. L'article soutient que ces méthodes sont insuffisantes car le RL est fortement façonné par la distribution a priori du modèle de base. De plus, les techniques d'exploration passives dégradent souvent la qualité de la génération, entraînant des déroulements bruités qui ne contribuent pas à un raisonnement efficace.

Le défi central consiste à encourager une exploration efficace — découvrir des comportements de raisonnement diversifiés, pertinents pour la tâche et de haute qualité — sans que la politique ne s'effondre vers des trajectoires redondantes ou de faible qualité, et sans le coût computationnel de la mise à jour des paramètres complets du modèle.

2. Méthodologie : Cadre IMAX

Les auteurs proposent l'Exploration Augmentée Maximisant l'Information (IMAX), un cadre léger qui déplace la cible d'optimisation des paramètres du LLM vers le réglage de prompts souples. Au lieu de modifier le processus RL au-dessus d'un prior fixe, IMAX apprend un pool de préfixes souples qui remodelent activement le prior du modèle sur les trajectoires de raisonnement.

2.1 Raisonnement Conditionné par Préfixe

Le mécanisme central implique un pool de $C$ préfixes souples (tokens virtuels continus) $\mathcal{E} = \{e_1, \dots, e_C\}$. Pour un prompt d'entrée donné $x$, un préfixe $e_z$ est préfixé à la séquence d'embeddings d'entrée. Le LLM de base reste figé, et la politique $\pi_E(y|x, z)$ est conditionnée à la fois par le prompt et l'identité du préfixe sélectionné $z$.

• Initialisation : Les préfixes sont initialisés à partir d'un pool de prompts textuels sémantiquement significatifs (par exemple, « Introduire explicitement les variables », « Partitionner en cas ») générés par un LLM externe. Cela garantit que les préfixes commencent avec des stratégies de raisonnement distinctes.
• Contrôle : Chaque préfixe agit comme un « bouton de contrôle », induisant une distribution de déroulement distincte et orientant le modèle vers différentes régions de son espace de raisonnement (par exemple, solutions algébriques vs diagrammatiques vs de type code).

2.2 Récompense Maximisant l'Information (InfoMax)

Pour empêcher le pool de préfixes appris de s'effondrer vers des contrôles redondants (où différents préfixes induisent des trajectoires identiques), les auteurs introduisent une récompense InfoMax. L'objectif est de maximiser l'information mutuelle (MI) entre l'identité du préfixe $Z$ et la trajectoire générée $Y$, conditionnée par le prompt $X$ :
$$I(Z; Y | X) = H(Y | X) - H(Y | X, Z)$$
Intuitivement, cela encourage différents préfixes à produire des comportements de raisonnement valides et distinguables pour la même question.

2.3 Limite Inférieure Variationnelle et Entraînement

Le calcul direct de l'information mutuelle est intraitable. Les auteurs dérivent une limite inférieure variationnelle en utilisant une postérieure variationnelle $q_\phi(z | x, y)$, qui est un classifieur léger entraîné pour prédire quel préfixe a généré une réponse spécifique $y$ étant donné le prompt $x$.

• Forme Inverse : Ils utilisent la décomposition inverse de l'information mutuelle : $I(Z; Y | X) = H(Z | X) - H(Z | X, Y)$.
• Formulation de la Récompense : La récompense InfoMax est définie comme $\log q_\phi(z | x, y)$. L'objectif total pour le pool de préfixes devient :
  $$\max_E \mathbb{E} [R(x, y) + \beta \log q_\phi(z | x, y)]$$
  où $R(x, y)$ est la récompense vérifiable standard (exactitude) et $\beta$ est un hyperparamètre.

2.4 Optimisation Alternée

La procédure d'entraînement suit un schéma d'optimisation alternée (Inférence Variationnelle par Ascension de Coordonnées) :

1. Mise à jour de la Postérieure ($q_\phi$) : Le pool de préfixes étant fixe, mettre à jour le classifieur $q_\phi$ pour prédire l'identité du préfixe à partir de la paire prompt-réponse (générée sans le préfixe souple pour éviter la fuite d'information).
2. Mise à jour du Pool de Préfixes ($E$) : $q_\phi$ étant fixe, mettre à jour les préfixes souples en utilisant un algorithme d'Optimisation de Politique Relative de Groupe (GRPO) augmenté de la récompense InfoMax.
3. Échantillonnage Stratifié : Pendant l'entraînement, chaque prompt est apparié à des identités de préfixe distinctes pour garantir une couverture uniforme de l'espace des préfixes au sein de chaque groupe GRPO.

3. Contributions Clés

1. Paradigme RLVR Léger : L'article introduit une méthode pour améliorer le raisonnement en optimisant un pool de prompts souples plutôt que les paramètres complets du modèle. Cela remodelle la distribution de génération conditionnelle d'un modèle figé sans la surcharge computationnelle du fine-tuning complet.
2. Récompense InfoMax pour la Diversité : Les auteurs proposent une récompense InfoMax au niveau de la réponse qui encourage explicitement différents préfixes à induire des trajectoires de raisonnement distinguables et pertinentes pour la tâche, empêchant l'effondrement du pool de préfixes vers des contrôles redondants.
3. Validation Empirique : Des expériences extensives sur des benchmarks de raisonnement mathématique (GSM8K, MATH-500, Minerva, AMC) à travers trois échelles de modèles (Qwen2.5-1.5B, Qwen3-4B, Qwen3-8B) démontrent que IMAX améliore à la fois la qualité et la diversité du raisonnement.

4. Résultats Expérimentaux

Les auteurs ont évalué IMAX par rapport à plusieurs bases de référence, notamment le GRPO de politique complète, l'initialisation par prompt textuel, et d'autres méthodes RLVR orientées vers la diversité (DIVER, MERCI).

• Gains de Performance : IMAX a constamment surpassé les bases de référence RLVR standard ajustées par préfixe.
  • Sur Qwen3-4B, IMAX a amélioré le Pass@4 jusqu'à 11,60 % et l'Avg@4 de 10,57 % par rapport aux bases de référence standard sur des benchmarks spécifiques.
  • Sur Qwen3-8B, DAPO+IMAX a montré des améliorations significatives par rapport à DAPO-préfixe sur GSM8K, MATH-500 et Minerva.
• Diversité et Couverture :
  • Analyse t-SNE : Les visualisations des embeddings de réponse ont montré que les préfixes induits par IMAX ont résulté en des clusters plus séparés par rapport aux autres méthodes, indiquant une plus grande diversité intra-groupe.
  • Composantes de Raisonnement : Les cartes de chaleur ont révélé que différents préfixes ont induit avec succès des motifs de raisonnement distincts (par exemple, le Préfixe 1 favorisant la manipulation algébrique, le Préfixe 2 favorisant la partition des cas et la vérification), s'alignant sur leur initialisation sémantique.
  • Performance Oracle : La performance combinée de plusieurs préfixes (Oracle) a dépassé le meilleur préfixe unique, confirmant que les préfixes couvrent des régions de succès complémentaires de l'espace du problème.
• Efficacité : IMAX nécessite un coût d'entraînement significativement inférieur (3,0 à 6,0 fois moins d'heures GPU) par rapport au GRPO à paramètres complets, car il ne met à jour que les embeddings de préfixe légers et la tête de classifieur.

5. Signification et Revendications

L'article revendique que la manière dont un modèle commence un processus de raisonnement (le prior) dicte fondamentalement comment il raisonne. En traitant les prompts souples comme des variables de contrôle réutilisables, IMAX offre une solution évolutive au problème d'exploration en RLVR.

• Agnostique de l'Algorithme : Le cadre est conçu pour être intégré de manière transparente dans les pipelines RLVR existants (par exemple, GRPO, DAPO).
• Diversité Structurée : Contrairement aux méthodes qui augmentent simplement l'aléatoire, IMAX encourage une diversité de raisonnement structurée, garantissant que l'exploration mène à des chemins sémantiquement distincts et de haute qualité.
• Limites Modestes : Les auteurs reconnaissent que la méthode est principalement conçue pour les tâches de raisonnement, montrant des gains limités sur les tâches générales de suivi d'instructions. De plus, l'efficacité dépend de la qualité des prompts textuels initiaux, qui peuvent nécessiter une conception manuelle ou un réajustement à travers les domaines. Les préfixes appris introduisent également une surcharge persistante, bien que faible, lors de l'inférence.

En résumé, IMAX démontre que l'optimisation du « début » du processus de génération via des préfixes souples est une stratégie hautement efficace pour piloter l'exploration et améliorer les capacités de raisonnement des LLM, offrant une alternative computationnellement efficace à l'apprentissage par renforcement de politique complète.