Résumé Technique : LARK (Learnability-Grounded Trajectory Selection)

Énoncé du Problème

Le papier traite du défi de la distillation de raisonnement, un paradigme où un modèle de langage (LLM) étudiant plus petit est affiné (fine-tuned) pour imiter les trajectoires de raisonnement générées par des modèles enseignants plus larges. Alors que les méthodes existantes de sélection de données opèrent souvent au niveau de la question ou s'appuient sur des heuristiques telles que la qualité de la trajectoire, la correction ou la confiance du modèle, elles négligent fréquemment un facteur critique : la capacité d'apprentissage (learnability).

Plus précisément, pour une seule question, plusieurs trajectoires de raisonnement candidates peuvent exister (générées par différents enseignants, échantillonnages ou styles). Même si ces trajectoires mènent à la même réponse finale correcte, elles fournissent des signaux d'entraînement très différents. Une trajectoire qui semble « de haute qualité » ou « naturelle » n'est pas nécessairement celle à partir de laquelle le modèle étudiant peut apprendre le plus efficacement. La question centrale posée est la suivante : Pouvons-nous concevoir un critère de sélection de trajectoire fondé sur des principes qui garantisse que les trajectoires de raisonnement distillées soient apprenables par le modèle étudiant ?

Méthodologie : LARK

Les auteurs proposent LARK (Learnability-grounded Anchor-time Ranking), une méthode qui sélectionne les trajectoires en se basant sur une définition de la capacité d'apprentissage issue de la théorie de l'optimisation plutôt que sur des métriques de qualité statiques.

1. La capacité d'apprentissage comme objectif d'optimisation

LARK cadre la sélection de trajectoire comme un problème d'optimisation de politique. Il définit un facteur de capacité d'apprentissage $\rho(\theta; q)$, qui caractérise le taux de diminution de la perte d'entraînement de l'étudiant. Ce facteur est défini comme le rapport du carré de la norme du gradient sur la perte :
$$\rho(\theta; q) \triangleq \frac{\|\nabla_\theta \mathcal{L}(\theta; q)\|^2}{\mathcal{L}(\theta; q)}$$
où $q$ est le vecteur de poids sur $K$ trajectoires candidates. Sous la condition de Polyak–Łojasiewicz (PL), un $\rho$ plus élevé implique une convergence plus rapide.

2. Approximation par temps d'ancrage (Anchor-Time)

L'optimisation directe de l'objectif de capacité d'apprentissage intégré dans le temps sur l'ensemble de la trajectoire d'ajustement fin est intraitable. LARK introduit la condition relative à l'ancre, qui postule que pour les modèles sur-paramétrés dans le régime de « l'entraînement paresseux » (lazy training), le taux de capacité d'apprentissage reste proportionnel à sa valeur aux paramètres d'ancrage initiaux $\theta_{ref}$. Cela permet d'approximer l'objectif en maximisant $\rho(\theta_{ref}; q)$.

3. Linéarisation locale et proxy de passage direct (Forward-Pass)

Maximiser $\rho(\theta_{ref}; q)$ directement est problématique car c'est une fonction quasi-convexe qui s'effondre vers une solution « one-hot » dégénérée (sélectionnant une seule trajectoire). Pour éviter cela, LARK :

• Linéarisation locale : Développe $\rho$ autour de la distribution uniforme $p$ (SFT standard sans sélection) en utilisant un développement de Taylor au premier ordre.
• Proxy de passage direct ($\hat{g}_k$) : Au lieu de calculer les gradients exacts (qui nécessitent des passages arrière/backward pour chaque candidat), LARK dérive un proxy calculable $\hat{g}_k$ en utilisant uniquement des passages directs/forward. Ce proxy estime le gradient local du taux de capacité d'apprentissage en se basant sur les résidus de prédiction par jeton (token-level) de l'étudiant et la perte d'entropie croisée.
  • Plus précisément, $\hat{g}_k$ favorise les trajectoires où l'étudiant est « confiant dans son erreur » de manière structurée : la perte $\ell_k$ est élevée (indiquant une mauvaise prédiction), mais la masse résiduelle est concentrée sur quelques jetons « presque corrects » (fournissant un signal de correction clair).

4. Politique de sélection régularisée par $\chi^2$

