Résumé technique : Auto-cohérence par affinement marginal

Énoncé du problème

Les modèles de langage modernes résolvent de plus en plus de tâches complexes en générant des traces de raisonnement intermédiaires (par exemple, des dérivations, des esquisses de code, des ébauches de preuves) avant d'émettre une réponse finale. Bien que les méthodes d'inférence standard comme l'approche par chaîne de pensée (Chain-of-Thought) suscitent ce comportement, elles reposent souvent sur l'échantillonnage d'une seule trajectoire ou sur un vote majoritaire a posteriori (auto-cohérence) pour sélectionner la réponse finale.

Les auteurs identifient une limitation fondamentale dans les méthodes existantes d'affinement au moment de l'inférence, spécifiquement l'échantillonnage par puissance. L'échantillonnage par puissance affine la distribution sur les séquences générées complètes (la distribution conjointe de la trace de raisonnement $z$ et de la réponse $a$). L'article soutient que ce n'est pas le bon niveau d'abstraction pour le raisonnement :

1. Enchevêtrement des variables latentes : Une trace de raisonnement est un calcul latent. De nombreuses traces distinctes peuvent conduire à la même réponse correcte.
2. Compétition vs Soutien : L'affinement au niveau de la séquence traite les dérivations alternatives ou les implémentations d'une même solution comme des concurrents, concentrant la probabilité sur des transcripts complets spécifiques plutôt que sur la réponse elle-même.
3. Fragilité du vote : L'auto-cohérence standard (vote majoritaire) nécessite une correspondance exacte des chaînes de caractères. Dans des domaines structurés comme le code ou les preuves, une même solution sous-jacente peut se manifester sous de nombreuses formes de surface (noms de variables différents, flux de contrôle ou ordres de lemmes différents), ce qui provoque une fragmentation du soutien du vote à travers des variations valides.

Le problème central consiste à savoir comment exploiter plusieurs traces de raisonnement pour amplifier la probabilité de la réponse correcte sans exiger un accord exact sur les chaînes de caractères ni engendrer le coût computationnel de la recherche sur des trajectoires complètes.

Méthodologie : Affinement marginal

Les auteurs proposent l'affinement marginal, un objectif au moment de l'inférence qui déplace la cible de la distribution conjointe $\pi(z, a | x)$ vers la marge de réponse affinée $p_{\alpha}^{marg}(a | x)$.

Cadre théorique

1. Décomposition : La complétion du modèle $y$ est analysée en une trace de raisonnement $z$ et un segment de réponse finale $a$. Le modèle induit une marge de réponse naturelle $m(a | x) = \sum_z \pi(z, a | x)$.
2. Distribution cible : Au lieu d'affiner la distribution conjointe, la méthode définit une distribution de puissance sur la marge de réponse :
   $$p_{\alpha}^{marg}(a | x) = \frac{m(a | x)^\alpha}{\sum_{a'} m(a' | x)^\alpha}$$
   où $\alpha > 1$ est la force d'affinement. Lorsque $\alpha \to \infty$, cela converge vers la réponse ayant le soutien marginal maximal (la cible idéale d'auto-cohérence).
3. Représentation par puissance entière : Pour une force d'affinement entière $K$, les auteurs dérivent une décomposition exacte utilisant $K$ traces de raisonnement indépendantes. La distribution cible peut être exprimée comme une espérance sur $K$ traces échantillonnées indépendamment $z_{1:K}$ :
   $$p_{K}^{marg}(a | x) \propto \mathbb{E}_{z_{1:K}} \left[ \prod_{i=1}^K \pi(a | x, z_i) \right]$$
   Cette formulation relie l'auto-cohérence à l'échantillonnage à partir d'une distribution cible au niveau de la réponse.

Algorithme : Décodage autorégressif parallèle

L'échantillonnage direct à partir de la distribution dérivée est intraitable car il nécessite de sommer sur toutes les continuations futures possibles (le « normalisateur de continuation »). Les auteurs proposent un algorithme léger de décodage autorégressif parallèle (Algorithme 1) qui approxime cette cible :

1. Génération de traces : Générez $M = K \times S$ traces de raisonnement, regroupées en $S$ ensembles de $K$ traces chacun.
2. Approximation au niveau du token : Au lieu de calculer la probabilité de la séquence complète, l'algorithme décode la réponse token par token. À chaque étape $t$, il calcule une distribution pondérée sur le vocabulaire :
   $$q_t^{(S)}(v | x, a_{<t}) \propto \sum_{s=1}^S w_s(a_{<t}) \exp\left( \sum_{i=1}^K \log \pi(v | x, z_i^{(s)}, a_{<t}) \right)$$
   • Poids de préfixe ($w_s$) : Ces poids suivent la vraisemblance du préfixe de réponse déjà généré sous chaque groupe de $K$ traces. Ils sont mis à jour en ligne à mesure que les tokens sont échantillonnés.
   • Approximation locale : L'algorithme ignore le « normalisateur de continuation » intraitable (la probabilité des tokens futurs), le traitant comme un facteur commun qui s'annule lors de la normalisation locale.
