Learning to Answer from Correct Demonstrations (2510.15464v1)

Abstract: We study the problem of learning to generate an answer (or completion) to a question (or prompt), where there could be multiple correct answers, any one of which is acceptable at test time. Learning is based on demonstrations of some correct answer to each training question, as in Supervised Fine Tuning (SFT). We formalize the problem as offline imitation learning in contextual bandits, with demonstrations from some optimal policy, without explicitly observed rewards. Prior work assumes that the demonstrator belongs to a low-complexity policy class, which motivates maximum likelihood estimation (i.e., log-loss minimization). In contrast, we propose relying only on the reward model (specifying which answers are correct) being in a low-cardinality class, which we argue is a weaker assumption. We show that likelihood maximization methods can fail in this case, and instead devise an alternative novel approach that learns with sample complexity logarithmic in the cardinality of the reward class. Our work motivates looking beyond likelihood maximization when learning from correct demonstrations.

• The paper introduces a sample-optimal algorithm that minimizes the error in learning correct answers from demonstrations.
• It demonstrates that maximum likelihood estimation fails under low-cardinality reward models, leading to arbitrarily high generalization errors.
• The proposed method achieves logarithmic sample complexity and shifts focus from distribution matching to direct reward maximization in imitation learning.

Learning to Answer from Correct Demonstrations: A Technical Analysis

Problem Formulation and Motivation

The paper formalizes the problem of learning to generate correct answers in settings where multiple responses may be equally valid, such as mathematical problem solving, code generation, or recommendation systems. The central abstraction is a contextual bandit: contexts $x$ (questions/prompts) and actions $y$ (answers/completions), with a binary reward function $r_*(x, y) \in \{0, 1\}$ indicating correctness. The learner receives demonstrations from an optimal policy $*$, i.e., for each $x$, a sampled $y$ such that $r_*(x, y) = 1$. The objective is to learn a policy $\widehat{\pi}(y|x)$ that, for unseen $x$, outputs a correct $y$ with high probability, minimizing the error $$L_{,*}(\pi) = \mathbb{E}_{x, y \sim \pi(\cdot|x)}[1\{y \notin *(x)\}]$$.

This setting is directly relevant to supervised fine-tuning (SFT) of LLMs, where models are trained on curated question–answer pairs, and the goal is to produce any correct answer, not to match the distribution of the demonstrator.

Policy Class vs. Reward Model Class Assumptions

Prior work typically assumes the demonstrator belongs to a low-complexity policy class $\Pi$, motivating maximum likelihood estimation (MLE) or log-loss minimization. Under this assumption, MLE achieves sharp $O(\log|\Pi|/m)$ convergence rates for the error, and the learned policy closely matches the demonstrator's distribution.

The paper departs from this by proposing a weaker and more realistic assumption: the reward model class (the set of functions specifying which answers are correct) has low cardinality. This reflects practical dataset curation, where correctness is well-defined but the set of valid answers is large and diverse. The key technical insight is that the policy class induced by a low-cardinality reward model class is much larger (often infinite), making MLE guarantees vacuous.

Failure Modes of Likelihood Maximization

The authors rigorously demonstrate that MLE can fail under the low-cardinality reward model class assumption, even in simple cases. For example, with $|\mathcal{S}| = 2$ (two possible support functions for correct answers), MLE may generalize poorly, outputting incorrect answers on unseen contexts with probability arbitrarily close to 1, regardless of sample size. Restricting the policy class to match the reward model class cardinality (e.g., uniform policies over correct answers) does not resolve the issue, as the demonstrator need not be uniform, and the error can remain high if the set of correct answers is large.

Sample-Optimal Learning Algorithms

The paper introduces a novel learning algorithm that achieves sample complexity logarithmic in the reward model class size, $O(\log|\mathcal{S}|)$. The approach is based on a weighted majority vote among consistent hypotheses, with an online-to-batch conversion for the statistical setting. The online algorithm maintains weights over hypotheses, updating them based on observed demonstrations. Crucially, the update rule increases the weight of hypotheses under which the predicted action is incorrect, exponentially reducing the number of mistakes.

The main result is that, for any finite reward model class $\mathcal{S}$, the algorithm makes at most $\log_2|\mathcal{S}|$ mistakes in the online setting, and achieves expected error at most $\log_2|\mathcal{S}|/m$ in the batch setting. High-probability bounds are established via Freedman's inequality, and the guarantees hold even if demonstrations are adaptively chosen.

Extensions: Bounded Rewards, $pass@k$, and Suboptimal Demonstrators

The framework is extended to general bounded reward functions $r(x, y) \in [0, 1]$, with analogous sample complexity guarantees. For the $pass@k$ objective—where the policy outputs $k$ answers and is correct if any is valid—the sample complexity improves to $O(\log_{k+1}|\mathcal{S}|)$, which is minimax optimal. The algorithm is further generalized to handle suboptimal demonstrators (not always producing correct answers), achieving error within a constant factor of the demonstrator's loss.

Reward Maximization vs. Distribution Matching

A central conceptual contribution is the argument that reward maximization (utility) should be the primary objective when learning from demonstrations, rather than matching the demonstrator's distribution. The paper provides formal examples where reward maximization is trivial but distribution matching is impossible. Likelihood maximization is inherently tied to cloning the demonstrator, which may be infeasible or unnecessary in practice. The authors advocate for direct reward-based learning, especially in LLM SFT, where producing any correct answer suffices.

Theoretical and Practical Implications

The results have significant implications for imitation learning, offline contextual bandits, and SFT of LLMs. The distinction between policy class and reward model class assumptions clarifies when likelihood-based methods are appropriate and when alternative approaches are required. The sample-optimal algorithms presented are information-theoretically optimal under the reward model class assumption, but their computational efficiency remains an open question, especially for continuous or parametric classes.

The analysis also highlights the limitations of deterministic and proper learning rules, which incur superlinear sample complexity, and raises questions about the necessity of stochasticity and improperness for optimal learning.

Open Problems and Future Directions

Several open issues are identified:

• Continuous Model Classes: Extending the analysis to infinite or parametric classes, characterizing sample complexity in terms of VC/Pollard subgraph dimension or other complexity measures.
• Computational Tractability: Developing efficient implementations for large or continuous model classes, avoiding exponential runtime in the number of parameters.
• Randomization and Properness: Determining whether optimal sample complexity can be achieved with deterministic or proper policies.
• Generalization to MDPs: Investigating the implications for general Markov Decision Processes beyond contextual bandits.
• Value Suboptimality for Suboptimal Demonstrators: Achieving reward suboptimality guarantees when the demonstrator is not optimal, especially for non-binary rewards.

Conclusion

The paper provides a rigorous analysis of learning from correct demonstrations under a low-cardinality reward model class assumption, demonstrating the failure of likelihood-based methods and introducing sample-optimal algorithms. The work motivates a shift from distribution matching to direct reward maximization in imitation learning and SFT, with broad implications for the design and evaluation of AI systems. Future research should address computational efficiency, continuous model classes, and generalization to more complex decision processes.