DEVOURED
On-Policy Distillation
On-policy distillation trains a student model using its own policy's sampled trajectories, with a teacher providing token-level supervision via KL-based regularization, effectively addressing train-inference distribution mismatch common in off-policy methods.
What: On-policy distillation is a machine learning technique that trains a smaller "student" model on data sampled from its own behavior (policy), while a larger "teacher" model provides detailed, token-level guidance using KL-based regularization. This method, which can be implemented with a one-line code change on an RL stack like Tinker, unifies forward-KL, reverse-KL, and JSD losses.
Why it matters: This technique is significant for optimizing AI models, especially in reinforcement learning contexts, by ensuring the student model learns effectively from data relevant to its own evolving behavior, thus making smaller models more robust and efficient.
Takeaway: If you are developing or fine-tuning AI models using distillation, particularly in reinforcement learning, consider integrating on-policy distillation with reverse-KL regularization to improve student model performance and close distribution gaps.
Decoder
- On-policy distillation: A machine learning method where a smaller "student" model is trained using data (trajectories) generated by its own current policy, guided by a larger "teacher" model.
- Student model: A smaller, often less complex AI model that is being trained to replicate the behavior or knowledge of a larger, more powerful "teacher" model.
- Teacher model: A larger, more performant AI model whose knowledge and behavior are transferred to a smaller "student" model during distillation.
- Policy: In reinforcement learning, the strategy that an agent uses to decide what actions to take in a given state.
- Trajectories: Sequences of states, actions, and rewards experienced by an agent in an environment.
- Token-level supervision: Guidance provided by the teacher model at the granularity of individual tokens (e.g., words or sub-word units) in the output sequence.
- KL-based regularization (Kullback-Leibler divergence): A measure of how one probability distribution diverges from a second, expected probability distribution, used here to guide the student model's outputs towards the teacher's.
- Train-inference distribution mismatch: A problem where the data distribution encountered during model training differs from the distribution encountered during actual deployment (inference), leading to performance degradation.
- Off-policy methods: Reinforcement learning methods that can learn from data generated by a different policy than the one being optimized.
- Forward-KL: A specific form of KL divergence where the student's distribution is compared against the teacher's, useful for mode-covering.
- Reverse-KL: A specific form of KL divergence where the teacher's distribution is compared against the student's, useful for mode-seeking, especially for smaller student models.
- JSD (Jensen-Shannon Divergence): A method of measuring the similarity between two probability distributions, a symmetric and smoothed version of KL divergence.
- RL stack: A software framework or set of libraries used for developing and deploying reinforcement learning algorithms (e.g., Tinker).
Original article
On-policy distillation trains a student model on trajectories sampled from its own policy while a teacher provides dense token-level supervision through KL-based regularization, closing the train-inference distribution mismatch that off-policy methods suffer. The canonical formulation unifies forward-KL, reverse-KL, and JSD losses with reverse-KL emerging as the default for mode-seeking smaller students, and a one-line code swap of the regularizer model on top of an RL stack like Tinker implements the technique.
Explore all 40 articles in today's digest
Fresh dev news, AI-digested daily. Never miss a trend at devoured.fyi