Résumé Technique : Troncature de Noyau Adaptative pour le Raisonnement Long-Form

Énoncé du Problème

Dans le raisonnement des modèles de langage à long format (long-form), le budget de génération est souvent suffisamment large pour permettre un raisonnement multi-étapes, l'auto-affinement et l'utilisation d'outils. Dans ces régimes, l'échantillonneur de décodage n'est pas simplement un détail de post-traitement, mais un mécanisme de contrôle critique qui détermine les chemins de raisonnement accessibles et la stabilité des trajectoires d'entraînement. Les méthodes de troncature existantes — telles que le top-$p$, le min-$p$ et le top-$n_\sigma$ fixe — améliorent l'échantillonnage non restreint mais souffrent de limitations significatives :

• Seuils Fixes : Les méthodes comme le top-$p$ et le min-$p$ reposent sur des seuils de probabilité fixes qui ne peuvent pas s'adapter aux changements d'entropie, à la difficulté de la tâche, à l'étape d'entraînement ou au budget de génération.
• Distorsions de l'Espace de Probabilité : Les méthodes basées sur la probabilité opèrent après le softmax et l'échelle de température, héritant des distorsions de la normalisation des probabilités. Par conséquent, la même géométrie de logits sous-jacente peut produire des ensembles de candidats différents lorsque la température change.
• Manque d'Adaptabilité : Un paramètre de troncature unique et fixe (par exemple, un $n$ fixe dans le top-$n_\sigma$) ne peut pas servir simultanément des contextes décisifs nécessitant une élagage tranchant et des contextes ambigus nécessitant un support plus large. De plus, dans les phases d'apprentissage par renforcement (RL) instables, une troncature agressive peut entraîner une instabilité de l'entraînement (dépassement de l'entropie, de la divergence KL et des normes de gradient).

Méthodologie : Échantillonnage de Troncature de Noyau Adaptatif (ANTS)

Les auteurs proposent l'Échantillonnage de Troncature de Noyau Adaptatif (ANTS), qui étend le l'échantillonnage top-$n_\sigma$ d'une heuristique de décodage fixe vers un mécanisme de contrôle de déploiement (rollout) adaptatif et conscient du budget. ANTS opère en trois étapes :

1. Estimation du Support dans l'Espace des Logits

Contra-irement aux méthodes basées sur la probabilité, ANTS sélectionne l'ensemble des candidats dans l'espace des logits pré-température. Il définit un voisinage $N_t(n)$ autour du logit maximum $\ell_{t, \text{max}}$ basé sur l'écart-type par vocabulaire $\sigma(\ell_t)$ :
$$N_t(n) = \{ v \in V \mid \ell_{t,v} > \max_{u \in V} \ell_{t,u} - n \cdot \sigma(\ell_t) \}$$
Cette sélection est invariante à la mise à l'échelle positive des logits (mise à l'échelle de la température), garantissant que l'ensemble des candidats reste cohérent quel que soit le changement de température.

2. Seuillage Conditionné par l'Entropie

La largeur du voisinage de troncature, $n_t$, n'est pas fixe mais s'adapte en fonction de l'incertitude locale du modèle. Elle est calculée comme suit :
$$n_t = n_0 + \gamma \cdot H(p^{(0)}_t)$$
où $H(p^{(0)}_t)$ est l'entropie de la distribution à température unitaire. Les contextes à haute entropie (ambigus) conservent plus d'alternatives (plus grand $n_t$), tandis que les contextes à faible entropie (décisifs) sont élagués plus sévèrement.

3. Contrôleur de Thompson-Sampling en Ligne

Pour déterminer dynamiquement la force de troncature optimale $\gamma$, ANTS emploie un contrôleur de bandit de type Thompson-sampling sur un ensemble fini de bras :

• Bras Finis : Un ensemble de $K$ bras avec des valeurs de $\gamma$ espacées logarithmiquement ($\gamma_k = 10^{\eta_k}$).
• Bras de Repli (Fallback) : Un $K+1$-ième bras critique avec $\gamma_{K+1} = +\infty$, qui désactive effectivement la troncature ($N_t = V$), récupérant ainsi l'échantillonneur de base.
• Signal de Récompense : Le contrôleur utilise une récompense intrinsèque basée sur l'entropie de la distribution induite par chaque bras. Les bras qui préservent un support suffisant (entropie supérieure à la moyenne) reçoivent des mises à jour postérieures plus importantes.
• Mécanisme de Stabilité : Le bras de repli agit comme une "issue de secours sécurisée". Si la troncature devient dangereuse lors des déploiements de RL (provoquant des dérives de métriques comme la divergence KL ou les normes de gradient), le contrôleur peut apprendre à sélectionner le bras de repli pour stabiliser l'entraînement.

