PRISM: Demystifying Retention and Interaction in Mid-Training
Organizations: IBM Research, MIT-IBM Watson AI Lab
Abstract
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.
Figures & tables
| Dataset | Type | Tokens (B) |
| DCLM-EDU Allal et al. (2025) | General web data | 111.46 |
| Open-R1 (MoT) Lozhkov et al. (2025) | Math reasoning | 0.60 |
| Nemotron Post-Training v1 Nathawani et al. (2025) | Math | 35.93 |
| Megamath-Web-Pro Zhou et al. (2025) | Math web | 14.73 |
| Open-R1 (MoT) Penedo et al. (2025) | Code reasoning | 1.18 |
| OpenCodeReasoning-2 Ahmad et al. (2025) | Code reasoning | 1.12 |
| Category | Benchmarks | What it measures | Why it matters |
|---|---|---|---|
| General ability | Leaderboard-V1 (LB-V1) (ARC, HellaSwag, MMLU, TruthfulQA, Winogrande, GSM8K), Leaderboard-V2 (LB-V2) (IFEval, BBH, MATH, GPQA, MUSR, MMLU-Pro) | Broad multitask knowledge and robustness | Detects generalization regressions hidden by domain-specific gains. |
| Long-context | RULER | Long-context retrieval | Ensures mid-training does not degrade long-context retrieval capabilities. |
| Code | LiveCodeBench Jain et al. (2024) , Codeforces Penedo et al. (2025) | Executable program synthesis and reasoning | Captures real-world coding ability. |
| Math | AIME AIME () , MATH500 Lightman et al. (2023) | Mathematical reasoning | Highly sensitive to data quality and mid-training composition. |
| Science | GPQA-Diamond Rein et al. (2023) | Expert-level scientific reasoning | Probes scientific reasoning capabilities |
| Leaderbds. | Code | Sci. | Math | |||||
|---|---|---|---|---|---|---|---|---|
| Stage | V1 | V2 | LCB | CF | GPQA | AI24 | AI25 | M500 |
| Phase 3 | 63.30 | 19.44 | 7.05 | 8.61 | 19.53 | 9.38 | 16.09 | 65.88 |
| Phase 4 | 62.84 | 20.85 | 7.89 | 7.95 | 17.85 | 10.00 | 14.06 | 61.70 |
| After LC | 62.91 | 20.53 | 10.39 | 6.18 | 25.93 | 23.59 | 20.94 | 77.44 |
| Leaderboard V1 | Leaderboard V2 | |||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| Mixture | ARC | HellaSwag | MMLU | TruthfulQA | Winogrande | GSM8K | OpenLLM V1 Avg | IFEval | BBH | MATH | GPQA | MUSR | MMLU-Pro | OpenLLM V2 Avg |
| Base | 61.95 | 83.46 | 62.56 | 52.24 | 80.35 | 56.33 | 66.15 | 46.62 | 24.68 | 10.20 | 6.38 | 8.88 | 23.82 | 20.10 |
| Math only | 62.54 | 78.72 | 64.29 | 46.04 | 75.30 | 71.95 | 66.47 | 46.46 | 25.57 | 17.75 | 5.59 | 9.08 | 29.86 | 22.39 |
| Math + Code | 61.01 | 78.09 | 62.65 | 47.36 | 74.74 | 73.46 | 66.22 | 45.56 | 26.87 | 18.43 | 5.93 | 10.60 | 28.40 | 22.63 |
| Math + Code + Science | 61.69 | 78.12 | 62.98 | 46.96 | 74.90 | 74.22 | 66.48 | 46.44 | 26.32 | 20.02 | 7.27 | 8.60 | 29.55 | 23.03 |
| Mixture | Code | Math | GPQA |
|---|---|---|---|
| Base | 2.07 | 8.95 | 22.56 |
| Math | 2.81 | 36.43 | 17.34 |
| Math+Code | 10.71 | 44.99 | 19.02 |
| Math+Code+Sci | 10.58 | 48.75 | 29.12 |
| Leaderboards | Code | Science | Math | ||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|
| Model | Variant | LB V1 | LB V2 | LCB | CF | Code Avg | GPQA-D | AIME24 | AIME25 | MATH500 | Math Avg |
| Granite-3.3 (8B) | Base | 66.15 | 20.10 | 2.15 | 1.99 | 2.07 | 22.56 | 0.46 | 0.31 | 26.09 | 8.95 |
| PRISM | 66.48 | 23.03 | 10.63 | 10.52 | 10.58 | 29.12 | 37.18 | 27.96 | 81.11 | 48.75 | |
| Granite-4 Micro (3B) | Base | 66.01 | 21.82 | 0.24 | 2.28 | 1.26 | 21.55 | 16.09 | 12.34 | 50.42 | 26.28 |
| PRISM | 62.91 | 20.53 | 10.87 | 6.25 | 8.56 | 34.34 | 27.19 | 22.29 | 79.40 | 42.96 | |
| Granite-4-H Micro (3B) | Base | 64.49 | 18.99 | 0.60 | 0.88 | 0.74 | 20.88 | 7.08 | 2.70 | 30.17 | 13.32 |
| RULER | Code / Science | Math | |||||||||||
| Model Variant | 8k | 16k | 32k | 64k | 128k | LCB | CF | Code Avg | GPQA-D | AIME24 | AIME25 | MATH500 | Math Avg |
| Granite-3.3 Base | 85.81 | 82.40 | 75.53 | 64.91 | 59.09 | 2.15 | 1.99 | 2.07 | 22.56 | 0.46 | 0.31 | 26.09 | 8.95 |
| Mid-Train (Math+Code) | 89.02 | 60.44 | 21.52 | 11.71 | 6.46 | 11.11 | 10.30 | 10.71 | 19.02 | 32.44 | 28.33 | 74.22 | 44.99 |
| Mid-Train + LC (Attention) | 90.04 | 82.56 | 71.47 | 54.63 | 36.32 | 23.78 | 15.53 | 19.65 | 17.85 | 36.56 | 32.55 | 67.20 | 45.44 |
| Mid-Train + LC (Full) | 89.29 | 80.74 | 70.86 | 56.02 | 38.41 | 29.99 | 21.04 | 25.52 | 14.48 | 35.21 | 30.36 | 62.30 | 42.62 |
| Merge (15% Base + 85% Mid-Train) | 89.12 | 69.76 | 32.63 | 15.44 | 11.32 | 10.75 | 10.96 | 10.86 | 22.22 | 28.39 | 24.90 | 72.97 | 42.09 |
| Context | LB-V1 | LB-V2 | LCB | CF | GPQA | AIME24 | AIME25 | M500 |
|---|---|---|---|---|---|---|---|---|
| Base | 66.01 | 21.82 | 0.24 | 2.28 | 21.55 | 16.09 | 12.34 | 50.42 |
| 8k | 62.91 | 20.53 | 10.87 | 6.25 | 34.34 | 27.19 | 22.29 | 79.40 |
| 16k | 64.23 | 20.37 | 12.19 | 8.90 | 38.89 | 31.82 | 25.26 | 82.47 |
| 32k | 64.48 | 21.05 | 14.93 | 7.50 | 39.89 | 30.98 | 21.87 | 82.70 |
| Tok. (B) | LB-V1 | LB-V2 | Code | GPQA | Math |
|---|---|---|---|---|---|
| Base | 66.01 | 21.82 | 1.26 | 21.55 | 26.28 |
| 10.49 | 63.45 | 19.50 | 9.59 | 19.19 | 40.21 |
| 15.73 | 63.24 | 19.79 | 9.02 | 23.06 | 42.07 |
| 26.21 | 63.28 | 19.63 | 8.69 | 19.19 | 42.22 |
| 31.46 | 63.16 | 20.05 | 7.62 | 21.38 | 42.42 |
| Domain | Sources | Count |
|---|---|---|
| Math | DeepScaleR-Preview | 294K |
| INTELLECT-2-RL | ||
| Skywork-OR1-RL-Data | ||
| Science | Nemotron-PT-v1-stem | 100K |
| Code | DeepCoder-Preview | 142K |
| Skywork-OR1-RL-Data |
| Model | MT | RL | LCB | CF | Code Avg | AIME24 | AIME25 | MATH500 | Math Avg | GPQA | AVG |
|---|---|---|---|---|---|---|---|---|---|---|---|
| LLaMA-3.1 | – | – | 0.00 | 0.07 | 0.04 | 0.05 | 0.15 | 6.51 | 2.24 | 20.20 | 7.49 |
| MC | – | 6.93 | 6.03 | 6.48 | 20.67 | 19.58 | 73.70 | 37.98 | 19.53 | 21.33 | |
| MCS | – | 6.09 | 5.45 | 5.77 | 16.45 | 19.32 | 73.47 | 36.41 | 21.04 | 21.07 | |
| MC | MC | 12.31 | 11.85 | 12.08 | 25.47 | 23.23 | 78.99 | 42.56 | 23.06 | 25.90 | |
| MC | MCS | 11.83 | 12.80 | 12.32 | 24.43 | 23.12 | 78.62 | 42.06 | 24.75 | 26.38 | |
| MCS | MC | 13.62 | 11.41 | 12.51 | 20.47 | 21.67 | 77.10 | 39.75 | 34.01 | 28.76 |
| Model | MT | Transition | Attn | MLP | Mamba | Total | Sparsity |
|---|---|---|---|---|---|---|---|
| Granite-3.3 (8B) | MCS | Base MT | 0.175 | 0.329 | – | 0.175 | 9.3% |
| MT RL | 0.0003 | 0.0006 | – | 0.0003 | 95.9% | ||
| Base RL (no MT) | 0.0004 | 0.0007 | – | 0.0004 | 96.0% | ||
| MC | Base MT | 0.177 | 0.333 | – | 0.177 | 9.3% | |
| MT RL | 0.0003 | 0.0006 | – | 0.0003 | 95.8% | ||
| Nemotron-H (8B, Hybrid) | MCS | Base MT | 0.230 | 0.289 | 0.138 | 0.112 | 2.7% |
| Model | Mix | Attn | MLP | Mamba | Total |
|---|---|---|---|---|---|
| Granite-3.3 (8B) | MC | 0.177 | 0.333 | – | 0.177 |
| MCS | 0.175 | 0.329 | – | 0.175 | |
| Nemotron-H (8B) | MC | 0.232 | 0.292 | 0.140 | 0.113 |
| MCS | 0.230 | 0.289 | 0.138 | 0.112 |
| Model | Arch. | Wiki | C4 | GSM8K |
|---|---|---|---|---|
| Granite-3.3 (8B) | Dense | 0.9999 0.0000 | 0.9999 0.0000 | 0.9997 0.0000 |
| LLaMA-3.1 (8B) | Dense | 0.9999 0.0000 | 0.9999 0.0000 | 0.9996 0.0001 |
| Nemotron-H (8B) | Hybrid | 0.9999 0.0000 | 0.9998 0.0000 | 0.9993 0.0001 |
| Model | Stage | Pass | Med. Len | Neg-LP | Corr. | Incorr. |
|---|---|---|---|---|---|---|
| Granite-3.3 (8B) | Base | 16.9% | 120 | 0.382 | – | 0.383 |
| MT | 75.5% | 2,254 | 0.138 | 0.128 | 0.153 | |
| RL | 79.5% | 1,700 | 0.141 | 0.135 | 0.160 | |
| LLaMA-3.1 (8B) | Base | 2.6% | 158 | 0.758 | – | 0.780 |
| MT | 43.1% | 1,052 | 0.377 | 0.146 | 0.469 | |
| RL | 64.6% | 1,188 | 0.267 | 0.149 | 0.320 |
| Dataset | Type | Tokens (B) |
| DCLM-EDU | General Web Data | 2.51 |
| Nemotron Post-Training v1 | Math (QA/Reasoning) | 5.08 |
| Megamath-Web-Pro | Math (web) | 4.33 |
| StarCoder2 | Code (web) | 37.52 |
| xenArcAI-codex | Code (QA/Reasoning) | 3.860 |
| Code | Science | Math | ||||||
|---|---|---|---|---|---|---|---|---|
| Model / Mixture | LiveCodeBench | Codeforces | Code Avg | GPQA-Diamond | AIME 24 | AIME 25 | MATH 500 | Math Avg |
| Base | 2.15 | 1.99 | 2.07 | 22.56 | 0.46 | 0.31 | 26.09 | 8.95 |
| Math only | 2.15 | 3.46 | 2.81 | 17.34 | 26.72 | 22.08 | 60.50 | 36.43 |
| Math + Code | 11.11 | 10.30 | 10.71 | 19.02 | 32.44 | 28.33 | 74.22 | 44.33 |
| Math + Code + Science | 10.63 | 10.52 | 10.58 | 29.12 | 37.18 | 27.96 | 81.11 | 48.75 |
| Token Budget (B) | LB V1 | LB V2 | LCB | CF | Code Avg | GPQA-D | AIME24 | AIME25 | MATH500 | Math Avg |
|---|---|---|---|---|---|---|---|---|---|---|
| Base | 66.01 | 21.82 | 0.24 | 2.28 | 1.26 | 21.55 | 16.09 | 12.34 | 50.42 | 26.28 |
| 10.49 | 63.45 | 19.50 | 10.51 | 8.68 | 9.59 | 19.19 | 23.95 | 19.17 | 77.52 | 40.21 |
| 15.73 | 63.24 | 19.79 | 10.75 | 7.28 | 9.02 | 23.06 | 26.14 | 21.30 | 78.76 | 42.07 |
| 26.21 | 63.28 | 19.63 | 9.80 | 7.58 | 8.69 | 19.19 | 28.49 | 20.10 | 78.08 | 42.22 |
| 31.46 | 63.16 | 20.05 | 8.24 | 6.99 | 7.62 | 21.38 | 28.02 | 22.08 | 77.15 | 42.42 |
| Hyperparameter | Value |
|---|---|
| Sampling Parameters | |
| Responses per Prompt | 16 |
| Temperature | 1.0 |
| Top- | 1.0 |
| Training Dynamics | |
| Total Batch Size | 1024 |
| Threshold (%) | Base MT unchanged | MT RL unchanged |
|---|---|---|
| 0.1 | 2.1% | 82.3% |
| 0.5 | 2.4% | 87.3% |
| 1.0 | 3.5% | 95.0% |
| 2.0 | 5.3% | 97.5% |
| 5.0 | 11.1% | 99.0% |
| 10.0 | 20.1% | 99.5% |
| Granite-3.3 (8B) | Nemotron-H (8B) | |||||||
| Sub-component | Base MT | MT RL | Base RL | Base MT | MT RL | Base RL | ||
| Attn: Q | 83.5% | 4.4% | 5.2% | +0.8 | 97.5% | 5.5% | 4.9% | 0.6 |
| Attn: K | 83.2% | 4.6% | 5.2% | +0.6 | 96.9% | 5.5% | 4.8% | 0.7 |
| Attn: V | 97.9% | 5.7% | 7.5% | +1.8 | 97.1% | 8.5% | 7.6% | 0.9 |
| Attn: O | 98.0% | 5.6% | 6.7% | +1.1 | 97.2% | 7.2% | 6.4% | 0.8 |
| MLP: gate | 98.3% | 5.4% | 6.1% | +0.7 | – | – | – | – |
LLM Training
Large language model training studies the data, objectives, optimizers, and systems used to build and adapt language models.
See all 342 papers