To address the high cost of large-scale language model pre-training, we propose a method to optimize datasets by leveraging small proxy models. Specifically, to address the inference performance challenges inherent only in large-scale models, we introduce rBridge, demonstrating that small proxy models (<1 billion) can effectively predict the inference performance of large-scale models by more closely aligning with (1) the pre-training objective and (2) the target task. rBridge weights the negative log-likelihood with task alignment and uses the inference traces of state-of-the-art models as gold labels. Experimental results demonstrate that rBridge reduces dataset ranking costs by over 100x compared to conventional methods, achieves the strongest correlation between models ranging from 1 billion to 32 billion across six inference benchmarks, and achieves zero-shot transfer of prediction relationships between pre-trained datasets ranging from 1 billion to 7 billion.