Principales Contributions

1. Échantillonneur Adaptatif dans l'Espace des Logits : Formulation d'ANTS comme un échantillonneur invariant à la température qui sépare la sélection de l'ensemble de candidats de la stochasticité intra-ensemble.
2. Contrôle Piloté par l'Entropie : Introduction d'un contrôleur Thompson conditionné par l'entropie qui adapte la force de troncature en ligne et inclut un bras de repli sans troncature pour prévenir l'instabilité de l'entraînement.
3. Évaluation de l'Échelle du Budget de Génération : Évaluation complète sur des budgets de génération de 8K, 16K et 32K, démontrant que les gains de performance augmentent généralement avec le budget.
4. Troncature Sensible au Budget : Caractérisation d'un mode de défaillance spécifique dans la génération de code à budget court, motivant la nécessité de politiques de troncature apprises conjointement avec le budget de génération et le type de tâche.

Résultats Expérimentaux

La méthode a été évaluée sur un modèle de raisonnement de type Mixture-of-Experts (MoE) creux avec 33B de paramètres totaux / 4B actifs.

• Échelle avec le Budget : La performance moyenne sur les benchmarks basés sur des pourcentages s'est améliorée de +1,9 points à 8K, +3,8 à 16K, et +5,2 à 32K.
• Respect des Instructions (IFBench) : A montré les gains les plus forts, s'améliorant de +10,5 (Loose) et +10,8 (Strict) points à 32K. ANTS a empêché la dérive vers des continuations verbeuses ou tangentielles qui violent les contraintes.
• Raisonnement Mathématique (AIME 2025) : S'est amélioré de +7,0 points à 32K, avec des gains ayant presque doublé de 8K à 32K.
• Génération de Code (Codeforces) : A révélé une interaction critique avec le budget. À 8K, ANTS accusait un retard par rapport à la ligne de base de -59 ELO malgré la génération de plus de tokens de solution. Cependant, à 16K et 32K, il a inversé cet écart, atteignant respectivement des gains de +230 et +212 ELO. Cela suggère que la troncature est préjudiciable dans les régimes de code à contraintes serrées, mais bénéfique lorsque l'espace des programmes réalisables est plus large.
• Questions Scientifiques (GPQA) & Connaissances (MMLU Pro) : Ont montré une parité quasi totale ou des gains modestes, indiquant que la troncature dans l'espace des logits offre moins de levier pour les tâches nécessitant une sélection de réponse unique précise par rapport à la génération long-form.
• Allocation de Tokens : ANTS a modifié les modèles d'utilisation des tokens. À des budgets modérés, il a augmenté les tokens de "pensée" (délibération). À 32K, il a réduit les tokens de pensée tout en maintenant une précision élevée, suggérant qu'il supprime les continuations inutiles plutôt que de simplement générer des chaînes plus longues.

Signification et Revendications

L'article soutient que la conception de l'échantillonneur ne doit pas être traitée simplement comme un hyperparamètre de décodage fixe, mais comme une composante intégrale du scaling de l'inférence et du contrôle de déploiement (rollout control).

• Stabilité : L'inclusion du bras de repli est centrale au succès de la méthode en RL, permettant au système de se remettre d'états de troncature instables dont les méthodes fixes ne peuvent s'échapper.
• Sensibilité au Contexte : Les résultats démontrent que la stratégie d'échantillonnage optimale dépend du contexte, variant selon le budget de génération, le type de tâche (ex: code vs maths) et l'état d'entraînement.
• Efficacité : ANTS atteint souvent les bandes de performance plus tôt que la ligne de base, améliorant l'efficacité de calcul effective dans les contextes de RL en réduisant la fréquence des déploiements de faible qualité ou instables.

Les auteurs concluent que bien qu'ANTS ne domine pas universellement tous les échantillonneurs existants sur toutes les tâches, il présente un schéma de scaling robuste par rapport au budget de génération, particulièrement pour le raisonnement long-form et le respect des instructions, où supprimer les queues de distribution bruitées sans supprimer les chemins de raisonnement utiles est crucial.