# PRISM: Demystifying Retention and Interaction in Mid-Training

> Mid-training with ~27B tokens boosts math, code, and science scores by up to 40 points and is critical for RL, which alone fails on base models.

- **Source:** [arXiv](https://arxiv.org/abs/2603.17074)
- **Published:** 2026-08-17
- **Permalink:** https://picx.dev/p/dHgQ0Z
- **Whiteboard:** https://picx.dev/p/dHgQ0Z/image

## Summary

# PRISM: Demystifying Retention and Interaction in Mid-Training

## Summary (Overview)

- **Comprehensive empirical study**: PRISM systematically investigates mid-training design choices for LLMs across 7 base models, spanning 4 families (Granite, LLaMA, Mistral, Nemotron-H), 2 architecture types (dense Transformer, attention-Mamba hybrid), and scales from 3B to 24B parameters.
- **Key findings**: Mid-training with ~27B tokens yields consistent gains of **+15 to +40 points on math, +5 to +12 points on code, and +6 to +13 points on science (GPQA-Diamond)**, while preserving general performance.
- **Mid-training is critical for RL**: The PRISM → RL pipeline improves the macro-average across six reasoning benchmarks from under 12 to 29–42 (a 3–4× improvement), whereas RL applied directly to base models is substantially less effective, with AIME scores near zero.
- **Data composition matters most at mid-training**: Including science data during mid-training unlocks +17 to +28 point GPQA-Diamond gains during RL, while changing the RL mix produces <2 point differences.
- **Mechanistic difference**: Mid-training densely restructures >90% of model weights, while RL makes sparse, front-loaded refinements to ~5% of parameters. RL consistently preserves mid-training's representational geometry (>0.998 CKA) across both dense Transformers and hybrid architectures.

## 3. Introduction and Theoretical Foundation

### 3.1 Background

The training pipeline for LLMs has evolved beyond the traditional two-stage recipe of pre-training followed by alignment. State-of-the-art models now incorporate an additional intermediate stage, **mid-training**, in which higher-quality, domain-focused data mixtures are used to imbue reasoning capabilities before downstream fine-tuning and RL.

### 3.2 Limitations of Prior Mid-Training Approaches

The term "mid-training" has been used inconsistently in the literature, referring to long-context extension, data annealing, and domain-specific capability refinement. Key limitations of prior work:

- **Narrow evaluation hides regressions**: Many studies report gains on a limited set of domain-specific benchmarks without assessing whether these improvements preserve general-purpose capabilities.
- **Interaction with RL remains underexplored**: There is a lack of controlled investigation into how mid-training interacts with downstream optimization, particularly RL.
- **Concurrent work**: Recent studies (e.g., Liu et al., 2025; Zhang et al., 2025) highlight mid-training as a critical yet underexplored stage.

### 3.3 Theoretical Motivation

The central hypothesis is that mid-training places the model in a weight configuration from which RL can effectively improve performance. PRISM provides empirical evidence for this through controlled experiments, weight-level analysis, and representation analysis. The paper introduces the concept of **retention-aware mid-training**, which balances domain specialization with preservation of general capabilities.

## 3. Methodology

### 3.1 Data Mixtures for Mid-Training

The datasets used for mid-training are summarized in Table 1. For Math and Code domains, both general web documents and domain-specific reasoning datasets are used. For Science, only reasoning-focused datasets are included. Chat and instruction-following datasets are also included to stabilize RL.

| Dataset | Type | Tokens (B) |
|---|---|---|
| DCLM-EDU | General web data | 111.46 |
| Open-R1 (MoT) | Math reasoning | 0.60 |
| Nemotron Post-Training v1 | Math | 35.93 |
| Megamath-Web-Pro | Math web | 14.73 |
| Open-R1 (MoT) | Code reasoning | 1.18 |
| OpenCodeReasoning-2 | Code reasoning | 1.12 |
| RefinCode | Code web | 186.44 |
| StarCoder2 | Code web | 432.73 |
| Open-R1 (MoT) | Science reasoning | 0.42 |
| OpenThoughts3 | Science reasoning | 0.73 |
| WildChat-1M | Chat | - |
| Tulu-3 SFT Personas | Chat | 0.91 |
| UltraChat-200k | Chat | - |

**Table 1**: Datasets used in mid-training mixtures.

### 3.2 Benchmark Selection

The evaluation suite combines general leaderboards (LB-V1, LB-V2) with focused long-context, code, math, and science evaluations to surface both gains and regressions.

| Category | Benchmarks | What it measures |
|---|---|---|
| General ability | LB-V1 (ARC, HellaSwag, MMLU, TruthfulQA, Winogrande, GSM8K), LB-V2 (IFEval, BBH, MATH, GPQA, MUSR, MMLU-Pro) | Broad multitask knowledge and robustness |
| Long-context | RULER | Long-context retrieval |
| Code | LiveCodeBench, Codeforces | Executable program synthesis and reasoning |
| Math | AIME, MATH500 | Mathematical reasoning |
| Science | GPQA-Diamond | Expert-level scientific reasoning |

**Table 2**: Benchmark categories recommended for evaluating mid-training design choices.

### 3.3 Mid-Training Configuration

- **Token budget**: ~27B tokens
- **Context length**: 8k (with ablations at 16k and 32k)
- **Architectures**: Dense Transformer and attention-Mamba hybrid
- **Models**: Granite-3.3 (8B), Granite-4 Micro (3B), Granite-4-H Micro (3B), Nemotron-H (8B), Mistral-7B, Mistral-Small (24B), LLaMA-3.1 (8B)

### 3.4 RL Setup

- **Unbalanced mix**: Math prompts with exactly one correct sample out of 16 (hard subset); code and science prompts with at least one correct sample.
- **Balanced mix**: Equalized domains to 19k prompts each, with randomized instruction-format templates for science.

## 4. Empirical Validation / Results

### 4.1 When to Mid-Train

Applying mid-training after long-context extension yields the strongest results. On Granite-4 Micro (3B), AIME24 increases from 9.38 (Phase 3) to 23.59 (after LC), and MATH500 rises to 77.44.

### 4.2 Domain-wise Effects of Mid-Training Data

| 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 |

**Table 5**: Domain-specific results for Granite-3.3 (8B). Math data drives the largest single-domain gains (+27.48 points), code data is essential for programming benchmarks, and science data improves GPQA without sacrificing other domains.

### 4.3 PRISM Effects Across Model Families

| Model | Variant | LB V1 | LB V2 | Code Avg | GPQA-D | Math Avg |
|---|---|---|---|---|---|---|
| Granite-3.3 (8B) | Base | 66.15 | 20.10 | 2.07 | 22.56 | 8.95 |
| Granite-3.3 (8B) | PRISM | 66.48 | 23.03 | 10.58 | 29.12 | 48.75 |
| Mistral-7B | Base | 60.88 | 14.89 | 0.07 | 26.94 | 0.59 |
| Mistral-7B | PRISM | 59.99 | 19.68 | 9.79 | 24.07 | 41.28 |
| LLaMA-3.1 (8B) | Base | 62.76 | 14.09 | 0.04 | 20.20 | 2.24 |
| LLaMA-3.1 (8B) | PRISM | 65.21 | 21.46 | 5.77 | 21.04 | 36.41 |
| Mistral-Small (24B) | Base | 74.98 | 27.29 | 0.15 | 22.55 | 9.48 |
| Mistral-Small (24B) | PRISM | 69.52 | 27.42 | 10.06 | 22.05 | 47.02 |

**Table 6**: Base vs. PRISM (Math+Code+Science) mid-training results across model families.

### 4.4 The Full Pipeline: Broader RL Analysis

The full PRISM → RL pipeline improves the six-benchmark macro-average from under 12 to 29–42, a 3–4× improvement. The best configuration per model is:

- **LLaMA-3.1 (8B)**: MCS mid-training + MCS RL → AVG 29.69
- **Granite-3.3 (8B)**: MCS mid-training + MC RL → AVG 41.76
- **Mistral-7B**: MCS mid-training + MCS RL → AVG 35.14

### 4.5 Weight-Level Analysis: Dense Restructuring vs Sparse Refinement

- **Mid-training**: 90.7% of attention and 98.1% of MLP parameters change significantly (Granite-3.3). L2 divergence of 0.175 (attention) and 0.329 (MLP).
- **RL**: Only ~5% of parameters modified. L2 divergence is 580× smaller for Granite-3.3 (0.0003 vs. 0.175) and 370× smaller for Nemotron-H (0.0003 vs. 0.112).

### 4.6 Representation Analysis (CKA)

| Model | Arch. | Wiki | C4 | GSM8K |
|---|---|---|---|---|
| Granite-3.3 (8B) | Dense | $0.9999 \pm 0.0000$ | $0.9999 \pm 0.0000$ | $0.9997 \pm 0.0000$ |
| LLaMA-3.1 (8B) | Dense | $0.9999 \pm 0.0000$ | $0.9999 \pm 0.0000$ | $0.9996 \pm 0.0001$ |
| Nemotron-H (8B) | Hybrid | $0.9999 \pm 0.0000$ | $0.9998 \pm 0.0000$ | $0.9993 \pm 0.0001$ |

**Table 14**: MT vs. RL representational similarity (minimum linear CKA) across input distributions. RL consistently preserves mid-training's representational geometry (>0.998) across all models.

### 4.7 Behavioral Analysis

| Model | Stage | Pass | Med. Len | Neg-LP |
|---|---|---|---|---|
| Granite-3.3 (8B) | Base | 16.9% | 120 | 0.382 |
| Granite-3.3 (8B) | MT | 75.5% | 2,254 | 0.138 |
| Granite-3.3 (8B) | RL | 79.5% | 1,700 | 0.141 |
| LLaMA-3.1 (8B) | Base | 2.6% | 158 | 0.758 |
| LLaMA-3.1 (8B) | MT | 43.1% | 1,052 | 0.377 |
| LLaMA-3.1 (8B) | RL | 64.6% | 1,188 | 0.267 |
| Nemotron-H (8B) | Base | 66.6% | 452 | 0.167 |
| Nemotron-H (8B) | MT | 61.6% | 1,928 | 0.150 |
| Nemotron-H (8B) | RL | 83.0% | 1,780 | 0.127 |

**Table 15**: Correctness, response length, and prediction confidence across pipeline stages on 200 held-out MATH500 problems. The PRISM → RL pipeline consistently achieves the highest pass rates across all three model families.

## 5. Theoretical and Practical Implications

### 5.1 Key Findings

- **Data composition matters most at mid-training**: Changing the mid-training mix from Math+Code to Math+Code+Science shifts AVG by +3 to +6 points, while changing the RL mix produces <2 point differences.
- **Science data at mid-training unlocks large RL gains on GPQA**: For Granite-3.3, MCS mid-training + MC RL achieves GPQA 52.86 (vs. 35.52 with MC mid-training + MC RL).
- **Mid-training and RL operate through different mechanisms**: Mid-training densely restructures >90% of parameters, while RL sparsely refines ~5%. RL applies identical weight changes regardless of starting point, yet only succeeds on mid-trained models.
- **RL expands the solvability frontier**: The model progressively learns to solve prompts it could not handle at the start of RL, with non-saturating training curves.

### 5.2 Practical Guidance for Mid-Training Pipelines

1. **Use a balanced benchmark suite**: Combine general leaderboards with domain-specific benchmarks to expose global regressions.
2. **Measure long-context retention explicitly**: Mid-training at 8k context can degrade long-context capabilities; a brief extension phase can restore performance.
3. **Include general web data**: The consistent inclusion of DCLM-EDU helps preserve general capabilities.
4. **Science data is valuable**: Including science data during mid-training unlocks large GPQA-Diamond gains during RL.
5. **Mid-training is a launch point for RL**: The full pipeline yields 3–4× improvements over direct RL on base models.

### 5.3 Limitations and Future Directions

- **Model-specific RL data selection**: For consistency, RL prompts were filtered using a single model (Granite-3.3-8B); adaptive RL data curation is a natural next step.
- **Broader domain coverage**: Extending PRISM to multilingual reasoning, agentic tasks, and tool use.
- **Scaling beyond 24B**: Verifying findings at 70B+ scale.
- **Long-context mid-training**: Exploring mid-training at 32k+ with proportional token budgets.

## 6. Conclusion

PRISM demonstrates that retention-aware mid-training is a highly effective intermediate step for reliable reasoning enhancement and RL scaling. The key takeaway is that mid-training provides a "launch point" for RL by placing the model in a weight configuration from which RL can effectively improve performance, even though the exact causal mechanism remains to be established. The paper provides practical guidelines for designing robust mid-training pipelines, including data composition, context length, and token budget recommendations.

---

_Markdown view of https://picx.dev/p/dHgQ0Z, served by PicX — AI-generated visual whiteboard summaries of research papers._