Coach Models Improve Language Model Training

Summary & Methodology Analysis

When training language models on open-ended tasks using reinforcement learning, standard methods compress rubric-based evaluations into a scalar reward. This process discards rich textual feedback and conflates responses with distinct quality profiles, especially near the top of the reward scale where fine-grained preferences are lost. To address this, the paper repurposes the feedback large language model from an LLM-as-a-Judge to an LLM-as-a-Coach. For each on-policy response generated by the policy model to a given prompt, the coach analyzes the rubric-based assessment and distills transferable insights into experiential knowledge. This experiential knowledge is then provided as context to a teacher model. The policy model is trained using on-policy context distillation to minimize the token-level reverse KL divergence, which measures the difference between two probability distributions, between its own distribution and the context-conditioned teacher distribution. The policy parameters are updated by minimizing the loss function based on this KL divergence.

In terms of architecture and execution, the approach integrates several models and datasets including WildChat-IF, GPT-4o, Qwen3-8B, OLMo-3-7B-Instruct, Tulu3, AlpacaEval v2.0, WildBench, ArenaHard v2.0, CreativeWritingV3, GPT-4-Turbo, Qwen3-1.7B, Qwen3-4B, and IFEval. The authors note important mathematical bounds, observing that a discrete 1 to 10 scalar reward has a theoretical maximum bandwidth of 3.3 bits per sample, whereas a learned reward head producing a bfloat16 value provides at most 16 representational bits. Among the teacher context ablation variants, the default Experiential Learning setting that extracts transferable experiential knowledge performed best in improving the WildChat test set score. Iteratively updating the teacher improves the WildChat score from 80.0 to 80.7, but substantially reduces IFEval accuracy. Meanwhile, adding general-prompt On-Policy Distillation from the initial frozen policy checkpoint improves IFEval accuracy from 74.9 to 79.9 while preserving the WildChat score.

Despite these performance gains, the method faces several limitations. The calculated feedback-bandwidth is an intuition and theoretical maximum rather than a measure of usable supervision, because natural-language redundancy and imperfect knowledge extraction can make the usable information substantially smaller. Furthermore, iteratively updating the teacher model substantially reduces IFEval accuracy, indicating forgetting on out-of-distribution instruction following. Using Full Critique or Rubrics Only as teacher context degrades IFEval accuracy because evaluation-oriented context biases the teacher toward a critiquing distribution, causing the policy to imitate this misaligned behavior. Additionally, a separate source of reward hacking can occur if the feedback model itself is biased or miscalibrated, which is not addressed by Experiential Learning since it focuses on feedback bandwidth rather than feedback quality. Finally, the controlled distribution-matching analysis is not intended as a faithful simulation of full reinforcement learning and experiential learning training pipelines, nor as proof of a fundamental limitation shared by all scalar-reward methods.