🤖 machine learning

Compute Aligned Training: Optimizing for Test Time Inference

Ce papier présente l'« Entraînement Aligné au Calcul », un cadre novateur qui aligne les objectifs d'entraînement des grands modèles de langage avec les stratégies d'inférence au moment du test en dérivant de nouvelles fonctions de perte, améliorant ainsi considérablement l'évolution des performances par rapport aux approches SFT et RL standard.

Auteurs originaux : Adam Ousherovitch, Ambuj Tewari

Publié 2026-04-29

📖 6 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Adam Ousherovitch, Ambuj Tewari

Article original sous licence CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Ceci est une explication générée par l'IA de l'article ci-dessous. Elle n'a pas été rédigée ni approuvée par les auteurs. Pour une précision technique, consultez l'article original. Lire la clause de non-responsabilité complète

Imaginez que vous formiez un étudiant pour un type d'examen final très spécifique.

L'Ancienne Méthode (Entraînement Standard) :
Traditionnellement, lorsque nous enseignons aux modèles d'IA (comme ceux qui rédigent des essais ou résolvent des problèmes de mathématiques), nous agissons comme un enseignant strict qui note chaque devoir individuellement. Si l'étudiant répond correctement à une question, nous disons « Bien joué ! » et s'il se trompe, nous disons « Réessayez ». Nous voulons que l'étudiant soit parfait pour obtenir la première réponse juste à chaque fois.

Le problème est que cela ne les prépare pas au véritable examen. Dans le monde réel, nous ne demandons pas au modèle une seule réponse. Nous lui demandons de générer de nombreuses réponses différentes (comme demander à un étudiant de rédiger 10 brouillons différents d'un essai), puis nous choisissons la meilleure, ou nous votons pour la plus courante.

Si vous formez un étudiant à être parfait dès le premier brouillon, il pourrait devenir trop confiant et cesser d'explorer de nouvelles idées. Il pourrait rester coincé dans une façon de penser « sûre ». Lorsque vous lui demandez de générer 10 brouillons plus tard, il pourrait simplement écrire 10 versions légèrement différentes de la même réponse « sûre », et aucune d'elles ne pourrait être la solution brillante et créative dont vous aviez besoin.

La Nouvelle Méthode (Entraînement Aligné sur le Calcul - CAT) :
Les auteurs de cet article proposent une nouvelle méthode d'entraînement appelée Entraînement Aligné sur le Calcul (CAT). Au lieu de noter l'étudiant sur un seul devoir, ils le notent en fonction de la façon dont il se comporterait dans le format de l'examen réel.

Voici comment cela fonctionne, en utilisant quelques analogies :

1. L'Analogie du « Pass@N » (Le Ticket de Loto)

Imaginez que l'examen vous permet d'acheter N tickets de loto (disons 64 tickets) pour une seule question. Vous gagnez si l'un quelconque de vos 64 tickets est le numéro gagnant.

Entraînement Standard : Si l'étudiant a déjà 50 % de chances de gagner avec un ticket, l'enseignant continue de le pousser pour atteindre 90 % ou 99 %. Mais dans un loto à 64 tickets, une fois que vous avez 50 % de chances, acheter plus de tickets garantit une victoire de toute façon. L'enseignant gaspille de l'énergie à essayer de rendre un étudiant « bon » « parfait » sur une question facile, tout en ignorant les questions difficiles où l'étudiant n'a qu'1 % de chances.
Entraînement CAT : L'enseignant réalise : « Hé, cet étudiant a déjà de fortes chances de gagner avec 64 tickets. Je vais arrêter de noter cette question facile avec autant d'insistance. » Au lieu de cela, il concentre toute son énergie sur les questions difficiles où l'étudiant échoue actuellement. Il apprend à l'étudiant à répartir ses chances, garantissant que au moins un des 64 tickets est un gagnant, plutôt que d'essayer de rendre un ticket spécifique parfait.

2. L'Analogie du « Vote Majoritaire » (L'Élection)

Imaginez que l'examen demande au modèle de générer 10 réponses, et que la classe vote pour la plus populaire.

