Cola2 devance Llama-3-8B sur HellaSwag avec 43 fois moins de tokens

Une équipe de chercheurs a dévoilé Cola2, un framework de modélisation hiérarchique du langage dans lequel le sens vit dans un espace latent continu, tandis qu'un décodeur discret assemble le texte. Entraîné sur 346,03 milliards de tokens sources, un checkpoint qui n'a pas encore convergé obtient une moyenne de 60,10 sur huit tâches et devance Llama-3-8B sur HellaSwag et StoryCloze. Le budget de pré-entraînement annoncé pour Llama-3 est pourtant plus de 43 fois supérieur.

Que compresse exactement l'espace latent

L'idée de génération hiérarchique circule depuis longtemps : le modèle planifie d'abord le sens à grands traits, puis le réalise en tokens. Le point faible reste toujours le même. Si les latents compressent la mauvaise information, ou s'y prennent mal, le décodeur les ignore ou se réduit à une simple photocopieuse.

Cola2 s'attaque au problème en plusieurs temps. Le texte est découpé dynamiquement en spans sémantiques, et non en segments de longueur fixe. Le taux de compression est piloté par une rétroaction négative, sans loss auxiliaire : le système ajuste lui-même un seuil, sans ajouter de termes de pénalité superflus à l'entraînement. Pour obtenir des latents à la fois compacts et prédictifs, les auteurs combinent deux signaux : la reconstruction du texte source et la génération classique.

Décodeur à diffusion et équilibre entre copie et autonomie

Le décodeur est ici un décodeur à diffusion masquée, bloc-causal. Il est entraîné sous trois types de conditions : avec des latents propres, avec des latents bruités et avec des latents entièrement supprimés. Cela permet d'appliquer le classifier-free guidance. En modifiant la force du guidance, on peut déplacer le comportement du modèle entre la restitution fidèle du sens à partir du latent et la génération autonome.

Un détail d'ingénierie concerne le flux des gradients. Les gradients de la prédiction de tokens traversent les latents jusqu'à l'encodeur, et les représentations s'adaptent aux besoins du décodeur. En revanche, l'entraînement du prior par Flow Matching est découplé de celui des représentations par un stop-gradient. Le prior apprend à prédire les latents suivants, mais il n'entraîne pas l'encodeur dans son sillage et ne détruit pas l'espace que celui-ci construit.

L'inférence réserve aussi une subtilité. Le texte généré est réencodé avec un découpage en spans figé, afin que l'historique des latents corresponde au contexte textuel réel, et non à ce que le modèle « pensait » auparavant. Sans cette étape, les erreurs entre la partie continue et la partie discrète s'accumuleraient de bloc en bloc.

Les chiffres et leurs limites

• Budget d'entraînement dans les expériences principales : environ 210 milliards de tokens, cinq tâches downstream.
• Passage à l'échelle : les auteurs le décrivent comme compétitif et non saturé.
• Grand entraînement : 346,03 milliards de tokens sources, checkpoint non convergé.
• Moyenne sur huit tâches : 60,10.
• Devance Llama-3-8B sur HellaSwag et StoryCloze avec un budget de pré-entraînement plus de 43 fois inférieur.

Il faut lire la comparaison avec Llama-3-8B avec prudence. HellaSwag et StoryCloze mesurent la vraisemblance de la suite d'un texte et le bon sens, c'est-à-dire précisément ce qu'un modèle doté d'une couche explicite de structure sémantique capte bien. Sur des benchmarks de connaissances et de raisonnement, le tableau pourrait être différent, et d'après le texte de l'article, la victoire n'est revendiquée que sur deux tâches sur huit. En revanche, la courbe de passage à l'échelle ne plafonne pas, ce qui signifie que le potentiel avec des budgets plus importants n'est pas épuisé.

Où cela peut servir

Pour ceux qui construisent des modèles, l'enseignement pratique est le suivant : on peut chercher l'efficacité en données ailleurs que dans l'architecture du transformer ou la curation du corpus, à savoir dans la conception de la couche sémantique intermédiaire. Si les latents compressent le texte en spans et que le décodeur génère par blocs, le coût d'inférence peut aussi baisser : moins d'étapes par unité de sens.

Pour les chercheurs travaillant sur les modèles de langage à diffusion, Cola2 propose une recette prête à l'emploi pour raccorder les parties continue et discrète : rétroaction négative pour la compression, trois régimes de conditions pour le décodeur, stop-gradient pour le prior, recalcul des latents pendant la génération. Chacun de ces éléments peut être testé et transposé séparément.

Pour l'industrie, c'est le signe que les alternatives au schéma purement autorégressif gagnent du poids. Pour l'instant, il s'agit d'un modèle de la classe 8B, dont l'entraînement n'est pas terminé. Mais si le passage à l'échelle se confirme à des dizaines et des centaines de milliards de paramètres, les modèles latents hiérarchiques deviendront un candidat sérieux pour succéder à la prédiction classique du token suivant (next-token prediction).