Source-linked AI summary
Towards a theory of inference-time alignment with unknown rewards
Steve Hanneke, Hongao Wang, Mingyue Xu
TL;DR
The paper studies inference-time alignment as statistical weak-to-strong learning without a known reward model. It introduces the alignment dimension and proves that a reward class is alignment learnable exactly when this dimension is finite.
Problem
Inference-time alignment has few statistical learning guarantees, especially when learning from preference data without assuming access to a good reward estimate.
Method
The paper formulates alignment as weak-to-strong learning and uses one-inclusion graph prediction with boosting and sample compression to compare non-dominating response-label sets.
Results
A reward class is alignment learnable if and only if its alignment dimension is finite.
Takeaways & Limitations
The alignment dimension completely characterizes learnability for the paper’s realizable binary-reward formulation.
Takeaways & Limitations
The analysis is restricted to realizable settings with binary rewards and assumes the reference policy generates a good response with constant probability.
Abstract
from arXiv · showhide
Generative model alignment has received broad interest, and significant progress has been made in supervised fine-tuning and inference-time computation. Yet, alignment has remained poorly understood from a statistical learning perspective. We formulate inference-time alignment as a weak-to-strong learning problem, where a reference policy (weak learner) is assumed to be fairly good and the goal is to produce a strong learner that predicts a good response at test time with arbitrarily high probability. Our problem is formulated as learning from scratch --- everything is learned from data rather than assuming access to a good reward estimate, and thus differs from the existing inference-time alignment theory. Our model shares similarity to the recent work of arXiv:2510.15464, where for each prompt, there could be multiple good responses. Our definition of the alignment learnability follows the PAC learning principle. We introduce a novel combinatorial dimension of the reward class which we call the alignment dimension, and show that it completely characterizes the alignment learnability --- a reward class is alignment learnable if and only if its alignment dimension is finite. The core of our learning procedure works by invoking the ordinary one-inclusion graph algorithm to run a tournament over all pairs of label sets satisfying that neither is a subset of the other. We believe our results might shed light on establishing a complete theoretical understanding towards alignment.
1 Introduction
This section frames inference-time alignment as a statistical learning problem with unknown rewards, contrasting it with reward-model-based training and prior theory that assumes reward estimates. It introduces a binary-reward setting and positions the work as a general framework for learning aligned responses directly from preference data.
- Modern generative models can acquire impressive capabilities while still exhibiting unintended or harmful behavior, motivating alignment with human intentions.
- RLHF/RLAIF first learns a parameterized reward function from preference data and then optimizes the model, whereas inference-time alignment keeps the reference policy fixed and uses rewards only during inference.Inference-time methods include rejection sampling, token-level decoding guidance, and planning or search.
- The central question is which reward classes are learnable for inference-time alignment when the learner has access only to preference data.
- The paper studies learning from scratch with binary reward functions R ⊆ {0, 1}^X×Y and assumes the true reward function r* belongs to R.Unlike existing inference-time alignment theory, the framework jointly considers reward estimation and test-time computation without assuming an externally supplied good reward estimate.
- The work presents a statistical framework for general alignment, particularly inference-time alignment, while distinguishing it from related multiple-correct-answer learning settings and finite reward-class results.Related prior work may include a policy class or restrict the reward class to finite cardinality, whereas this work omits a policy class to capture inference-time methods.
2 Preliminaries
The paper formulates inference-time alignment as learning an unknown binary reward through accepted response sets from sampled slates. It defines alignment learnability via finite training and inference-time sampling, then characterizes it exactly using the alignment dimension.
- Problem formulation: A reward class induces set-valued acceptance functions that mark which responses in a size-k slate receive reward 1.For z=(x,y_1,...,y_k), A_r(z) contains exactly the indices j with r(x,y_j)=1.
- Problem formulation: Training data consist of i.i.d. instances and k-response slates sampled from an unknown reference policy, together with the target reward’s acceptance set.The learner has oracle access to the reference policy but no distribution information about it.
- Problem formulation: The learner predicts one index from each test slate, aiming for the selected response to satisfy the unknown target reward with high probability.This models inference-time alignment as selecting a good response from candidates rather than predicting a single response directly.
- Alignment learnability: Alignment learnability requires finite training size m and finite inference-time slate size k for every γ, ϵ, and δ.For constant γ, finite k suffices to guarantee a good response; finiteness of the alignment dimension ensures finite m.
- Alignment dimension: The alignment dimension d_ALN(R) is the largest number of instances that can be alignment shattered, and R is alignment learnable if and only if d_ALN(R) < ∞.Thus, the alignment dimension completely characterizes the learnability of the reward class.
3 Proof of Theorem 1
Theorem 1 is proved by matching upper and lower bounds: finite alignment dimension suffices for learnability and is necessary. The upper bound uses a non-dominating Natarajan dimension and a transductive One-inclusion Graph procedure, while the lower bound establishes a corresponding sample-complexity obstruction.
- Upper bound: Finite non-dominating Natarajan dimension suffices for alignment learnability and is bounded by the alignment dimension.The proof establishes dN(Ak(R)) ≤ dALN(R), then uses finite dN(Ak(R)) to derive learnability.
- Upper bound: The learning procedure applies the One-inclusion Graph algorithm to pairwise label-set subsamples where neither candidate acceptance set is a subset of the other.The procedure is transductive, and its high-probability guarantee combines OIG with boosting and sample compression.
- Upper bound: Theorem 2 states that every reward class with finite alignment dimension is alignment learnable.The theorem follows from the upper-bound construction and its high-probability error guarantee.
- Lower bound: Alignment learnability requires finite alignment dimension, completing the characterization in Theorem 1.The lower bound is supplied by Theorem 3, which complements the upper bound in Theorem 2.
- Lower bound: m = Ω((d + log(1/δ))/ϵ) samples are necessary for a reward class with alignment dimension d under the stated ϵ and δ ranges.Lemma 5 specifies ϵ ∈ (0, 1/60) and δ ∈ (0, 1/18), and constructs a target reward and distribution witnessing the lower bound.
4 Discussion
The paper frames inference-time alignment as weak-to-strong statistical learning and introduces alignment dimension as a complete learnability characterization. It also identifies assumptions and mismatches with practical alignment, motivating extensions to single-response data, direct prediction, and real-valued preferences.
- Contributions: The framework casts alignment as weak-to-strong learning that improves a reference policy using preference data, initially restricting analysis to realizable settings and binary rewards.The learner uses i.i.d. data to boost a reference policy’s prediction performance.
- Contributions: Alignment dimension completely characterizes learnability: a reward class is alignment learnable if and only if its alignment dimension is finite.The procedure additionally assumes an ordering of response tuples, although the resulting notion is permutation invariant.
- Limitations: The theory requires multiple responses sampled from the reference policy for each training input and at test time so a good response can appear in the candidate slate.This requirement follows from the assumption that the reference policy outputs a good response with constant probability.
- Limitations: Practical alignment commonly collects two responses per prompt, motivating the open problem of alignment learnability with single-response training data while retaining multiple test responses.For binary rewards, one versus two sampled responses makes almost no difference to the relevant failure event.
- Limitations: The selection-based formulation differs from practical methods that predict one response for each prompt and typically regularize the learned policy toward the reference policy with KL divergence.The paper studies selection because direct prediction can otherwise be irrelevant to the reference policy, unlike RLHF-style objectives.
- Future directions: A major extension is exact preference alignment with real-valued rewards, where training compares two reference-policy responses using a binary preference label and must address how to measure reference-policy quality.Existing approaches use coverage or discrepancy measures relative to comparator policies, including coverage coefficients and EM-divergence.