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.

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.

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

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

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)

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

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.