Test-Time Scaling Makes Overtraining Compute-Optimal

1. Problématique

Les lois d'échelle actuelles pour les grands modèles de langage (LLM) traitent séparément deux phases critiques :

1. L'entraînement (Pretraining) : Des lois comme Chinchilla (Hoffmann et al., 2022) optimisent la taille du modèle ($N$) et le nombre de tokens d'entraînement ($D$) pour minimiser la perte (loss) sous un budget de calcul fixe. Elles supposent implicitement une inférence à un seul passage (single-pass, $k=1$).
2. L'inférence (Test-Time) : Des travaux récents montrent que l'augmentation du calcul au moment du test (via l'échantillonnage répété, ou repeated sampling, noté $k$) peut améliorer les performances. Cependant, ces travaux traitent généralement le modèle pré-entraîné comme un donné fixe et ne relient pas cette stratégie d'inférence aux décisions d'entraînement.

Le problème central : Il existe un compromis fondamental non résolu. Les petits modèles sont moins chers par échantillon mais moins performants par échantillon. L'avantage de l'échantillonnage répété ($k$) est une fonction non linéaire de la qualité de chaque échantillon. Aucune loi d'échelle existante ne joint l'optimisation de la taille du modèle, de la durée d'entraînement et du nombre d'échantillons d'inférence sous un budget global unique. Cela conduit à une sous-optimisation : les modèles sont souvent entraînés selon les ratios Chinchilla (environ 20 tokens par paramètre), alors que l'inférence moderne (avec $k \gg 1$) pourrait bénéficier de modèles plus petits et beaucoup plus entraînés (overtrained).

2. Méthodologie : Les lois d'échelle Train-to-Test (T2)

Les auteurs proposent un cadre unifié, Train-to-Test (T2), qui optimise conjointement $N$ (taille), $D$ (tokens d'entraînement) et $k$ (nombre d'échantillons) sous des contraintes de budget de calcul pour l'entraînement ($C_{train}$) et l'inférence ($C_{inf}$).

Les coûts sont modélisés ainsi :

• $C_{train} \approx 6ND$
• $C_{inf} \approx 2Nk$ (coût d'un passage avant pour un token)

Pour un budget d'inférence fixe, un modèle plus petit ($N$ réduit) permet d'allouer plus d'échantillons ($k$ augmenté).

L'article explore deux approches de modélisation distinctes qui convergent vers les mêmes conclusions :

Approche 1 : Modélisation de la Perte (Loss-based)

Cette méthode étend la loi d'échelle Chinchilla pour inclure l'impact de $k$ sur la perte négative du logarithme de vraisemblance (NLL).

• L'objectif est de minimiser l'espérance de $-\log(\text{pass@}k)$.
• La perte est modélisée comme : $\hat{L}(N, D, k) = E + \frac{A}{N^\alpha} + \frac{B}{D^\beta} + \frac{G}{k^\gamma}$.
• Le terme $\frac{G}{k^\gamma}$ capture l'amélioration due à l'échantillonnage répété, supposant que la distribution de difficulté des tâches suit une loi Beta.

Approche 2 : Modélisation de la Précision (Accuracy-based)

Cette méthode modélise directement la métrique pass@k (probabilité d'avoir au moins une réponse correcte sur $k$ tentatives).

• Elle évite les biais d'estimation en modélisant la distribution des probabilités de réussite par question comme une distribution Beta.
• Les paramètres de la distribution Beta ($\mu, \nu$) sont liés à la perte estimée par Chinchilla via des fonctions de lien (logit et log).
• La précision attendue est calculée analytiquement à partir de cette distribution Beta : $\text{pass@}k = 1 - \frac{B(a, b+k)}{B(a, b)}$.

Dans les deux cas, l'optimisation consiste à trouver le triplet $(N, D, k)$ qui maximise la performance (ou minimise la perte) sous les contraintes de budget, en substituant $k = C_{inf} / 2N$.

3. Contributions Clés

1. Optimisation End-to-End : Formalisation d'un problème d'optimisation conjointe couvrant l'entraînement et l'inférence, brisant le silo entre les lois d'échelle de pré-entraînement et les stratégies de déploiement.
2. Deux Cadres de Modélisation : Introduction de deux approches complémentaires (basée sur la perte NLL et basée sur la précision pass@k) qui, bien que différentes, aboutissent aux mêmes recommandations stratégiques.
3. Validation Empirique sur le Sur-entraînement : Création d'un banc de tests avec plus de 100 modèles (de 5M à 901M paramètres) incluant des checkpoints "sur-entraînés" (au-delà des ratios standards) pour valider les prédictions théoriques.
4. Robustesse Post-Entraînement : Démonstration que les gains liés au sur-entraînement persistent même après des étapes de post-entraînement (Fine-Tuning et Supervised Fine-Tuning).

4. Résultats Expérimentaux

Les expériences ont été menées sur 8 tâches (4 réelles : LAMBADA, ARC-Easy, SciQ, OpenBookQA ; et 4 synthétiques) avec des modèles de moins de 1 milliard de paramètres.

• RQ1 : Les décisions de pré-entraînement doivent-elles changer ?

  • Oui. Lorsque le budget d'inférence (avec échantillonnage répété) est pris en compte, la frontière optimale de pré-entraînement se déplace radicalement vers le régime de sur-entraînement.
  • Contrairement à Chinchilla qui recommande ~20 tokens par paramètre, T2 recommande des modèles beaucoup plus petits mais entraînés sur beaucoup plus de tokens (parfois des centaines de fois le ratio standard).
  • Les modèles optimaux T2 sont significativement plus performants que les modèles optimaux Chinchilla pour un budget d'inférence donné.

• RQ2 : Les prédictions T2 s'extrapole-t-elles aux checkpoints sur-entraînés ?

  • Oui. Les modèles entraînés spécifiquement dans la région "sur-entraînée" prédite par T2 surpassent systématiquement les modèles Chinchilla optimaux sur toutes les tâches.
  • L'approche 1 (NLL) montre une erreur d'extrapolation relative de 2,8 %, tandis que l'approche 2 (Précision) enregistre 8,4 %, confirmant la validité des prédictions.

• RQ3 : Les résultats survivent-ils au post-entraînement ?

  • Oui. Après Fine-Tuning (FT) et Supervised Fine-Tuning (SFT) sur des tâches réelles, la tendance vers les petits modèles sur-entraînés persiste.
  • Bien que le sur-entraînement rende le fine-tuning légèrement plus difficile (conformément à la littérature récente), les modèles sur-entraînés restent supérieurs aux modèles Chinchilla optimaux même après ajustement.

5. Signification et Conclusion

Ce travail remet en question le dogme actuel de l'entraînement des LLM. Il démontre que l'optimisation du calcul doit être holistique, intégrant le coût de l'inférence dès la phase de conception du modèle.

• Changement de paradigme : Au lieu d'entraîner de grands modèles pour une inférence rapide et unique, il est plus efficace (en termes de coût total) d'entraîner des modèles plus petits sur des données massives et de compenser la perte de qualité par un échantillonnage répété au moment du test.
• Implication pratique : Pour les déploiements modernes où le calcul d'inférence est disponible (ex: recherche, raisonnement complexe), les praticiens devraient adopter des stratégies de "sur-entraînement" guidées par les lois T2, plutôt que de suivre strictement les ratios Chinchilla.
• Avenir : Les auteurs suggèrent que cette approche est particulièrement pertinente pour les modèles de pointe (frontier models) destinés à des tâches de raisonnement complexe, où le coût de l'inférence multiple est justifié par la précision accrue.

En résumé, l'article prouve que le sur-entraînement n'est pas un gaspillage, mais une stratégie optimale lorsque l'on considère le coût et les gains de l'inférence par échantillonnage répété.