Select to Think: Unlocking SLM Potential with Local Sufficiency

1. Énoncé du problème

Les petits modèles de langage (SLM) offrent une efficacité computationnelle pour un déploiement évolutif, mais manquent souvent des capacités de raisonnement des grands modèles de langage (LLM). Les stratégies actuelles visant à combler cet écart font face à un « dilemme » :

• Inférence collaborative : Utiliser un LLM pour intervenir aux points de divergence du raisonnement (routage au niveau des tokens) introduit une latence et un coût élevés en raison d'appels externes synchrones.
• Distillation standard : Entraîner un SLM à imiter la distribution générative ouverte d'un LLM échoue souvent en raison de l'« écart de capacité ». Les SLM peinent à reproduire avec précision les pics de probabilité complexes, conduisant à une imitation fragile plutôt qu'à une prise de décision robuste.

La question centrale abordée est la suivante : La génération ouverte par un LLM est-elle strictement nécessaire aux points de divergence, ou l'ensemble des candidats propres au SLM suffit-il pour capturer le choix préféré du LLM ?

2. Insight clé : la suffisance locale

Les auteurs identifient un phénomène appelé suffisance locale. Par une analyse empirique, ils démontrent qu'aux points où un SLM et un LLM divergent (erreurs de raisonnement), le token préféré du LLM est constamment contenu dans les K meilleurs prédictions de token suivant du SLM, même si le SLM échoue à le sélectionner comme son choix numéro 1.

• Preuve : Les 8 meilleurs candidats d'un SLM de 1,5 milliard de paramètres capturent le token préféré d'un LLM de 32 milliards avec un taux de réussite de 95 %.
• Implication : Le goulot d'étranglement pour les SLM n'est pas un manque de connaissances (potentiel de génération) mais un échec dans le classement des candidats. Le token correct existe dans la distribution locale mais est supprimé par le score interne du SLM.

3. Méthodologie : le cadre S2T

Le papier propose SELECT TO THINK (S2T), un cadre qui reformule le rôle du LLM de la génération de tokens à la sélection parmi les propositions du SLM.

A. S2T (Phase collaborative)

Au lieu que le LLM génère le token suivant, le SLM génère un ensemble de candidats $C_t$ (les K meilleurs tokens). Le LLM agit ensuite comme un discriminateur, notant ces candidats en fonction de leur probabilité conditionnelle $p_{\theta_L}(c | s_t)$.

• Avantage : Cela simplifie le signal de supervision, passant de l'ajustement à une distribution ouverte de haute dimension à une tâche de classement discrète, plus facile à distiller.

B. S2T-LOCAL (Phase sans enseignant)

Pour éliminer la dépendance au LLM lors de l'inférence, les auteurs distillent la logique de sélection dans le SLM lui-même.

• Mécanisme : Ils utilisent le cadre Zero-overhead Inference-time Prediction (ZIP).
• Implémentation : Le SLM est affiné pour réaffecter les logits d'un ensemble de tokens réservés ($R$) afin de coder des scores de préférence.
  • Pour chaque candidat $c \in C_t$, le modèle effectue un passage avant sur l'état augmenté $(x, y_{<t} \oplus c)$.
  • Les logits des tokens réservés sont lus pour calculer un score de sélection $S_{local}(c; s_t)$.
• Objectif d'entraînement : Le modèle minimise une perte de sélection (Entropie croisée sur l'ensemble des candidats) tout en imposant un régularisateur de stabilité (divergence KL sur le vocabulaire standard) pour empêcher le modèle d'oublier ses capacités linguistiques.

4. Contributions clés

1. Insight (Suffisance locale) : Formalisation et vérification empirique du fait que les SLM contiennent intrinsèquement les tokens préférés du LLM dans leurs K meilleures prédictions aux points de divergence, déplaçant le focus de l'optimisation de la génération vers la sélection.
2. Cadre (S2T) : Reformulation du raisonnement collaboratif comme une tâche de sélection discrète, contournant l'écart de capacité inhérent à la distillation standard.
3. Méthode (S2T-LOCAL) : Introduction d'une variante sans enseignant qui internalise la logique de sélection à l'aide de tokens réservés, permettant un re-classement autonome sans appels externes au LLM.
4. Efficacité : Démonstration que cette approche atteint des performances comparables à la cohérence multi-chemins (8 chemins) avec une efficacité de trajectoire unique.

5. Résultats expérimentaux

Les auteurs ont évalué S2T-LOCAL sur le raisonnement mathématique (GSM8K, MATH, OlympiadBench, AIME25), la programmation (HumanEval) et les connaissances générales (MMLU-Pro) en utilisant des modèles Qwen2.5 (élèves de 0,5B et 1,5B contre un enseignant de 32B).

• Gains de performance :
  • S2T-LOCAL a amélioré la précision du décodage glouton d'une moyenne de 24,1 % sur le modèle de 1,5B et de 38 % sur le modèle de 0,5B.
  • Il a égalé la précision de la cohérence à 8 chemins (qui nécessite 8 fois plus de calcul) tout en opérant avec une efficacité de trajectoire unique.
• Taux de réussite : Le SLM de 1,5B a atteint un Hit@8 de 95 % par rapport à l'enseignant de 32B, confirmant l'hypothèse de la suffisance locale.
• Efficacité de la distillation : Le sélecteur distillé a atteint un accord de 67 à 71 % avec le LLM sur des instances retenues, récupérant avec succès des tokens que le SLM de base sous-estimait.
• Efficacité : Par rapport aux méthodes collaboratives (comme « Takeover » où le LLM génère le token), S2T-LOCAL a réduit la latence d'inférence d'environ 75 % (par exemple, de 81,5 s à 21,0 s sur AIME25) tout en maintenant un nombre de tokens générés similaire.

6. Importance

Ce travail modifie fondamentalement le paradigme d'amélioration du raisonnement des SLM. Au lieu de considérer les SLM comme des imitateurs passifs devant reproduire les distributions des LLM (ce qui est souvent impossible en raison des limites de capacité), S2T les traite comme des sélecteurs actifs de leurs propres propositions.

• Efficacité des ressources : Il permet aux appareils périphériques d'exécuter des tâches de raisonnement complexes avec une précision proche de celle des LLM, mais à la vitesse et au coût d'un petit modèle.
• Évolutivité : En découplant la sélection de la génération, il offre une voie évolutive pour améliorer le raisonnement sans les coûts prohibitifs des appels externes aux LLM ou de l'expansion massive des modèles.
• Généralisabilité : La propriété de « suffisance locale » s'applique à différentes échelles de modèles (de 0,5B à 1,5B) et architectures (Qwen, Gemma), suggérant qu'il s'agit d'une caractéristique intrinsèque des distributions de tokens des SLM.