Entraînement Standard : L'enseignant essaie de faire en sorte que la première réponse de l'étudiant soit absolument la plus populaire, même si elle gagne déjà par un large écart. C'est comme un candidat qui remporte déjà 90 % des voix ; l'enseignant continue de lui dire de faire campagne plus fort, ce qui est une perte de temps.
Entraînement CAT : L'enseignant regarde le « point de bascule ». Si la réponse de l'étudiant perd actuellement de justesse, l'enseignant lui donne un énorme coup de pouce pour l'aider à franchir la ligne d'arrivée. Si l'étudiant gagne déjà confortablement, l'enseignant arrête de lui donner des points supplémentaires. Cela force le modèle à se concentrer sur les questions « de champ de bataille » où un peu d'effort supplémentaire peut inverser le vote.

3. L'Analogie du « Meilleur des N » (Le Concours de Talent)

Imaginez que le modèle génère 10 chansons, et que vous ne gardez que la seule meilleure.

Entraînement Standard : L'enseignant essaie de faire en sorte que la chanson moyenne sonne bien. Cela conduit à une musique sûre, ennuyeuse, « moyenne », qui n'est jamais mauvaise, mais jamais étonnante non plus.
Entraînement CAT : L'enseignant dit au modèle : « Ce n'est pas grave si 9 de tes chansons sont terribles, tant que la 10e est un chef-d'œuvre. » Cela encourage le modèle à prendre des risques et à essayer des idées étranges et à forte variance. Il apprend à produire une grande variété de sorties, sachant que le processus de « recherche » (choisir la meilleure) filtrera les échecs et gardera le gagnant.

Qu'ont-ils réellement fait ?

Les chercheurs ont testé cette idée de trois manières spécifiques :

Problèmes de Mathématiques : Ils ont entraîné des modèles d'IA à résoudre des problèmes de mathématiques. Lorsqu'ils ont utilisé le CAT pour préparer les modèles au « Pass@N » (générer de nombreuses réponses et choisir la bonne), les modèles sont devenus significativement meilleurs pour résoudre des problèmes difficiles par rapport à la méthode standard.
Vote : Ils ont entraîné des modèles à générer des réponses pour un « Vote Majoritaire ». Là encore, les modèles CAT ont mieux performé lorsque la stratégie de vote a été appliquée.
Conception de Protéines : Ils ont même testé cela sur un type d'IA complètement différent qui conçoit des protéines (les blocs de construction de la vie). Dans un scénario où l'IA devait trouver une structure protéique rare et de haute qualité parmi beaucoup de mauvaises, les modèles entraînés avec le CAT étaient bien meilleurs pour trouver la protéine « jackpot » que les modèles standard.

La Conclusion

L'article soutient que la façon dont vous entraînez un modèle doit correspondre à la façon dont vous l'utilisez.

Si vous prévoyez d'utiliser le modèle pour générer de nombreuses options et choisir la meilleure, vous ne devriez pas l'entraîner à être parfait du premier coup. Vous devriez l'entraîner à être bon pour créer un ensemble d'options où la meilleure a de fortes chances d'être trouvée.

Ils appellent cela « Entraînement Aligné sur le Calcul » car il aligne le processus d'entraînement sur le « calcul » (la puissance de réflexion supplémentaire) utilisé au moment du test. C'est un moyen d'obtenir de meilleurs résultats sans avoir besoin d'ordinateurs plus puissants ou de plus de temps d'entraînement ; vous changez simplement les règles du jeu pour correspondre à la réalité de la façon dont le modèle sera utilisé.

1. Énoncé du problème : Le désalignement entraînement-inférence

Le papier aborde une rupture fondamentale dans le développement moderne des grands modèles de langage (LLM) : les objectifs d'entraînement sont optimisés pour la performance sur un échantillon unique, tandis que le déploiement repose de plus en plus sur des stratégies de calcul au moment de l'inférence (recherche).

