Résumé Technique : Mise à l'échelle de l'inférence par routage de branches locales pour des modèles de langage efficaces et entraînables

Énoncé du Problème

Les progrès récents dans le raisonnement des modèles de langage (LLM) démontrent que l'allocation d'un calcul supplémentaire lors de l'inférence (mise à l'échelle au moment du test) peut améliorer considérablement les performances. Cependant, les approches existantes font face à un compromis critique :

• Échantillonnage de la Chaîne de Pensée (CoT) : Bien qu'efficace, la CoT standard reste monothreadée, s'engageant sur un jeton avant de pouvoir observer les conséquences des alternatives.
• Méthodes de Recherche : Les approches telles que l'Arbre de Pensées (Tree of Thoughts) ou la recherche arborescente de Monte Carlo (MCTS) introduisent de la largeur en explorant des arbres de pensées intermédiaires. Cependant, celles-ci sont souvent à grain grossier, coûteuses en calcul et difficiles à optimiser comme une politique de décodage unique et entraînables de bout en bout.
• Méthodes de Jetons Mous (Soft-Token) : Des techniques comme le "Soft Thinking" ou le "Multiplex Thinking" offrent une largeur au niveau du jeton en fusionnant les plongements (embeddings) de candidats en mélanges continus. Bien qu'entraînables, ces méthodes brouillent l'identité discrète des branches candidates, ce qui peut occulter l'information spécifique aux concepts.

L'article postule que les modèles de langage pourraient bénéficier d'un mécanisme de mise à l'échelle au moment du test, léger et au niveau du jeton, permettant aux décisions d'être éclairées par les états cachés de candidats locaux futurs sans encourir le coût d'une recherche complète de solution ou l'ambiguïté des mélanges mous.

Méthodologie : Routage de Branches Locales (LBR)

Les auteurs proposent le Routage de Branches Locales (LBR), un cadre de décodage par branchement et routage au niveau du jeton, entraînable. Le LBR augmente la génération autoregressive standard par un arbre de regard vers l'avant (lookahead) local et tournant.

Pipeline de Décodage

À chaque étape de décodage $t$, étant donné un préfixe engagé $x_{<t}$, le LBR exécute un processus "élagage–décalage–croissance" (prune–shift–grow) en quatre étapes :

1. Croissance (Grow) : Au lieu d'échantillonner un seul jeton suivant, le LBR développe un arbre local $\mathcal{A}_t$ de largeur $K$ et de profondeur $L$ enraciné à $x_{<t}$. Chaque nœud échantillonné est transmis à travers le modèle de langage pour calculer son état caché $h_t(v)$. Cela crée un arbre d'états cachés locaux plutôt qu'un simple ensemble de chaînes candidates.
2. Routage (Route) : Un routeur léger $\rho_\phi$ observe l'arbre transmis et sélectionne un sous-arbre de profondeur 1 à engager. Le routeur utilise un mécanisme d'attention d'ensemble (set-attention) pour encoder chaque sous-arbre candidat (par exemple, un jeton pour $L=1$, ou un jeton plus ses continuations pour $L>1$) en un vecteur. Ces vecteurs sont comparés via une attention inter-sous-arbres pour produire un score relatif, permettant au routeur de prendre une décision contrastive parmi les futurs frères (siblings).
3. Élagage (Prune) : Le routeur sélectionne le meilleur sous-arbre (indice $k^*_t$). Tous les autres sous-arbres sont rejetés. Le jeton racine du sous-arbre sélectionné devient le jeton suivant engagé $x_t$.
4. Décalage et Régénération (Shift and Regrow) : Le sous-arbre sélectionné est décalé vers l'avant (sa racine est supprimée et ses descendants deviennent le nouveau préfixe pour l'étape suivante). Une nouvelle couche est développée pour restaurer la profondeur de l'arbre $L$, et le cycle se répète.

Objectif d'Entraînement

Le LBR définit une vraisemblance de trajectoire d'arbre traçable qui permet un apprentissage par renforcement de bout en bout avec des récompenses vérifiables (RLVR).

• Stochasticité : Seules deux opérations sont stochastiques : l'échantillonnage des nouveaux nœuds $G_t$ générés à partir du modèle de base et l'échantillonnage de la décision du routeur $k^*_t$. L'élagage, le décalage et la réutilisation des nœuds sont déterministes.
• Vraisemblance (Likelihood) : La log-probabilité d'une trajectoire $F$ est la somme des log-probabilités des jetons nouvellement générés et des probabilités de sélection du routeur :
  $$\log p_{\theta,\phi}(F | q) = \sum_{t=1}^T \sum_{v \in G_t} \log \tilde{\pi}_\theta(\text{tok}(v) | \text{ctx}(v)) + \sum_{t=1}^T \log \rho_\phi(k^*_t | x_{<t}, \mathcal{A}_t; \theta)$$
