English

Learning to Answer from Correct Demonstrations

Machine Learning 2026-02-27 v2 Artificial Intelligence Machine Learning

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 imitation learning (i.e., apprenticeship learning) in contextual bandits, with offline demonstrations from some expert (optimal, or very good) policy, without explicitly observed rewards. In contrast to prior work, which assumes the demonstrator belongs to a bounded-complexity policy class, we propose relying only on the underlying reward model (i.e., specifying which answers are correct) being in a bounded-complexity class, which we argue is a strictly weaker assumption. We show that likelihood-maximization methods can fail in this setting, and instead present an approach that learns to answer nearly as well as the demonstrator, with sample complexity logarithmic in the cardinality of the reward class. Our method is similar to Syed and Schapire 2007, when adapted to a contextual bandit (i.e., single step) setup, but is a simple one-pass online approach that enjoys an "optimistic rate" (i.e., 1/ε1/\varepsilon when the demonstrator is optimal, versus 1/ε21/\varepsilon^2 in Syed and Schapire), and works even with arbitrarily adaptive demonstrations.

Keywords

Cite

@article{arxiv.2510.15464,
  title  = {Learning to Answer from Correct Demonstrations},
  author = {Nirmit Joshi and Gene Li and Siddharth Bhandari and Shiva Prasad Kasiviswanathan and Cong Ma and Nathan Srebro},
  journal= {arXiv preprint arXiv:2510.15464},
  year   = {2026}
}

Comments

Generalized some results. Updated the presentation in light of an important related work of Syed and Schapire. Improved discussions. Comments are welcome

R2 v1 2026-07-01T06:42:54.117Z