Predicting LLM Reasoning Performance with Small Proxy Model

TL;DR

This work tackles the cost and practicality of predicting large-scale reasoning from small proxy models. It introduces rBridge, which uses frontier-model reasoning traces $R^\phi$ as the gold $Y^*$ and applies token-level weighting of $NLL$ by frontier confidence to achieve both distributional and task alignment with pre-training objectives. Empirically, rBridge delivers over 100× compute savings in dataset ranking, exhibits the strongest cross-scale correlations across multiple reasoning benchmarks from $1\text{B}$ to $32\text{B}$, and enables zero-shot functional transfer of learned relationships across datasets with substantial cost reductions. The results suggest a practical path for low-cost, reasoning-oriented pre-training data optimization and dataset ranking, with wide potential for two-stage data curation and improved efficiency in scaling AI systems.

Abstract

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.

Figures (10)

• Figure 1: Schematic overview of rBridge, which is used to predict and rank performance at much larger model size. We use a frontier model $\pi^\phi$'s reasoning trace NEURIPS2022_9d560961 as the gold label $Y^*$ and compute weighted NLL for evaluation. Each token $i$'s NLL is weighted by the frontier model's confidence in that token (MinMax normalized). To handle tokenizer mismatches between proxy and frontier models, we compute weights at the letter level and average within tokens.
• Figure 2: Using MATH500 as an example benchmark, given the same data source OLMo-Mix-1124 olmo20242, smaller models exhibit more noise and get the direction wrong, making it challenging to use smaller models to proxy larger model performance. R$^2$ values are derived from linear curve fitting. Extended visualization across other reasoning benchmarks are available in Appendix \ref{['app:noisy']}.
• Figure 3: When evaluating a 1B pre-trained model with next token prediction, $\pi^\text{p}(y_\tau \vert x, y^*_{<\tau})$, how in-distribution the target $Y^*$ is becomes important. All visualized benchmarks demonstrate smooth improvements at larger (13$\times$, 32$\times$) scale with target metric Acc./p@k. For clarity, we visualize the benchmarks in our empirical study with the two smallest and largest average NLL values.
• Figure 4: Using reasoning trace $R^\phi$ over benchmark test set's $Y^*$ significantly reduces NLL, suggesting that $R^\phi$ is more in-distribution. Error bars indicate one standard deviation.
• Figure 5: Example visualization from a fold of proxy-target relationship study at 1B $\rightarrow$ 13B on MMLU Pro (STEM). Each data point represents equal trained tokens for the proxy and target model.