3. Échantillonnage : L'algorithme échantillonne le token suivant à partir de cette distribution, met à jour les poids de préfixe et répète jusqu'à la terminaison de la réponse.

Cette approche permet au modèle de générer une réponse unique qui est « soutenue » par plusieurs chemins de raisonnement, effectuant efficacement une auto-cohérence douce pendant la génération plutôt qu'en tant qu'étape de vote a posteriori.

Contributions clés

1. Formulation de l'affinement marginal : L'article définit un objectif au moment de l'inférence qui traite les traces de raisonnement comme des variables latentes et affine la distribution de réponse induite, plutôt que la distribution conjointe des traces et des réponses.
2. Décomposition multi-traces : Les auteurs démontrent que pour des forces d'affinement entières, la marge de réponse affinée admet une décomposition en une espérance sur $K$ traces de raisonnement indépendantes, reliant théoriquement l'auto-cohérence à l'échantillonnage au niveau de la réponse.
3. Algorithme de décodage efficace : Ils proposent un algorithme autorégressif parallèle qui échantillonne approximativement cette distribution affinée. Il est purement autorégressif, ne nécessite aucun canoniseur externe et évite le coût quadratique de génération de tokens de l'échantillonnage par puissance au niveau de la séquence.
4. Validation empirique : La méthode est évaluée sur des benchmarks de raisonnement mathématique (MATH-500) et de génération de code (HumanEval+, LiveCodeBench).

Résultats

Les expériences utilisent des modèles de la famille Qwen (1,7B, 4B, 8B) sur trois benchmarks.

• Performance vs Échantillonnage par puissance : L'affinement marginal surpasse significativement l'échantillonnage par puissance standard sur les benchmarks de génération de code (HumanEval+ et LiveCodeBench). Sur MATH-500, il se comporte de manière comparable ou légèrement inférieure au vote majoritaire (qui est une base forte pour les mathématiques en raison des petits espaces de réponses et de la vérification symbolique).
• Efficacité : La méthode est plusieurs ordres de grandeur plus rapide que l'échantillonnage par puissance pour les générations longues. Pour $L=8192$, l'affinement marginal était 38 fois plus rapide que l'échantillonnage par puissance sur le benchmark HumanEval+ tout en atteignant une précision supérieure. L'échantillonnage par puissance nécessite une génération de tokens en $O(L^2)$ dû au rééchantillonnage, tandis que l'affinement marginal nécessite $O(L)$.
• Hyperparamètres :
  • Force d'affinement ($K$) : L'augmentation de $K$ améliore généralement les performances en concentrant la probabilité sur les réponses soutenues par plusieurs traces, bien que cela augmente le calcul.
  • Nombre d'échantillons ($S$) : L'augmentation de $S$ (le nombre de groupes de traces) améliore la qualité de l'approximation au niveau du token, produisant des gains constants, en particulier avec des budgets de génération plus faibles.
  • Allocation du budget : Sous un budget de calcul total fixe ($K \times S = \text{const}$), différentes allocations entre la force d'affinement et le nombre d'échantillons produisent des résultats comparables, suggérant un compromis où un affinement agressif et une meilleure qualité d'approximation se compensent largement.

Signification et affirmations

L'article affirme que l'auto-cohérence doit être comprise comme un principe d'inférence au niveau de la réponse plutôt que comme une heuristique de vote a posteriori.

• Changement de focalisation : En ciblant la marge de réponse, la méthode agrège correctement le soutien à travers des chemins de raisonnement divers (par exemple, différentes implémentations de code ou stratégies de preuve) que les méthodes au niveau de la séquence traitent comme des sorties distinctes et concurrentes.
• Évolutivité : L'algorithme proposé permet les avantages de l'auto-cohérence dans des domaines avec de grands espaces de sortie structurés (comme le code) où la correspondance exacte des chaînes échoue, sans le coût computationnel prohibitif de la recherche complète de trajectoire ou du rééchantillonnage.
• Praticité : La méthode est une version « douce » de l'auto-cohérence qui peut être implémentée avec des modèles autorégressifs standard et un décodage parallèle, en faisant une alternative pratique à l'échantillonnage par température et à l'échantillonnage par puissance pour les tâches de raisonnement.

Les auteurs concluent que le temps de calcul au moment du test est plus efficacement dépensé à affiner la croyance en la réponse induite par le raisonnement latent qu'à sélectionner le seul transcript complet le plus probable.