@@ -129,4 +129,48 @@ An interesting issue arises with this design of iterative RL: why is RL aimed at
Since the optimization objective of each phase is different, the design of iterative RL can also be analyzed and understood from the perspective of multi-objective optimization \citep{wang-etal:2024hybrid}. In the context of multi-objective optimization, iterative RL for LLMs can be likened to interactive methods where the solution process is iterative, and preferences are actively defined and refined by the decision-maker during the search for the most preferred solutions \citep{miettinen-etal:2008introduction,deb-etal:2016multi}. More specifically, in this scenario, RL can be seen as a decision-maker, iteratively refining and enhancing the different capabilities of the LLM across different phases.
\subsection{On-Policy Distillation}
% TODO
\ No newline at end of file
\subsection{On-Policy Distillation for Reasoning}
While process supervision provides step-by-step reward signals to guide reasoning, it still relies on a scalar value to tell the model how well it is doing. To provide an even richer training signal, we can employ \mindex{on-policy distillation} (\mindex{OPD}). This method is generally used in post-training of LLMs, and aligns well with the goal of step-based supervision: making the training signals as dense as possible.
Standard knowledge distillation can be viewed a kind of \mindex{off-policy distillation}. A smaller ``student'' model learns by mimicking the reasoning paths generated by a stronger ``teacher'' model. The student is trained on a fixed dataset of these pre-recorded paths. However, reasoning is a long and fragile process. If the student makes a small mistake or chooses a slightly different valid approach during actual generation, it enters a state it has never seen in the training data. Because the student does not know how to recover from its mistakes, the reasoning process may collapse. This phenomenon is known as exposure bias or distribution shift.
The core idea of OPD is that the student model learns from its own generated trajectories, rather than from trajectories sampled from the teacher model. Formally, let the teacher model be $\mathrm{Pr}^{t}(\cdot)$ and the student model (parameterized with $\theta$) be $\mathrm{Pr}_{\theta}^{s}(\cdot)$. The typical procedure of OPD consists of two steps:
\begin{itemize}
\item\vspace{0.5em}\textbf{Student Rollout}. Given an input $\mathbf{x}$, the student model uses its current parameters $\theta$ to generate a reasoning path $\mathbf{y}=\{\bar{\mathbf{y}}_1, \dots, \bar{\mathbf{y}}_{n_p}\}$. Because the student generates this path itself, it naturally explores its own reasoning behaviors. This includes making slight deviations, taking alternative approaches, or even making mistakes that a teacher would never make.
\item\vspace{0.3em}\textbf{Teacher Guidance}. We then feed this student-generated path to the teacher model. At each step $t$, instead of providing a scalar reward, the teacher model calculates the probability distribution for the next step $\mathrm{Pr}^{t}(\cdot | \mathbf{x}, \bar{\mathbf{y}}_{<t})$.
\end{itemize}
\vspace{0.5em}
Ideally, the student model is updated to align its step-level probability distribution with the distribution generated by the teacher model. The objective is to minimize the expected KL divergence between the two distributions over the reasoning paths generated by the student model:
\noindent However, calculating this step-level distribution is computationally intractable in practice. Because a reasoning step $\bar{\mathbf{y}}_t$ typically consists of multiple tokens, the probability distribution over all possible multi-token steps resides in an exponentially large space.
To address this, we fall back to the standard OPD approach, which operates at the token level. By treating the generated reasoning path as a token sequence $\mathbf{y}=\{y_1,\dots, y_n\}$, the objective is reformulated to compute the KL divergence token by token over the vocabulary:
\noindent This objective is then approximated using sampling. For a given input $\mathbf{x}$, we simply perform a rollout to generate one (or a small batch of) reasoning paths from the student model, and calculate the token-level KL divergence on these sampled rollouts.
Figure \ref{fig:comparision-of-sft-ps-opd} compares OPD with supervised fine-tuning and process supervision. To further understand the difference among these methods, imagine a student trying to solve a complex geometry problem:
% \caption{Comparison of supervised fine-tuning, process supervision, and on-policy distillation methods. In supervised fine-tuning (a), the student model is trained off-policy to maximize the likelihood of golden target outputs provided by a human or a teacher model. In process supervision (b), the student model generates its own reasoning path, and a verifier evaluates these self-generated steps to compute advantages. In on-policy distillation (c), the student also generates its own path, but instead of scalar rewards, a teacher model provides a dense, token-level probability distribution at each step. The student is updated by minimizing the KL divergence between its own predicted distribution and the teacher's distribution.}
% \label{fig:comparision-of-sft-ps-opd}
% \end{figure}
\begin{itemize}
\item\vspace{0.5em}\textbf{Supervised Fine-tuning}. This is like giving the student a textbook filled with the correct solutions written by the teacher to memorize. If the student decides to draw an auxiliary line differently than the textbook did, the textbook offers no help on how to proceed from that new starting point.
\item\vspace{0.5em}\textbf{Process Supervision} (RL with PRMs). This is like having a strict grader who checks the student's work step-by-step and says, ``+1 point for step 1'', and ``-1 point for step 2''. It tells the student where they went wrong, but not how to fix it or what the correct next step should be.
\item\textbf{On-Policy Distillation}. This is like having an expert tutor sitting right next to the student. When the student draws that unusual auxiliary line (a state not in the standard textbook), the tutor looks at the student's current state and says, ``Okay, since you drew the line here, your next best step is to calculate this specific angle.''
\end{itemize}
Because the supervision of OPD comes in the form of a full probability distribution over the vocabulary, the student receives a very dense and rich learning signal. It learns not only what the single best next step is, but also the probabilities of other valid alternatives. This rich feedback is well suited for reasoning because it teaches the student model how to recover from deviations and maintain logical consistency across the entire output sequence, rather than just treating a rationale as an isolated component. However, applying on-policy distillation comes with a high computational cost. The process requires running a large teacher model alongside the student during training to calculate probability distributions in real time. This is more expensive than standard off-policy distillation, where the outputs of the teacher model are only generated once and stored in a static dataset.
During training, one challenge is trajectory sampling
\ No newline at end of file
During training, trajectory sampling introduces a unique challenge for flow matching models. Diffusion models naturally support stochastic sampling because Gaussian noise is progressively involved in the generation process. In contrast, flow matching models typically generate samples through deterministic ODEs, where the intermediate states are uniquely determined once the initial noise is given. This deterministic process creates two difficulties for online RL: it makes the transition probability required by policy-gradient optimization difficult to compute, and, more importantly, provides limited randomness for policy exploration.
To address these issues, we can convert the deterministic ODE into an equivalent stochastic differential equation (SDE) that preserves the marginal distributions while introducing stochasticity into the generation process \citep{liu-etal:flow}. The resulting sampling process can be generally written as:
where $\mathbf{v}^d_\theta(\cdot)$ is determined by the learned velocity field and $d\mathbf{w}$ denotes the Wiener process increment. Here, the drift function of the converted SDE can be constructed from the original flow velocity field as:
where $\mathbf{v}_{\theta}(\cdot)$ denotes the original velocity field and $\mathbf{v}^{d}_{\theta}(\cdot)$ denotes the drift function after the ODE-to-SDE conversion. Therefore, the stochastic generation process can be written as:
The last term injects stochasticity into each generation step, allowing the model to sample diverse trajectories from the same condition and thereby enabling effective exploration for online RL. After discretization, each generation step therefore corresponds to a stochastic Gaussian transition, allowing the model to sample diverse trajectories and compute the transition probabilities required by GRPO.
