PaperlayerView exact revision on arXiv

Source-linked AI summary

PRISM: Demystifying Retention and Interaction in Mid-Training

Bharat Runwal, Ashish Agrawal, Anurag Roy, Rameswar Panda

arXiv:2603.17074v3cs.LG

TL;DR

PRISM addresses the limited systematic understanding of mid-training by studying its data choices, timing, RL interaction, and architectural generality. Controlled experiments show that retention-aware mid-training improves reasoning and makes downstream RL substantially more effective while preserving general-purpose performance.

  • Problem

    Mid-training is increasingly used to add reasoning capabilities, but systematic guidance on its data, timing, interaction with RL, and cross-architecture generality remains limited.

  • Method

    PRISM uses a retention-aware empirical framework spanning multiple domains, benchmark axes, model families, architectures, and downstream RL behavior.

  • Results

    The full PRISM→RL pipeline raises macro-average performance across six reasoning benchmarks from under 12 to 29–42, while RL applied directly to base models remains substantially less effective.

  • Takeaways & Limitations

    Mid-training is the primary leverage point for reliable reasoning enhancement, with data composition mattering more during mid-training than during RL.

  • Takeaways & Limitations

    Whether the reported ordering generalizes across larger models or different architectures remains an open question.

Abstract

from arXiv · show

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.

1 Introduction

PRISM addresses the limited systematic understanding of mid-training through controlled, cross-model experiments. It finds that targeted mid-training improves reasoning, strengthens downstream RL, and produces distinct mechanistic changes across architectures and scales.

  • Study scope: PRISM evaluates mid-training across seven models from four families, two architecture types, and 3B–24B parameter scales using approximately 27B high-quality tokens.The study decomposes design choices across retention, domain interaction, benchmark selection, RL compatibility, and scaling behavior.
  • Key findings: +15 to +40 points on math, +5 to +12 points on code, and +6 to +13 points on science benchmarks, while preserving general-purpose performance.Science gains are reported for Granite and hybrid models.
  • Key findings: The full PRISM→RL pipeline raises macro-average performance across six reasoning benchmarks from under 12 to 29–42, a 3–4× improvement.RL applied directly to base models is substantially less effective, with AIME scores near zero.
  • Data and RL interaction: Science data during mid-training unlocks +17 to +28 point GPQA-Diamond gains during RL, whereas changing the RL mix produces less than 2 point differences.Changing the mid-training mix from Math+Code to Math+Code+Science shifts AVG1 by +3 to +6 points.
  • Mechanisms: RL progressively solves initially unsolvable prompts on PRISM-mid-trained Granite-3.3 models, with non-saturating training curves across hundreds of steps.This finding supports the paper’s analysis of RL expanding the solvability frontier.
  • Mechanisms: Mid-training densely restructures over 90% of parameters, while RL sparsely refines approximately 5% and preserves representational geometry above 0.998 CKA.RL weight changes are front-loaded in the first approximately 200–400 steps, and the full pipeline raises held-out MATH500 pass rates from 2.6–66.6% to 64.6–83.0%.

2 Limitations of Prior Mid-Training Approaches

Prior mid-training research lacks unified terminology, standardized evaluation, and controlled analysis of interactions with downstream RL. These gaps motivate PRISM’s retention-aware, cross-domain and cross-family framework.

  • Conceptual limitations: Mid-training has been used for long-context extension, data annealing, and domain-specific refinement without a unified framework or standardized evaluation.The inconsistent usage limits direct comparison across studies.
  • Evaluation limitations: Narrow domain evaluations can hide regressions in general-purpose capabilities and fail to measure cross-domain reasoning robustness.Prior work often emphasizes math or code while omitting broad generalization and other reasoning dimensions.
  • RL limitations: Controlled investigation of how mid-training interacts with downstream RL remains limited across model families, domains, and benchmark suites.Existing claims about RL facilitation are typically tested in narrow experimental settings.
  • Motivation: PRISM responds with a retention-aware framework spanning multiple domains, benchmark axes, and downstream RL behavior across model families.Its purpose is to uncover trade-offs overlooked by prior work.

3 Data Mixtures for Mid-Training