• Optimisation : Le cadre est entraîné à l'aide d'un objectif de rapport de vraisemblance (par exemple, GRPO) où l'avantage $A(F)$ issu d'un vérificateur récompense à la fois la génération de jetons du modèle de base et les décisions de sélection du routeur.

Contributions Clés

1. Branchement au Niveau du Jeton avec Identité Discrète : Contrairement aux méthodes de jetons mous qui fusionnent les candidats en mélanges continus, le LBR préserve les candidats en tant que branches transmises discrètes. Cela permet au routeur d'exploiter les états cachés post-candidats — une information révélée uniquement après qu'un jeton candidat a été généré et traité par le modèle.
2. Entraînement de Bout en Bout Traçable : En définissant une vraisemblance spécifique pour la trajectoire de l'arbre, le LBR permet au modèle de base et au routeur d'être optimisés conjointement sous le même principe RLVR utilisé dans les systèmes de raisonnement à jetons discrets.
3. Mise à l'Échelle Efficace au Moment de l'Inférence : Le cadre introduit de la largeur au niveau du jeton avec un regard vers l'avant peu profond ($L$ est petit, par exemple 1 ou 2), évitant la surcharge de calcul d'une recherche de pleine profondeur tout en tirant parti des preuves au-delà de la distribution du jeton suivant racine.
4. Routeur par Attention d'Ensemble : Les auteurs introduisent une architecture de routeur qui utilise l'attention inter-sous-arbres pour comparer les candidats frères, démontrant que le routage est une décision contrastive plutôt qu'un reclassement indépendant.

Résultats Expérimentaux

Planification Hiérarchique Synthétique

Sur une tâche synthétique impliquant la de l'accessibilité de graphes par translation de radix :

• Performance : Le LBR a atteint la précision cible la plus élevée (0,967), surpassant la CoT discrète (0,914) et le "Soft Thinking" (0,845).
• Analyse du Mécanisme : Les expériences de sondage ont révélé que les états cachés post-candidats dans le LBR sont nettement plus prédictifs de la cible finale que les états pré-branchement. Cela confirme que la transmission des candidats permet au modèle d'accéder à des informations en aval qui sont indisponibles pour le décodage autoregressif standard.
• Représentation : Contra contrairement au "Soft Thinking", qui brouille les identités conceptuelles via des mélanges continus, le LBR préserve les identités discrètes des branches, rendant le chemin correct plus récupérable.

Benchmarks de Raisonnement Mathématique

Évalué sur des modèles DeepSeek-R1-Distill-Qwen (1,5B et 7B) à travers six benchmarks (Minerva, AIME, MATH500, etc.) :

• Gains de Performance : Le LBR a amélioré à la fois le Pass@1 et le Pass@32 par rapport à :
  • La Chaîne de Pensée (CoT) discrète.
  • Le RLVR de jetons discrets classique.
  • Les bases de branchement à jetons mous compatibles avec le RL (ex: Multiplex Thinking).
• Ablation de la Profondeur : Augmenter la profondeur du regard vers l'avant de $L=1$ à $L=2$ améliore davantage les performances, particulièrement sur le backbone 7B, suggérant que des continuations locales plus profondes fournissent plus de preuves de routage.
• Ablation du Routeur : Supprimer l'attention inter-sous-arbres (en utilisant un routeur indépendant) a entraîné des performances moindres, validant que la comparaison contrastive des branches sœurs est cruciale.
• Dynamique d'Entraînement : Pendant l'entraînement RLVR, l'entropie du routeur a diminué à mesure que le routeur devenait plus décisif, tandis que l'entropie du modèle de base aux positions routées a également diminué, indiquant une co-adaptation où le modèle apprend à générer des branches locales plus distinguables.

Signification et Revendications

L'article affirme que le Routage de Branches Locales offre une forme efficace, entraînable et discrète de mise à l'échelle du modèle de langage au moment de l'inférence. Sa principale signification réside dans la démonstration que :

1. Les preuves post-candidats sont précieuses : Les décisions de routage bénéficient des états cachés induits par les jetons candidats, et non seulement de la distribution pré-branchement.
2. Le branchement discret est supérieur au mélange mou : Préserver les identités de branches discrètes permet une prise de décision plus claire par rapport aux mélanges de jetons mous, qui peuvent occulter l'identité des concepts.
3. Scalabilité : Le cadre offre un juste milieu entre la CoT monothreadée et la recherche d'arbre complète coûteuse, permettant une mise à l'échelle efficace via un branchement local léger.

Les auteurs reconnaissent la limite de la localité : le LBR est contraint par la fenêtre de regard vers l'avant peu profonde et ne réalise pas de recherche globale sur des solutions complètes. Des travaux futurs sont suggérés pour combiner le LBR avec une profondeur adaptative ou une recherche globale sélective.