Predicting LLM Reasoning Performance with Small Proxy Model
This provides a practical, cost-effective method for optimizing pre-training datasets for reasoning capabilities, which is incremental but addresses a specific bottleneck in model development.
The paper tackles the problem of predicting large language model reasoning performance using small proxy models, showing that rBridge reduces dataset ranking costs by over 100x and achieves the strongest correlation across reasoning benchmarks at scales up to 32B parameters.
Given the prohibitive cost of pre-training large language models, it is essential to leverage smaller proxy models to optimize datasets before scaling up. However, this approach becomes challenging for reasoning capabilities, which exhibit emergent behavior that only appear reliably at larger model sizes, often exceeding 7B parameters. To address this, we introduce rBridge, showing that small proxies ($\leq$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 100x relative to the best baseline, (ii) achieves the strongest correlation across six reasoning benchmarks at 1B to 32B scale, and (iii) zero-shot transfers predictive relationships across pre-training datasets at 1B to 7B scale. These findings indicate that rBridge offers a practical path for exploring reasoning-oriented pre-training at lower cost.