PRISM uses empirically tuned mixtures that combine broad web and instruction data with domain-specific reasoning sources. Three progressively richer configurations support targeted capability gains while retaining broad pretraining knowledge.

  • Mixture design: PRISM uses general web documents for retained knowledge, domain-specific reasoning datasets for problem-solving, and chat or instruction data to stabilize downstream RL.Science uses reasoning-focused datasets, while Math and Code combine general web and domain-specific sources.
  • Preprocessing: Data quality is supported by deterministic preprocessing, quality filtering for web data, evaluator-based filtering for OpenCodeReasoning-2, and normalization of chat conversations.The supplied passages describe these preprocessing choices as part of the data construction pipeline.
  • Mixture design: The three configurations are Math-only, Math+Code, and Math+Code+Science, with progressively richer domain coverage.Figure 2 reports per-source sampling weights for these configurations.
  • Mixture selection: The adopted sampling splits provide the best observed balance between retaining broad pretraining knowledge and inducing targeted domain improvements.These empirically selected splits become the default policy for all reported experiments.

4 What to Evaluate: Benchmark Selection

PRISM evaluates mid-training with a broad benchmark suite designed to reveal both domain gains and capability regressions. The setup combines general-purpose, long-context, and targeted reasoning evaluations.

  • Evaluation scope: PRISM combines general leaderboards with focused long-context, code, math, and science evaluations so single-domain gains cannot conceal capability loss elsewhere.Table 2 organizes these benchmark categories by evaluation role.
  • Evaluation scope: Effective evaluation requires both breadth and depth: general-purpose leaderboards expose global regressions, while targeted benchmarks measure domain-specific gains.The recommended setup explicitly mixes broad and focused evaluations.
  • Long-context retention: Long-context reasoning should be measured separately because short-context-dominated mid-training can degrade long-context capabilities.A lightweight fine-tuning stage may be needed to recover such performance.

5 When to Mid-Train

PRISM is most effective when mid-training follows long-context extension, producing the strongest combined gains across math, code, and science while maintaining stable general-purpose performance. This timing result is preliminary because it was established on one 3B model and may be confounded by differences in starting-model quality.

  • Timing results: 19.44 to 20.85 Leaderboard V2 improvement occurs from Phase 3 to Phase 4, while both phases underperform the final base model on math and science.
  • Timing results: After long-context extension, mid-training yields the strongest overall performance across math, code, and science.AIME24 rises to 23.59, MATH500 reaches 77.44, LiveCodeBench reaches 10.39, and GPQA-Diamond reaches 25.93.
  • Timing results: 23.59 AIME24 after long-context extension exceeds 9.38 after Phase 3 and 10.00 after Phase 4.
  • General capabilities: General-purpose leaderboards remain relatively stable across timing choices, indicating no large broad-capability regressions from later mid-training.
  • Scope and caveat: The timing conclusion is preliminary because experiments use only Granite-4 Micro (3B), and post-long-context starting models may be stronger in absolute terms.

6 Domain-wise Effects of Mid-Training Data

PRISM’s domain-specific mixtures produce broad reasoning gains across model families, while balanced coverage preserves general performance more effectively than narrow mixtures. Short-context mid-training can damage long-context ability, but targeted restoration and moderate context lengths recover or improve the overall trade-off.

  • Domain-wise data effects: Balanced mixtures across math, code, and science achieve the best overall trade-off, improving domain reasoning while preserving general capabilities.Across model families, PRISM consistently improves mathematical, coding, and scientific reasoning, while leaderboard changes are smaller and more model dependent.
  • Domain-wise data effects: +27.48 points: Math-only mid-training raises Math average from 8.95 to 36.43.The result identifies math-specific data as the largest single-domain driver of mathematical reasoning gains.
  • Domain-wise data effects: +8.64 points: Math+Code raises Code average to 10.71 relative to the baseline.Math-only mid-training produces only a +0.74 code gain, while adding science maintains a similar Code average of 10.58.
  • Domain-wise data effects: +10.10 points: adding science to Math+Code increases GPQA-Diamond from 19.02 to 29.12 without sacrificing math or code gains.Math average rises from 44.99 to 48.75, while Code average remains similar at 10.58 versus 10.71.
  • Cross-model effects: +15 to +40 points on math and +5 to +12 points on coding occur across all tested models, with +6 to +13 science gains on Granite and hybrid models.The gains generalize across families, architectures, and scales, although science improvements primarily emerge after RL when science data is included during mid-training.
  • Context length and restoration: Mid-training can reduce long-context performance, but 1k steps of long-context training and a 15% base/85% mid-trained merge recover RULER while retaining reasoning gains.For Granite-3.3, RULER at 128k rises from 6.46 to 38.41 with full-parameter restoration and to 42.16 after merging; the latter reaches Code Avg 25.54 and Math Avg 44.48.
  • Context length and budget: 16k context provides the most favorable balance between reasoning gains and training efficiency, while gains largely saturate beyond 16k.From 8k to 16k, MATH500 rises from 79.40 to 82.47 and GPQA-Diamond from 34.34 to 38.89; extending to 32k yields small or mixed additional changes.
  • Context length and budget: Beyond 26.21B tokens, gains largely saturate: Math Avg remains nearly constant from 42.22 to 42.42 while Code Avg declines from 8.69 to 7.62.General-purpose leaderboard scores remain stable across budgets, and GPQA-Diamond shows no consistent trend.