Le statu quo : Les méthodes post-entraînement standard comme le fine-tuning supervisé (SFT) et l'apprentissage par renforcement (RL) optimisent la vraisemblance d'une seule sortie correcte (par exemple, minimiser l'entropie croisée ou maximiser le reward attendu d'un échantillon).
Le conflit : Les stratégies d'inférence modernes (par exemple, Pass@N, Vote majoritaire, Best-of-N) génèrent plusieurs solutions candidates et les agrègent pour trouver la meilleure réponse.
- Exemple (Pass@N) : Si un modèle a 50 % de chances d'être correct, augmenter cette probabilité à 90 % rapporte des rendements décroissants pour une métrique Pass@100 (où le succès est déjà quasi certain). L'entraînement standard continue cependant de pousser vers une vraisemblance de 100 %, gaspillant ainsi la capacité sur des échantillons "faciles".
- Exemple (Vote majoritaire) : Un modèle n'a besoin que d'une pluralité (par exemple, 35 % de probabilité) pour l'emporter sur ses rivaux. L'entraînement standard pousse vers $p \to 1.0$ , encourageant des victoires "écrasantes" plutôt qu'un consensus robuste, ce qui peut dégrader les performances lorsque le modèle est contraint d'agir comme un ensemble.
Conséquence : Ce désalignement provoque un effondrement de la politique (réduction de la diversité) et empêche les modèles d'utiliser efficacement des budgets de calcul au moment de l'inférence accrus.

2. Méthodologie : Compute Aligned Training (CAT)

Les auteurs proposent le Compute Aligned Training (CAT), un cadre qui aligne les objectifs d'entraînement avec la stratégie d'inférence spécifique utilisée au moment du test.

Conceptualisation de base

