Source-linked AI summary
MIRA: Mid-training Rubric Anchoring for Source-Aware Data Selection
Haowen Wang, Yaxin Du, Jian Yang, Jiajun Wu, Shukai Liu, Yuxuan Zhang, Pingjie Wang, Siheng Chen, Tuney Zheng, Ming Zhou, Xianglong Liu, Bryan Dai
TL;DR
Heterogeneous mid-training data require scalable selection with source-adaptive semantic criteria, but existing approaches typically provide only implicit signals or assume fixed rubrics. MIRA discovers group-specific rubrics, distills them into student scorers, and on 21 sources in 5 groups matches the unfiltered 50B-token corpus at half the token budget.
Problem
Heterogeneous mid-training sources make it difficult to select data with both scalable scoring and source-specific semantic quality criteria.
Method
MIRA discovers source-group-specific rubrics, distills anchored teacher judgments into group-specific student scorers, and applies reliability-aware source-preserving filtering.
Results
On 21 sources organized into 5 groups, MIRA-Group achieves the best 25B-token macro average across nine code benchmarks and matches the unfiltered 50B-token corpus.
Takeaways & Limitations
The results support source-adaptive quality criteria as a useful complement to scalable scoring for mid-training data selection.
Takeaways & Limitations
MIRA addresses filtering but does not cover source discovery, mixture-ratio design, curriculum scheduling, deduplication, or contamination control.
Abstract
from arXiv · showhide
Mid-training has become an important stage in modern LLM development, using large-scale curated mixtures to strengthen capabilities before final post-training. Its data selection problem is distinct: the data are optimized under a pretraining-style objective at near-pretraining scale, but are curated toward downstream capabilities and drawn from heterogeneous sources with different formats and training roles. As a result, effective selection requires both scalability and source-adaptive semantic criteria. Existing model-based methods scale well, but provide only implicit quality signals. Semantic selection methods offer stronger judgments, but usually assume fixed rubrics or standardized data formats. To address this mismatch, we propose MIRA, a source-aware filtering framework based on self-anchored rubric discovery. The key idea is to make rubric construction part of data selection: MIRA first discovers what should be evaluated for each source group, then distills those judgments into scalable student scorers for full-corpus filtering. On code-oriented mid-training with 21 sources and 5 source groups, MIRA outperforms selection baselines across nine code benchmarks and matches the full-corpus run while using only half the tokens.
1 Introduction
Mid-training data selection must combine scalable scoring with semantic criteria that adapt to heterogeneous source formats and capability roles. MIRA discovers group-specific rubrics, distills them into scalable scorers, and achieves strong code-benchmark performance with half the full-corpus token budget.
- Mid-training combines a pretraining-style objective and near-pretraining scale with curated, heterogeneous sources targeting downstream capabilities.
- Existing scalable filters provide implicit signals, while semantic filters generally assume fixed criteria, leaving a mismatch for heterogeneous mid-training data.
- MIRA separates rubric discovery from scalable scoring by distilling frontier-teacher judgments into group-specific student scorers for full-corpus filtering.
- Across 9 code benchmarks and 21 sources organized into 5 groups, MIRA outperforms strong 50B-token baselines and slightly exceeds the unfiltered run.
- MIRA derives source-group-specific quality criteria instead of applying a single global score or fixed rubric.
- MIRA’s source-conditioned aggregation and source-aware retention thresholds support calibrated filtering while preserving source diversity.
2 Related Work
Prior data-selection methods trade off scalability, interpretability, and semantic specificity. MIRA differs by inducing source-adaptive rubrics and applying them through distilled scorers across heterogeneous sources.
- Heuristic and corpus-level filters scale to large web corpora but provide limited semantic evidence for structured examples.
- Optimization- and target-aware methods are more interpretable but typically rely on fixed or globally shared criteria.
- MIRA induces group-specific anchor rubrics from sampled records, distills them into student scorers, and filters the full corpus source-aware.
3 Method
MIRA makes rubric construction part of filtering: it discovers group-specific dimensions, anchors teacher labels, distills group-specific scorers, and aggregates scores using reliability-aware source-preserving selection.
- Source grouping: MIRA groups related sources by content-embedding similarity so each group contains records with similar capability roles and quality patterns.
- Self-anchored rubric discovery: A frontier judge freely proposes quality dimensions for sampled records, avoiding a pre-specified rubric during rubric discovery.
- Self-anchored rubric discovery: Judgment points are clustered within each group, with centroid-nearest points selected as anchor dimensions representing the teacher’s quality vocabulary.
- Self-anchored rubric discovery: The discovered rubric is an empirical artifact of the teacher’s assessment behavior rather than a normative rubric authored in advance.
- Anchored scoring: Fixed anchors let the teacher produce comparable numerical scores and rationales across a larger sample for structured supervision.
- Student distillation: Those anchored labels are distilled into group-specific generative students that output scores and rationales for every group dimension.
- Reliability aggregation: Reliability masking addresses source–dimension disagreement by excluding unreliable dimensions before trimmed-mean aggregation.
- Source-preserving selection: Source-aware retention thresholds avoid a single global cutoff, mitigating source attrition when score distributions differ across sources.
4 Experiments
MIRA is evaluated after shared mid-training and SFT under matched token budgets, with comparisons spanning four benchmark groups. The 25B-token MIRA-Group selection achieves the strongest reported overall trade-off while retaining half the full corpus.
- Setup: The evaluation compares full 50B-token mid-training with 25B-token selection methods after a shared SFT stage across four benchmark groups.The groups cover code generation, multilingual code generation, executable SQL, and software-engineering repair.
- Main results: 54.53 is MIRA-Group’s best code-generation average, with leading MBPP, MBPP+, and BigCodeBench-hard scores of 88.90, 73.80, and 33.11.Different MIRA variants specialize in different capability groups.
- Main results: 72.84 is MIRA-Source’s best Multipl-E average, while its SQL average of 94.38 is close to DSIR’s best score of 95.20.MIRA-Group reaches 36.33 on SWE-Multi, exceeding the other 25B-token selection baselines.
5 Analysis
The analyses show that MIRA’s source-aware scorers are length-robust, require group- and dimension-specific reliability handling, and discover format-dependent rubrics. Its Agent scores emphasize trajectory-level structural correctness rather than fluency alone.
- 5.1 Scorer Analysis: MIRA remains usable around the training cutoff, maintaining a smooth score profile across short and medium-length records instead of collapsing with length.PPL and DSIR show stronger length dependence, while DataMan lacks scores for very long structured traces.
- 5.1 Scorer Analysis: Agent, QA, and Text students exhibit different teacher–student residual profiles, so a single global correction rule would over-mask reliable regions or leave unreliable ones active.The reliability diagnostics use validation-split MAE across 15 group-specific dimensions.
- 5.1 Scorer Analysis: Unreliability is sparse and dimension-specific, with post-hoc masking removing selected source–dimension cells while preserving the remaining scoring signal.The diagnostics identify high-residual axes such as A2, A8, A12, and A13 depending on the group.
- 5.3 Rubric Space Visualization: QA, Text, and Agent rubrics occupy distinct regions, while all 14 DataMan dimensions lie inside the broader MIRA rubric space.Nearest-neighbor distances from DataMan rubrics to MIRA rubrics range from 0.160 to 0.449; 13/14 fall within MIRA’s P95 distance of 0.426.
- 5.4 Rubric Space: MIRA assigns low Agent scores to invalid tool-call payloads and absent error recovery, distinguishing structurally incorrect traces from superficially fluent ones.The inspected low-scoring traces concatenate JSON objects, trigger parser errors, and repeat invalid calls without recovery.
6 Conclusion
MIRA combines group-specific rubric discovery, student distillation, reliability-aware aggregation, and source-aware retention for heterogeneous mid-training data. On code-oriented mid-training, MIRA-Group matches the unfiltered corpus at half the token budget while outperforming the selection baselines.
- Conclusion: MIRA discovers group-specific anchor rubrics, distills them into efficient student scorers, and applies reliability-aware aggregation with source-aware retention thresholds.The framework is designed for full-corpus selection across heterogeneous sources.
- Conclusion: MIRA-Group achieves the best 25B-token macro average across nine code-related benchmarks on 21 sources organized into 5 groups.It outperforms PPL, DSIR, DataMan, and Random selection.
- Conclusion: MIRA-Group matches the unfiltered 50B-token corpus while using half the token budget.The conclusion attributes the result to source-adaptive quality criteria rather than scalable scoring alone.
Limitations
MIRA addresses source-aware filtering but does not cover broader mid-training data-management decisions. These include source discovery, mixture-ratio design, curriculum scheduling, deduplication, and contamination control.
- Limitations: MIRA’s scope is limited to source-aware filtering and semantic quality scoring, leaving broader mid-training data management outside the study.The stated out-of-scope decisions include source discovery, mixture-ratio design, curriculum scheduling, deduplication, and contamination control.
A Experimental Setup
The experiments compare MIRA with several data-selection baselines and the full raw corpus under matched token budgets. Filtered corpora are controlled for total token count, isolating selection effects.
- The study compares Base Model, Raw Mixture, Random, and other selection baselines.Base Model omits mid-training, Raw Mixture uses the full unfiltered corpus, and Random samples within each source.
- All methods operate on the same raw corpus.
- Filtered corpora are matched in total token count for a fair comparison.
A.2 Data Sources and Grouping
The data comprise 21 sources organized into five capability-coherent groups. Sources are filtered and sampled using group-specific quality bins and proportions.
- 21 sources are organized into 5 capability-coherent groups, each containing 3 to 5 sources.The groups cover two QA categories, one Text category, and two Agent categories.
- Group boundaries follow source-level content-embedding similarity and reflect a single capability theme per group.
- Sources outside the code-focused evaluation, with high overlap, or with broadly miscalibrated student scorers are dropped.
- Each source is independently down-sampled at p ∈ {1.0, 0.5, 0.2, 0.1} from quality-stratified stage-2 data.For p<1.0, the low/mid/high distribution becomes 50/30/20 rather than the natural approximately 30/40/30 split.
A.3 Mid-training Configuration
Mid-training conditions share the same launch and evaluation setup while differing in their data paths and token budgets. MIRA uses teacher-derived, group-specific scoring to support corpus filtering.
- Mid-training configuration: All mid-training conditions use the same launch script, with only the training-data path differing.This isolates post-SFT performance differences to the mid-training data.
- Mid-training configuration: The 50B-token full-corpus condition uses approximately 1,490 iterations, while 25B-token subset conditions use approximately 745.TRAIN_ITER is computed from total tokens, global batch size, and sequence length.
- Post-training evaluation: All mid-training checkpoints receive the same SFT recipe before evaluation.The recipe is also applied to the SFT-only baseline initialized from Qwen2.5-Coder-14B.
- Post-training evaluation: SFT packs multiple short samples into 128K-token sequences and trains for two epochs over the packed dataset.The packed sample count determines TRAIN_ITER and the cosine learning-rate decay.
- Post-training evaluation: Loss is computed only on assistant response tokens, while system, user, and tool-result tokens are masked.
- Teacher and student scoring: MIRA’s scoring stack combines one frontier teacher with five group-specific student scorers.Students are distilled from anchored teacher labels for full-corpus scoring.
A.6 Student Scorer Training Configuration
The five student scorers are trained with a shared full-parameter recipe, while sequence-length budgets reflect group-specific data distributions. Distillation adds relatively little compute compared with mid-training.
- Student architecture and training: Each student is fully fine-tuned from a shared 35B-parameter MoE base with approximately 3B active parameters per token.The same base weights are used across groups, while training data paths and sequence-length budgets differ.
- Shared training recipe: All five students share the same optimization and runtime recipe, including sequence packing, gradient checkpointing, ZeRO-3, BF16, and flash-attention v2.The uniform parallelism setting avoids confounding cross-group comparisons.
- Validation: Student validation uses held-out records and teacher/student mean absolute error averaged across 15 anchor dimensions.
- Group-specific settings: Group-specific cutoff lengths are approximately 24,576 for qa_math_reasoning and 16,384 for the other four groups.Per-device micro-batch is 1 with gradient accumulation 8 throughout.
- Compute: Approximately 250 H800 GPU-hours are required across all five students, placing distillation below 4% of per-condition mid-training cost.The paper states that adding rubric discovery and student distillation does not materially change the overall compute budget.
A.7 Reliability Diagnostics
MIRA’s reliability diagnostics identify source/slot disagreements between teacher and student scorers, then mask unreliable cells before aggregation. The diagnostics show uneven reliability across sources and dimensions, including a worst-case source whose score depends on only seven surviving slots.
- Masking rule: MAE ≥1 masks source/slot cells, while cells below 1 contribute to trimmed-mean aggregation.The threshold is defined on the 0–10 scoring scale and corresponds to roughly one 5-point training-utility band.
- Clean sources: Nine of 21 sources have no flagged slot and contribute all 15 dimensions to trimmed-mean aggregation.No text_code_doc source is clean, whereas clean sources occur in both QA groups and both Agent groups.
- Aggregate mask coverage: 37 of 315 source/slot cells are masked, or approximately 12% of the 21 × 15 A1..A15 budget.Masking is concentrated in particular dimensions, including Technical Precision, format-compliance axes, and Code Reference Accuracy.
- Worst-case source: funccall_agent_b has the highest disagreement, with 8 of 15 slots flagged and an overall MAE of 1.00.Its retained score uses the seven surviving slots, making it effectively narrower than scores for other agent sources.
- Interpretation: Slot indices are anchor-list-specific, so the same index can represent different dimensions across sources.For example, A8 denotes Technical Precision in qa_math_reasoning but Action Efficiency in agent_toolcall.