Predicting LLM Reasoning Performance with Small Proxy Model
Woosung Koh, Juyoung Suk, Sungjun Han, Se-Young Yun, Jay Shin
Abstract
Given the prohibitive cost of pre-training large language models, it is essential to leverage smaller proxy models to optimize recipes before scaling up. However, this approach becomes challenging for reasoning capabilities, which exhibit emergent behavior that only appears reliably at larger model sizes, often exceeding 7B parameters. To address this, we introduce rBridge, showing that small proxies (1B) can effectively predict large-model reasoning by aligning more closely with (1) the pre-training objective and (2) the target task. rBridge achieves this by weighting negative log-likelihood with task alignment, using reasoning traces from frontier models as gold labels. In our experiments, rBridge (i) reduces dataset ranking costs by over 100 relative to the best baseline, (ii) achieves the strongest correlation across six reasoning benchmarks at 1B to 32B scale, and (iii) transfers predictive relationships across pre-training recipes at 1B to 7B scale. These findings indicate that rBridge offers a practical path for exploring reasoning-oriented pre-training at lower cost.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Builds on11
- Language Models are Few-Shot LearnersTom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah et al.NeurIPS 2020 · 64,255 citations
- Chain-of-Thought Prompting Elicits Reasoning in Large Language ModelsJason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma et al.NeurIPS 2022 · 22,562 citations
- Measuring Massive Multitask Language UnderstandingDan Hendrycks, Collin Burns, Steven Basart, Andy Zou et al.ICLR 2021 · 7,905 citations
- Are Emergent Abilities of Large Language Models a Mirage?Rylan Schaeffer, Brando Miranda, Sanmi KoyejoNeurIPS 2023 · 796 citations
- DoReMi: Optimizing Data Mixtures Speeds Up Language Model PretrainingSang Michael Xie, Hieu Pham, Xuanyi Dong, Nan Du et al.NeurIPS 2023 · 457 citations
Related papers
- Beyond Two-Stage Training: Cooperative SFT and RL for LLM ReasoningLiang Chen, Xueting Han, Li Shen, Jing Bai et al.ICML 2026 · 24 citations
- MobileLLM-R1: Exploring the Limits of Sub-Billion Language Model Reasoners with Open Training RecipesChangsheng Zhao, Ernie Chang, Zechun Liu, Chia-Jung Chang et al.ICLR 2026 · 10 citations
- Predicting Emergent Tool Use in LLMs Before It Emerges: A Proxy PerspectiveBowen Zhang, Yan Yan, Guang Liu, Xu-Cheng YinAAAI 2026
- Mind the Gap: Bridging Thought Leap for Improved Chain-of-Thought TuningHaolei Xu, Yuchen Yan, Yongliang Shen, Wenqi Zhang et al.NeurIPS 2025 · 2 citations
- Finding the Minimal Parameter Budget for Implicit Reasoning: A Data Complexity Driven Scaling Law for Language ModelsXinyi Wang, Shawn Tan, Shenbo Xu, Mingyu Jin et al.ICML 2026 · 4 citations
