🤖 machine learning

Towards a theory of inference-time alignment with unknown rewards

This paper establishes a PAC learning framework for inference-time alignment under unknown rewards by defining a novel "alignment dimension" that fully characterizes learnability and proposing a tournament-based algorithm using the one-inclusion graph to transform a weak reference policy into a strong learner.

Original authors: Steve Hanneke, Hongao Wang, Mingyue Xu

Published 2026-08-18

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Steve Hanneke, Hongao Wang, Mingyue Xu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

In the rapidly evolving world of artificial intelligence, a central challenge is ensuring that computer programs behave in ways that match human intentions. While modern systems can generate fluent text and solve complex problems, they sometimes produce outputs that are nonsensical, harmful, or simply unhelpful. To fix this, researchers have developed methods to "align" these models with human values. One common approach involves training the model on vast amounts of data where humans or other AIs have ranked different responses, teaching the system to prefer good answers over bad ones. Another approach, known as inference-time alignment, does not change the model's internal code. Instead, it acts like a filter at the moment of use: the system generates several possible answers, and a separate scoring mechanism picks the best one to show the user. This method is popular because it is flexible and does not require the expensive process of retraining the entire system. However, while these techniques work well in practice, scientists have struggled to explain exactly why they work or what limits their success from a mathematical standpoint.

A team of researchers at Purdue University has now taken a major step toward filling this gap by building a new statistical framework to understand inference-time alignment. They approached the problem as a learning task where a "weak" starting point is improved into a "strong" result using data. Imagine a reference model that is generally competent but occasionally makes mistakes; it serves as a baseline that generates a list of candidate answers. The goal is to use a dataset of human preferences to teach a new system how to reliably pick the single best answer from that list every time. The researchers asked a fundamental question: under what conditions is it actually possible to learn this selection skill from data alone, without any prior knowledge of how the scoring system works? They found that the answer depends entirely on the complexity of the rules used to judge the answers.

The team discovered that not all sets of judging rules can be learned. They introduced a new way to measure the complexity of these rules, which they call the "alignment dimension." Think of this dimension as a measure of how many different ways the rules can contradict or confuse the system. If this number is finite, meaning the rules have a manageable level of complexity, then it is possible to design an algorithm that will eventually learn to pick the correct answer with near-perfect accuracy, given enough data. If the dimension is infinite, the rules are too chaotic to be learned from data no matter how much is collected. This finding provides a complete mathematical guarantee: a reward system is learnable if and only if its alignment dimension is finite. This is a significant shift from previous theories, which often assumed that researchers already had a perfect understanding of the scoring system or that the rules were simple enough to be described by a fixed number of parameters.

To prove this, the researchers designed a specific learning procedure that operates like a tournament. When the system needs to choose an answer, it does not simply look at the data once and make a guess. Instead, it compares pairs of possible answer groups against each other. For every pair of groups where neither is clearly a subset of the other, the system runs a specialized comparison algorithm to decide which group is more likely to contain the correct answer. By running these comparisons across all possible pairs, the system narrows down the field until it is left with a small, highly reliable set of candidates from which to select the final answer. This method works by leveraging the fact that the starting model, while imperfect, has a constant chance of generating a good answer. By sampling enough candidates and using the tournament logic to filter them, the system can boost its success rate to an arbitrarily high level.

The paper also clarifies what this new theory rules out. It shows that simply trying to memorize the best answers from a training set, a method known as empirical risk minimization, is not sufficient on its own. In some cases, a system might need to rely on the specific structure of the data and the ability to sample new candidates at the moment of testing, rather than just recalling what it saw during training. The researchers demonstrated that for certain types of complex reward systems, no amount of training data would allow a standard learning algorithm to succeed without this additional sampling step. Their work suggests that the key to successful alignment lies in the interplay between the complexity of the reward rules and the ability to generate multiple options at test time.

This research represents a foundational step toward a rigorous theory of AI alignment. By defining the exact conditions under which alignment is possible, the authors provide a clear target for future development. They have moved the field beyond trial and error, offering a mathematical proof that tells engineers exactly when their alignment strategies will work and when they will fail. While the current study focuses on binary rewards—where an answer is either good or bad—the framework opens the door to understanding more complex, real-valued scoring systems. The ultimate goal is to establish a set of principles that can guide the creation of safer and more reliable AI systems, ensuring that as these models become more powerful, their ability to follow human intent remains mathematically guaranteed.

Technical Summary: Towards a theory of inference-time alignment with unknown rewards

Problem Formulation
This work addresses the statistical learning foundations of inference-time alignment, a paradigm where a generative model's responses are selected or steered at test time based on a reward signal, without updating the model parameters (unlike Reinforcement Learning from Human/AI Feedback, RLHF). The authors formulate this as a weak-to-strong learning problem under the assumption of unknown rewards.

The setting involves:

Reference Policy ( $\pi_{ref}$ ): A "weak learner" (e.g., a pre-trained or SFT model) that generates candidate responses. It is assumed to have a constant probability $\gamma$ of generating a "good" response (reward $r^*(x, y) = 1$ ) for any input $x$ .
Reward Class ( $\mathcal{R}$ ): A set of binary reward functions $r: \mathcal{X} \times \mathcal{Y} \to \{0, 1\}$ . The true reward $r^*$ is unknown but lies within $\mathcal{R}$ .
Data Generation: Training data consists of $m$ i.i.d. samples. Each sample includes an input $x_i$ and a "slate" of $k$ responses $\{y_i^{(1)}, \dots, y_i^{(k)}\}$ drawn from $\pi_{ref}(\cdot|x_i)$ . The learner observes the set of indices of good responses $A_{r^*}(z_i) = \{j : r^*(x_i, y_i^{(j)}) = 1\}$ .
Goal: Design a learning algorithm that, given a test input $x$ and a slate of $k$ candidates from $\pi_{ref}$ , selects an index $j$ such that $r^*(x, y^{(j)}) = 1$ with high probability ( $1-\epsilon$ ).

The authors distinguish their work from existing theoretical literature by removing the assumption of access to a pre-estimated reward model. Instead, they treat the problem as learning from scratch, where the reward function itself must be inferred from finite preference data.

Methodology and Core Concepts
The paper introduces a novel combinatorial dimension to characterize the learnability of the reward class:

Alignment Dimension ( $d_{ALN}(\mathcal{R})$ ):
A set of inputs $\{x_1, \dots, x_d\}$ is alignment shattered by $\mathcal{R}$ if there exist pairs of responses $(y_i^{(0)}, y_i^{(1)})$ such that for any binary vector $b \in \{0, 1\}^d$ , there exists a reward function $r_b \in \mathcal{R}$ satisfying $r_b(x_i, y_i^{(b_i)}) = 1$ and $r_b(x_i, y_i^{(1-b_i)}) = 0$ . The alignment dimension is the size of the largest such shattered set.
Non-Dominating Natarajan Dimension:
To prove learnability, the authors relate the problem to the One-Inclusion Graph (OIG) algorithm used for partial concept classes. They define a "non-dominating Natarajan dimension" for the induced set-valued acceptance class. They prove that this dimension is upper-bounded by the alignment dimension.
Learning Algorithm (Tournament over Non-Dominating Pairs):
The proposed learning procedure is transductive and operates at test time:
- For a test slate $Z$ , the algorithm considers all pairs of non-empty subsets of indices $(A, B)$ such that neither is a subset of the other ( $A \not\subseteq B$ and $B \not\subseteq A$ ).
- For each pair, it constructs a training subset containing instances where the true acceptance set was either $A$ or $B$ .
- It runs the OIG algorithm on this subset to predict whether the true acceptance set is $A$ or $B$ .
- This process forms a "tournament" that narrows down the possible acceptance sets to a chain. The final prediction is any response index contained in the intersection of the surviving sets.

Key Results
The paper establishes a complete characterization of alignment learnability:

Theorem 1: A reward class $\mathcal{R}$ $R$ is alignment learnable if and only if its alignment dimension $d_{ALN}(\mathcal{R})$ $d_{A L N} (R)$ is finite.
- Upper Bound (Theorem 2): If $d_{ALN}(\mathcal{R}) < \infty$ , the proposed tournament algorithm achieves $(\epsilon, \delta)$ -learnability with a sample size $m$ and inference-time slate size $k$ that are finite functions of the dimension, $\epsilon$ , $\delta$ , and $\gamma$ .
- Lower Bound (Theorem 3): If $d_{ALN}(\mathcal{R}) = \infty$ , no learning algorithm can achieve alignment learnability with finite samples. The proof constructs a hard instance using the probabilistic method, showing that infinite dimensionality implies the existence of distributions where the learner cannot distinguish the true reward function.

Significance and Claims
The authors position this work as a pioneering step toward a rigorous statistical learning theory for general alignment problems. Key claims regarding the paper's contribution include:

Novel Framework: It is the first work to formulate inference-time alignment as a statistical learning problem where the reward is unknown and learned entirely from data, distinguishing it from prior works that assume access to a good reward estimate.
Theoretical Characterization: The introduction of the alignment dimension provides a necessary and sufficient condition for learnability, analogous to the VC dimension in binary classification but tailored to the specific constraints of selecting one good response from multiple candidates.
Distinction from Existing Paradigms: The authors clarify that their setting differs from multiclass learning (where one label is correct) and list learning (where a list of labels is output). In their setting, multiple responses may be correct, but the learner must successfully identify one of them.
Modesty on Scope: The authors acknowledge limitations, noting that their current formulation assumes binary rewards and realizable settings. They explicitly state that their learning procedure is non-standard (relying on test-time computation and tuple ordering assumptions) and that extending the theory to single-response training data, real-valued rewards, or preference-based learning (Bradley-Terry models) remains an open direction for future work.

The paper concludes that establishing the alignment dimension as the fundamental complexity measure offers a theoretical foundation for understanding when and how inference-time alignment is possible, potentially guiding the development of more systematic principles for practical alignment.

Technical Summary: Towards a theory of inference-time alignment with unknown rewards

More like this