Les auteurs conceptualisent les stratégies d'inférence non pas comme des heuristiques a posteriori, mais comme des opérateurs différentiables ( $T$ ) qui transforment la politique de base $\pi_\theta$ en une distribution effective au moment du test $\tilde{\pi}_\theta$ :
$\tilde{\pi}_\theta(y|x) = T(\pi_\theta(\cdot|x), \phi)$
où $\phi$ représente les hyperparamètres (par exemple, le budget d'échantillonnage $N$ ).

La fonction de perte unifiée

Au lieu de minimiser le négatif de la vraisemblance de la vérité terrain sous la politique de base ( $\pi_\theta$ ), le CAT minimise la perte sous la politique effective au niveau du système ( $\tilde{\pi}_\theta$ ) :
$L_{CAT}(\theta) = -\mathbb{E}_{(x, y^*)} [\log \tilde{\pi}_\theta(y^*|x)]$

Mécanisme de mise à l'échelle des gradients

Pour rendre cela calculable sans échantillonner tout l'espace de sortie, les auteurs dérivent une approximation diagonale du gradient. Ils montrent que le gradient de l'objectif CAT peut être approximé en pondérant le gradient standard par un facteur de poids $w(p, \phi)$ :
$\nabla_\theta L_{CAT} \approx w(p, \phi) \cdot \nabla_\theta L_{CE}$
Ici, $p$ est la probabilité de base de la réponse correcte, et $w$ représente l'utilité marginale d'augmenter $p$ par rapport à la stratégie au moment du test.

Facteurs d'échelle clés dérivés :

Pass@N (SFT & RL) :
- $w(p, N) = \frac{N p (1-p)^{N-1}}{1 - (1-p)^N}$
- Comportement : Supprime les gradients pour les échantillons faciles (p élevé) où le succès est déjà garanti, en concentrant l'apprentissage sur les échantillons "difficiles" où les gains marginaux améliorent la couverture.
Vote majoritaire (SFT & RL) :
- $w(p, N, k) \approx \frac{k \binom{N}{k} p^k (1-p)^{N-k}}{\sum_{i=k}^N \binom{N}{i} p^i (1-p)^{N-i}}$
- Comportement : Concentre la masse du gradient près du seuil de consensus ( $p \approx k/N$ ). Il ignore les échantillons qui gagnent déjà ou perdent sans espoir, se concentrant uniquement sur le "point de bascule".
Best-of-N (RL) :
- $w \propto N (P_{<y})^{N-1}$
- Comportement : Crée une dynamique de "vainqueur prend tout". Les gradients renforcent uniquement les échantillons à la récompense la plus élevée dans un lot, encourageant le modèle à explorer les queues à forte variance et à forte récompense de la distribution plutôt que de s'effondrer vers une moyenne sûre.

Garanties théoriques

Le papier prouve que pour les stratégies compétitives (où l'augmentation de la probabilité d'un rival nuit au succès), l'approximation diagonale est une sous-estimation conservative du vrai gradient. Cela garantit que l'optimiseur ne fait jamais un pas dans la mauvaise direction, agissant comme un stabilisateur naturel contre l'effondrement de la politique.

3. Contributions clés

Cadre d'opérateurs unifié : Formalise les stratégies au moment du test comme des opérateurs différentiables, dérivant un mécanisme général pour aligner l'entraînement sur des stratégies d'inférence arbitraires (SFT et RL).
Extension au-delà de Pass@N : Généralise les travaux précédents (par exemple, Direct Coverage Optimization) au Vote majoritaire et au Best-of-N, couvrant à la fois les paradigmes SFT et RL.
Agnosticisme de la modalité : Démontre l'applicabilité du cadre au-delà des LLMs vers les modèles de langage protéique (PLM), montrant une amélioration des performances dans des paysages de récompense complexes et multi-modaux.
Analyse de la variance : Fournit une analyse théorique et empirique de l'inflation de la variance des gradients, proposant des techniques de normalisation (par exemple, estimateurs pondérés par le logarithme) pour stabiliser l'entraînement.

4. Résultats expérimentaux

Les auteurs ont évalué le CAT sur le benchmark MATH (utilisant Mistral-7B) et la Génération de protéines (utilisant ProtGPT2).

A. Modèles de langage (Benchmark MATH)

Pass@N (SFT) : Les modèles entraînés avec CAT pour $N=64$ ont sacrifié la précision sur un échantillon unique (Pass@1 a chuté d'environ 2,3 %) mais ont réalisé une amélioration absolue de 7,8 % dans la couverture Pass@64 par rapport au SFT standard.
Vote majoritaire (SFT) : Les modèles CAT ont systématiquement surpassé les bases pour tous les budgets d'inférence. Un "point idéal" a été observé à $N=16$ , où le modèle a équilibré biais et variance de manière optimale.
Apprentissage par renforcement (GRPO) : Le RL standard a souffert d'un effondrement de la politique, saturant rapidement (Pass@32 $\approx$ 35,8 %). Les modèles entraînés avec CAT ont préservé la variance nécessaire à la recherche, atteignant 40,0 % à Pass@32.
Effet de préchauffage (Warmup) : L'utilisation de CAT pour la phase de préchauffage SFT a considérablement amélioré les performances RL subséquentes, augmentant même la précision en tir unique (Pass@1) en enrichissant les chemins de raisonnement du modèle.

B. Génération de protéines (Modèles de langage protéique)

Inconditionnel (Vallée de la mort) : Dans un paysage de récompense multi-modal (optimum local à $h \approx 0.35$ , optimum global à $h \approx 0.75$ , séparés par une "vallée" d'agrégation), le RL standard est resté coincé dans l'optimum local. Le Best-of-N CAT a réussi à "sauter" la vallée, trouvant le maximum global en maintenant une forte variance distributionnelle.
Génération conditionnelle : Les modèles CAT ont appris à générer des séquences avec une hydrophobicité inverse par rapport aux entrées, démontrant une mise à l'échelle supérieure et un meilleur alignement avec la tâche relationnelle par rapport au RL standard.

5. Signification et implications

Changement de paradigme : Le papier soutient que les objectifs d'entraînement doivent être conçus en tenant compte de la procédure d'inférence en aval. À mesure que le calcul au moment du test devient le levier principal pour l'augmentation des capacités des modèles, l'entraînement doit optimiser le résultat agrégé, et non l'échantillon unique.
Efficacité : Le CAT réalise ces gains sans augmenter les coûts de calcul d'entraînement. Il pondère simplement les gradients existants en fonction de l'utilité marginale de l'échantillon pour la stratégie au moment du test.
Stabilité vs Exploration : Le cadre fournit une manière principée d'échanger la précision sur un échantillon unique contre la diversité distributionnelle, permettant aux modèles d'utiliser efficacement des procédures de recherche comme Best-of-N et Vote majoritaire sans s'effondrer vers un comportement de recherche de mode.
Généralisabilité : En prouvant l'efficacité à la fois dans les LLMs et les modèles de langage protéique, le papier suggère que le "Compute Aligned Training" est un principe universel pour tout modèle génératif utilisant une recherche au moment du test.

En conclusion, le Compute Aligned Training comble le fossé entre l'apprentissage et la recherche, garantissant que les modèles ne sont pas seulement "intelligents" pour générer une réponse, mais sont optimisés pour être "intelligents" afin de trouver la meilleure réponse lorsqu'ils disposent du budget de calcul pour rechercher.