PRISM : Démystifier la rétention et l'interaction en milieu d'entraînement

PRISM: Demystifying Retention and Interaction in Mid-Training

March 17, 2026

Auteurs: Bharat Runwal, Ashish Agrawal, Anurag Roy, Rameswar Panda

cs.AI

Résumé

Nous présentons PRISM, une étude empirique exhaustive des choix de conception en mi-entraînement pour les grands modèles de langage. Par des expériences contrôlées sur sept modèles de base couvrant quatre familles (Granite, LLaMA, Mistral, Nemotron-H), deux types d'architecture (Transformers dense et hybride attention-Mamba) et des échelles de 3 à 24 milliards de paramètres, nous montrons qu'un mi-entraînement sur environ 27 milliards de tokens de haute qualité produit des gains constants de +15 à +40 points en mathématiques, +5 à +12 points en code et +6 à +13 points sur des benchmarks scientifiques, tout en préservant les performances générales. Le pipeline complet PRISM vers RL améliore la moyenne macro sur six benchmarks de raisonnement de moins de 12 à 29-42 (une amélioration de 3 à 4x), tandis que le RL appliqué directement à la plupart des modèles de base reste nettement moins efficace, avec des scores AIME proches de zéro. La composition des données est plus déterminante lors du mi-entraînement que lors du RL : inclure des données scientifiques pendant le mi-entraînement permet des gains de +17 à +28 points sur GPQA-Diamond lors du RL, tandis que modifier le mélange de RL produit des différences inférieures à 2 points. Mécaniquement, le mi-entraînement restructure densément plus de 90 % des poids du modèle, tandis que le RL effectue des ajustements épars et frontaux sur environ 5 % des paramètres. L'analyse des représentations (CKA) confirme que le RL préserve systématiquement la géométrie représentationnelle du mi-entraînement (CKA supérieur à 0,998) across les architectures. Fait crucial, le RL applique des modifications de poids identiques quel que soit le point de départ, mais ne réussit que sur les modèles ayant subi un mi-entraînement, ce qui concorde avec l'hypothèse que le mi-entraînement place le modèle dans une configuration à partir de laquelle le RL peut efficacement améliorer les performances. Nos résultats démontrent que le mi-entraînement avec conservation des connaissances est très efficace pour l'amélioration fiable du raisonnement et fournissent des conseils pratiques pour concevoir des pipelines de mi-entraînement robustes.

English

We present PRISM, a comprehensive empirical study of mid-training design choices for large language models. Through controlled experiments across seven base models spanning four families (Granite, LLaMA, Mistral, Nemotron-H), two architecture types (dense Transformer and attention-Mamba hybrid), and scales from 3B to 24B parameters, we show that mid-training on approximately 27B high-quality tokens yields consistent gains of +15 to +40 points on math, +5 to +12 points on code, and +6 to +13 points on science benchmarks while preserving general performance. The full PRISM to RL pipeline improves macro-average across six reasoning benchmarks from under 12 to 29-42 (a 3-4x improvement), whereas RL applied directly to most of the base models remains substantially less effective, with AIME scores near zero. Data composition matters most at mid-training, not RL: including science data during mid-training unlocks +17 to +28 point GPQA-Diamond gains during RL, while changing the RL mix produces less than 2 point differences. Mechanistically, mid-training densely restructures over 90% of model weights, while RL makes sparse, front-loaded refinements to approximately 5% of parameters. Representation analysis (CKA) confirms that RL consistently preserves mid-training's representational geometry (over 0.998 CKA) across architectures. Crucially, RL applies identical weight changes regardless of starting point, yet only succeeds on mid-trained models, consistent with mid-training placing the model in a configuration from which RL can effectively improve performance. Our results demonstrate that retention-aware mid-training is highly effective for reliable reasoning enhancement and provide practical guidance for designing robust mid-training pipelines.