Demystifying LLM-as-a-Judge: Analytically Tractable Model for Inference-Time Scaling

TL;DR

This paper introduces an analytically tractable model for understanding inference-time scaling in large language models, revealing how sampling strategies and reward misspecification affect generalization error and optimal inference-time computation.

Contribution

The paper provides a theoretical framework for inference-time scaling in LLMs using Bayesian linear regression, analyzing the effects of sampling and reward misspecification on generalization.

Findings

Generalization error decreases as 1/k^2 with optimal reward and sampling temperature.

Reward misspecification can lead to a finite optimal number of inference samples.

Inference-time compute advantage diminishes with increasing task difficulty.

Abstract

Recent developments in large language models have shown advantages in reallocating a notable share of computational resource from training time to inference time. However, the principles behind inference time scaling are not well understood. In this paper, we introduce an analytically tractable model of inference-time scaling: Bayesian linear regression with a reward-weighted sampler, where the reward is determined from a linear model, modeling LLM-as-a-judge scenario. We study this problem in the high-dimensional regime, where the deterministic equivalents dictate a closed-form expression for the posterior predictive mean and variance. We analyze the generalization error when training data are sampled from a teacher model. We draw $k$ inference-time samples and select via softmax at a temperature applied to a quadratic reward. When the reward is not too different from the teacher, the generalization error decreases monotonically with increasing inference time samples $k$. However, the specific reward that optimizes inference-time selection generally differs from the teacher. In contrast, substantial reward misspecification induces a finite optimal $k$ beyond which more sampling can increase the generalization error. For fixed $k$, there exists an optimal sampling temperature. We experimentally verify these facts in large language model inference with an additional large language model as a judge. In the "best-of-$k$" limit with the teacher as reward, we theoretically show that the generalization error decays as $\Theta(1/k^2)$ and determine the leading coefficient via extreme value theory. These formulas delineate domains where scaling inference-time computation is provably preferable to collecting more data. Finally, we demonstrate that when task difficulty increases, the previously mentioned advantage of inference-time compute degrades.