STAPLE ajoute des milliards de paramètres à OLMo-2 et décode deux fois plus vite que le MoE

Une équipe de chercheurs a publié sur arXiv STAPLE, une méthode qui permet d'augmenter le nombre de paramètres d'un modèle de langage déjà entraîné, sans passer par un coûteux préentraînement sparse à partir de zéro. Sur OLMo-2 1B, elle apporte +4,27 points sur HumanEval, +2,80 sur MBPP, +2,27 sur GSM8K et +1,65 sur DROP par rapport à une simple poursuite de l'entraînement dense. Selon les auteurs, la génération ralentit de 1,5% au maximum et l'entraînement de 3% au maximum, alors que des milliards de paramètres sont ajoutés tardivement au modèle.

Pourquoi pas du MoE

Le calcul sparse est depuis longtemps considéré comme un moyen d'accroître la capacité d'un modèle sans gonfler le calcul actif pour chaque token. Le problème tient au coût : le préentraînement d'architectures sparse exige beaucoup de données et d'argent. Les méthodes d'upcycling existantes, c'est-à-dire la transformation d'un checkpoint dense en modèle sparse, visent surtout les Mixture-of-Experts. Les auteurs y identifient trois goulots d'étranglement : l'homogénéité des experts, la faible sample efficiency, et le surcoût système lié au routage et à la mémoire.

La mémoire sparse statique fonctionne autrement. Son accès est prévisible : les paramètres sont sélectionnés par l'identifiant du token, et non par la décision d'un routeur. Les mises à jour sont elles aussi liées à un token précis. Mais, comme le notent les auteurs, l'insertion naïve d'une telle mémoire dans un modèle préentraîné donne un rendement faible.

Comment fonctionne STAPLE

La méthode s'appuie sur STEM, où les projections up des blocs FFN sont remplacées par des embeddings indexés par token. Pour un entraînement from scratch, cela fonctionne. Sur un checkpoint existant, en revanche, le remplacement casse les calculs déjà accumulés. STAPLE procède autrement : les projections up restent en place et les embeddings s'ajoutent sous forme de chemin résiduel. La partie préentraînée continue de fonctionner comme avant, tandis que la nouvelle mémoire apprend ce que le modèle sait le moins bien sur les tokens : une information supplémentaire propre à chaque token.

Un détail d'ingénierie important : l'implémentation décharge les embeddings vers l'hôte. Le GPU ne conserve qu'environ 5% des embeddings les plus fréquents, le reste réside dans la mémoire de l'hôte. Compte tenu de la distribution des tokens, cela suffit pour que le débit ne chute presque pas.

Les chiffres

Principaux résultats du préprint :

• OLMo-2 1B, gain par rapport au continued dense training : HumanEval +4,27, MBPP +2,80, GSM8K +2,27, DROP +1,65.
• Le résultat moyen augmente avec la capacité statique, et les améliorations se maintiennent sur OLMo-2 7B.
• STAPLE surpasse les baselines sparse statiques : STEM et un modèle dans le style d'Engram.
• Face à l'upcycling MoE : score moyen plus élevé, décodage 2 fois plus rapide, moins de la moitié de la mémoire GPU consacrée au modèle et moins d'énergie par token.
• Surcoût : génération à moins de 1,5% de la base dense, entraînement à moins de 3%.

Une réserve importante : la comparaison avec le MoE sur la vitesse et l'énergie repose sur une évaluation centrée sur le débit de décodage. Pour les scénarios soumis à des exigences strictes de latence sur une requête isolée, le résultat peut différer. En outre, il ne s'agit pour l'instant que d'un préprint sur des modèles de 1B et 7B, et le passage à l'échelle de dizaines ou de centaines de milliards de paramètres n'est pas démontré.

Ce que cela change en pratique

Pour les équipes qui poursuivent l'entraînement de modèles ouverts, STAPLE offre une voie économique pour gagner en capacité : on prend un checkpoint existant, on y ajoute la mémoire statique, et le modèle continue de s'entraîner avec une faible perte de vitesse. Le gain est particulièrement net sur le code et les mathématiques, où la connaissance exacte des tokens et des motifs est particulièrement utile. Côté inférence, le profil mémoire est intéressant : l'essentiel des paramètres vit sur l'hôte, et le GPU n'a besoin que d'un petit cache. Cela abaisse les exigences en mémoire vidéo par rapport à un MoE de qualité comparable.

Avant de transposer l'approche en production, trois vérifications s'imposent : le comportement sur des modèles de plus de 7B, la robustesse face à un changement de domaine où la distribution des tokens fréquents diffère, et la latence avec de petits batchs. Si ces vérifications sont concluantes, l'upcycling par embeddings statiques deviendra une alternative sérieuse au MoE pour étendre des modèles déjà entraînés.