A.2 PRISM Training Details

This section documents PRISM’s training datasets, benchmark coverage, evaluation settings, and long-context sequence-packing procedure.

  • Training configuration: PRISM mid-training hyperparameters are summarized across models, with dataset inventories provided separately for the long-context extension phase.Tables 16 and 18 organize the training configuration and datasets used.
  • Training configuration: Long-context extension hyperparameters and datasets are listed in separate tables for reproducibility.Table 17 covers extension-phase hyperparameters, while Table 18 lists datasets and token counts.
  • Sequence packing: Long-context extension training uses Best-Fit Decreasing packing to construct fixed-length sequences while reducing truncation and wasted capacity.Long documents may be split into overlapping chunks before length-sorted greedy assignment.
  • Evaluation: Evaluation covers general ability, long-context reasoning, and contamination-free code performance through LB-V1, LB-V2, RULER, LiveCodeBench, and Codeforces.The benchmark suite isolates specialization, retention, and long-context capabilities.
  • Evaluation: Math and code evaluations use 32k-token generation limits, with task-specific sampling settings and dedicated evaluation frameworks.Math uses 64 samples at temperature 0.6; code uses 3 samples at temperature 0.7.

C Extended Results Tables

The extended-results section supplies detailed domain-ablation results and token-budget scaling results that complement the main text.

  • Extended results tables: Table 19 reports per-benchmark Granite-3.3-8B results across Math-only, Math+Code, and Math+Code+Science mid-training mixtures.The table expands the domain ablation summarized in the main text.

D RL Training Details

This section specifies the RL algorithm and configuration used alongside benchmark summary tables and implementation details for training rewards.

  • RL method: PRISM uses Group Relative Policy Optimization as the RL algorithm for all models.The section introduces the objective and advantage computation.
  • Supporting tables: The extended tables include domain-specific evaluation breakdowns and token-budget comparisons used to contextualize the RL experiments.Table 19 covers Granite-3.3-8B domain mixtures, while Table 20 covers Granite-4 Micro token scaling.
  • RL configuration: RL training configuration includes on-policy updates and soft reference-policy updates for convergence.The training batch size equals the inference batch size.
  • RL configuration: A thinking-format reward assigns zero reward when responses do not follow the required <think> thoughts </think> response structure.Correctness is evaluated only when the specified format is followed.

E RL Training Curves for Granite-4 Micro Dense

RL training after PRISM mid-training shows stable, steadily improving correctness for Granite-4 Micro Dense at 3B parameters, with related curves extending the pattern across models and domains.

  • Granite-4 Micro Dense: ∼0.48 to ∼0.72 overall verifiable correctness rises steadily over 2k RL steps for Granite-4 Micro Dense 3B.The model is evaluated after PRISM mid-training with the Math+Code+Science mixture.
  • Granite-4 Micro Dense: ∼0.52 to ∼0.70 math correctness improves over 2k RL steps for Granite-4 Micro Dense 3B.The domain-specific curve is monotonically increasing.
  • Cross-model curves: Monotonically increasing, non-saturating correctness curves extend across Mistral-7B, LLaMA-3.1-8B, and Granite-4 Micro Dense.The additional curves support generalization of the PRISM →RL training pattern across families and scales.
  • Cross-model curves: Figure 18 reports consistent code and math improvements for Mistral-7B using the unbalanced MCS RL mix.The figure complements the broader cross-model learning-curve results.

F.2 RL on Base Models (No Mid-Training)