Pour empêcher le sélecteur de « hacker » le critère de capacité d'apprentissage en s'éloignant trop de la distribution originale des données (ce qui nuit à la généralisation), LARK impose un terme de régularisation $\chi^2$. Cela conduit à une règle de sélection à forme fermée et paramétrée par le budget :

• Les trajectoires sont classées par leur score $\hat{g}_k$.
• Les $B$ meilleures trajectoires sont sélectionnées.
• Contra lecture d'une troncature stricte top-$B$, LARK assigne des poids souples proportionnels à la marge du score au-dessus de la $(B+1)$-ième trajectoire. Cette règle de pondération est dérivée pour maximiser la borne inférieure de l'amélioration de la capacité d'apprentissage tout en contrôlant la divergence $\chi^2$ par rapport au préalable uniforme.

Contributions Clés

1. Formulation : Les auteurs formulent la sélection de trajectoire de raisonnement comme un problème d'optimisation de politique fondé sur la capacité d'apprentissage, introduisant le taux de capacité d'ancrage $\theta$ comme un objectif rigoureux.
2. Algorithme : Ils dérivent une règle de sélection à forme fermée, paramétrée par le budget, qui utilise un proxy de passage direct $\hat{g}_k$ et une régularisation $\chi^2$. Cela évite les passages arrière coûteux par trajectoire tout en garantissant théoriquement des bornes sur l'erreur d'estimation et la généralisation.
3. Validation Empirique : Le papier démontre que LARK surpasse systématiquement les lignes de base existantes (incluant les méthodes basées sur des heuristiques, la qualité ou le scoring côté étudiant) sur plusieurs modèles de base (Qwen-2.5-7B/1.5B, Llama-3.2-3B) et des benchmarks de raisonnement (AIME, AMC, GPQA, MATH-500).

Résultats Expérimentaux

• Performance : LARK a obtenu la meilleure moyenne Acc@5 sur tous les réglages de budget de modèle testés. Par exemple, sur Qwen-2.5-7B avec un budget de $B=3$, LARK a amélioré la performance de 7,07 points de pourcentage par rapport à la meilleure ligne de base (RSR).
• Efficacité : LARK est une méthode de passage direct uniquement (forward-pass-only), égalant le coût de calcul des lignes de base les moins coûteuses côté étudiant, alors que les méthodes basées sur le gradient exact nécessiteraient beaucoup plus de calculs.
• Analyse Diagnostique :
  • Décroissance de la perte : Les trajectoires sélectionnées par LARK ont induit la réduction la plus rapide de la perte SFT pendant l'entraînement.
  • Validité prédictive : Le score LARK $\hat{g}_k$ a montré une forte corrélation positive (Pearson $r=0,69$) avec la précision en aval lorsqu'il est agrégé par modèle enseignant.
  • Comportement au niveau des jetons : Les trajectoires sélectionnées présentaient des probabilités de jetons corrects plus faibles mais une concentration de « jetons erronés » plus élevée, confirmant que LARK sélectionne des exemples où l'étudiant se trompe, mais où l'erreur est structurée et corrigeable.

Signification et Revendications

Le papier affirme que la capacité d'apprentissage (learnability) est un principe distinct et efficace pour une distillation de raisonnement efficace, séparé de la qualité ou de la difficulté de la trajectoire. En ancrant la sélection dans la dynamique d'optimisation du modèle étudiant, LARK identifie les trajectoires qui fournissent le signal d'apprentissage le plus efficace.

Les auteurs soulignent que leur méthode :

• Préserve la Généralisation : Grâce à la régularisation $\chi^2$, elle empêche la politique de sélection de faire du surapprentissage (overfitting) sur un ensemble étroit de trajectoires « faciles » ou « hackées ».
• Est Fondée sur la Théorie : La règle de sélection est dérivée des premiers principes de la théorie de l'optimisation (conditions PL, flot de gradient) avec des bornes d'erreur explicites.
• Est Pratique : Elle ne nécessite aucun réglage d'hyperparamètres pour la règle de sélection elle-même (la force de régularisation est déterminée par le budget $B$) et évite la surcharge de calcul des passages arrière.

Le travail suggère que de futures extensions pourraient appliquer cette perspective fondée sur la capacité d'apprentissage à des ensembles de candidats plus bruités et à des domaines de raisonnement plus larges, mais le champ actuel est strictement limité à la distillation de raisonnement dans la résolution de problèmes mathématiques.