RL applied directly to base models produces weak or unstable reasoning gains, unlike RL after PRISM mid-training. Balanced RL mixes generally perform comparably to unbalanced mixes, while the mid-training/RL sparsity asymmetry remains robust across thresholds.

  • RL on Base Models: AIME24 and AIME25 remain near zero when RL is applied directly to LLaMA-3.1-8B and Mistral-7B base models.GPQA-Diamond stagnates or regresses below the base model’s level in these no-mid-training settings.
  • RL on Base Models: RL after PRISM mid-training shows large, sustained reasoning gains across benchmarks, contrasting with direct RL on base models.LLaMA-3.1-8B exhibits stable monotonic gains, while Granite-4 Micro Dense shows consistent but smaller absolute gains than 8B models.
  • RL Data Mix: Balanced RL sampling produces comparable overall gains to the unbalanced mix across most models, with modest domain-specific differences.For Mistral-Small 24B, the balanced-mix GPQA-Diamond gain is +25.93, slightly below the unbalanced mix.
  • Weight-Change Sparsity: At the 1% threshold, RL leaves 95.0% of Granite-3.3 parameters unchanged, while mid-training changes most parameters.The dense/sparse asymmetry holds from 0.1% to 10%, making the conclusion robust to threshold choice.

H RL Sub-component Weight Analysis

RL targets similar weight sub-components whether or not mid-training precedes it. It preferentially changes residual-stream projections while leaving core Mamba state-dynamics parameters effectively frozen.

  • Analysis Setup: The analysis measures the fraction of parameters changing by more than 1% relative change across attention, MLP, and Mamba matrices.It compares MT→RL and Base→RL conditions for Granite-3.3 and Nemotron-H.
  • Sub-component Targeting: RL’s sub-component targeting differs by at most 1.8 percentage points between mid-trained and base starting points.Most differences are below 1 percentage point, indicating similar targeting across initialization conditions.
  • Mamba Parameters: Mamba A and dt parameters are completely frozen during RL, with less than 0.1% changed.Learned in_proj and out_proj matrices change at rates comparable to attention projections.

I Extended CKA Representation Analysis

Across dense and attention-Mamba models and three input distributions, RL preserves the representational geometry established by mid-training. Divergence between base and mid-trained models varies by model and input type.

  • Analysis Setup: The CKA analysis uses 200 prompts per input type and batch-size-1 encoding, with layer index on the x-axis and linear CKA score on the y-axis.Each figure contains separate panels for Wikipedia, C4, and GSM8K.
  • Cross-model CKA: MT versus RL representation similarity remains above 0.998 at every layer across all tested models and input distributions.This holds for Granite-3.3, LLaMA-3.1, and Nemotron-H on Wikipedia, C4, and GSM8K inputs.
  • Base-to-Mid-training Divergence: Base-versus-mid-training divergence is model-specific across input distributions.Granite-3.3 dips most at the GSM8K output layer, LLaMA-3.1 on C4 web text, and Nemotron-H across broader later-middle layers.

J AIME 2026 Evaluation

PRISM mid-trained Granite-3.3 and Mistral-Small improve on the post-training AIME 2026 benchmark during RL. Gains remain visible on this benchmark despite its release after the training runs.

  • AIME 2026 Results: Granite-3.3 improves from ∼33% to ∼37% on AIME26 over RL training.AIME26 was released after completion of all training runs.
  • AIME 2026 Results: Mistral-Small improves from ∼30% to ∼38% on AIME26 over RL training.The evaluation uses two PRISM mid-trained models: Granite-3.3 at 8B and Mistral-Small at 24B.

K Model Generations

The examples show mid-training introducing structured reasoning and RL refining it, while model behavior varies across math, code, and science tasks, including looping failures and improved termination.

  • Mid-training introduces structured reasoning, while RL further refines reasoning chains into more accurate and complete solutions.
  • The examples also show failure modes: base and mid-trained models may loop, whereas RL training applies stop-token and repetition penalties and consistently terminates outputs correctly.
  • +11.95 on GPQA-Diamond is reported for Granite-4-H Micro under the balanced-mix PRISM →RL pipeline.
  • RL targets the same sub-components from base and mid-trained starting points, with value and output projections most affected while Mamba parameters and norms remain frozen.
  • >0.998 CKA between mid-training and RL representations is reported across dense and hybrid models, while base-to-mid-training divergence is strongest on math prompts.
Loading 2603.17074v3…