update.

section4/Figures/figure-examples-of-rule-based-reward-length.tex

@@ -2,6 +2,8 @@
 % !TEX encoding = UTF-8 Unicode
 
 \begin{center}
+  \begingroup
+  \newcommand{\jsonq}{\texttt{\char34}}
   \begin{tikzpicture}[background rectangle/.style={fill=black!5!white,rounded corners=5pt,minimum width=\textwidth}, show background rectangle]
     \begin{scope}
       \matrix (m) [matrix anchor=west,
@@ -9,21 +11,21 @@
       nodes={anchor=north west,}
       ] at (-\textwidth/2,0) {
         \color{gray}{Input} &
-        \begin{varwidth}{18em}\small{Give me three tips to improve my accuracy in solving math problems. Return in JSON format with ``tip1'', ``tip2'', ``tip3'' as the key and the corresponding tip as the value.  (\textbf{Expect a JSON output})}
+        \begin{varwidth}{18em}\small{Give me three tips to improve my accuracy in solving math problems. Return in JSON format with \jsonq{}tip1\jsonq{}, \jsonq{}tip2\jsonq{}, \jsonq{}tip3\jsonq{} as the key and the corresponding tip as the value.  (\textbf{Expect a JSON output})}
         \end{varwidth} \\
         \color{gray}{Output 1} &
         \begin{varwidth}{18em}
           \small{ \{\\
-            \hspace*{.4cm}"tip1": " Delve into the labyrinth of equations with the focus of a hawk hunting ...", \\
-            \hspace*{.4cm}"tip2": " Engage in the art of problem-solving like a master alchemist turning ...", \\
-            \hspace*{.4cm}"tip3": "Practice, practice, practice, until the numbers themselves bow in reverence ..." \\
+            \hspace*{.4cm}\jsonq{}tip1\jsonq{}: \jsonq{} Delve into the labyrinth of equations with the focus of a hawk hunting ...\jsonq{}, \\
+            \hspace*{.4cm}\jsonq{}tip2\jsonq{}: \jsonq{} Engage in the art of problem-solving like a master alchemist turning ...\jsonq{}, \\
+            \hspace*{.4cm}\jsonq{}tip3\jsonq{}: \jsonq{}Practice, practice, practice, until the numbers themselves bow in reverence ...\jsonq{} \\
             \} (\textbf{Valid: True})}
         \end{varwidth} \\
         \color{gray}{Output 2} &
         \begin{varwidth}{18em}
-          \small{tip1: " Delve into the labyrinth of equations with the focus of a hawk hunting ...", \\
-            tip2: " Engage in the art of problem-solving like a master alchemist turning ...", \\
-            tip3: "Practice, practice, practice, until the numbers themselves bow in reverence ..."\\
+          \small{tip1: \jsonq{} Delve into the labyrinth of equations with the focus of a hawk hunting ...\jsonq{}, \\
+            tip2: \jsonq{} Engage in the art of problem-solving like a master alchemist turning ...\jsonq{}, \\
+            tip3: \jsonq{}Practice, practice, practice, until the numbers themselves bow in reverence ...\jsonq{}\\
             (\textbf{Valid: False})}
         \end{varwidth} \\
       };
@@ -40,4 +42,5 @@
       \draw [->] ([xshift=.1cm]counter.east|-m-3-2.east) -- (res2.west);
     \end{scope}
   \end{tikzpicture}
+  \endgroup
 \end{center}

section4/section4.tex

@@ -93,7 +93,7 @@ where $r'(\cdot)$ is the transformed reward, $r(\cdot)$ is the original delayed
 f(\mathbf{x}, \mathbf{y}_{<t}, y_{t}) & = & \gamma \Phi(\mathbf{x}, \mathbf{y}_{<t+1}, y_{t+1}) - \Phi(\mathbf{x}, \mathbf{y}_{<t}, y_{t}) \label{eq:shaping-reward-function}
 \end{eqnarray}
 
-where $\Phi(\cdot)$ is called the potential value function. If we define $\Phi(\cdot)$ as the common value function and substitute Eq.~(\ref{eq:shaping-reward-function}) into Eq.~(\ref{eq:transformed-reward-function}), we obtain
+where $\Phi(\cdot)$ is called the potential value function, and $\gamma$ is the discount factor that controls how much the next-step potential contributes to the current shaping reward. If we define $\Phi(\cdot)$ as the common value function and substitute Eq.~(\ref{eq:shaping-reward-function}) into Eq.~(\ref{eq:transformed-reward-function}), we obtain
 \begin{eqnarray}
 r'_{t} & = & r_{t} + \gamma V_{t+1} - V_{t}
 \end{eqnarray}
@@ -114,15 +114,6 @@ In addition to reward shaping, another method to address the sparse reward issue
 \subsubsection{Improved Reward Generalization}
 \label{sec:improved-reward-generalization}
 
-\begin{figure*}[!t]
-  \centering
-  \input{section4/Figures/figure-methods-improving-reward-generalization}
-  \caption{
-    When training reward models, we can use parameter freezing and regularization methods to preserve the features of the LLM and thereby improve reward generalization.
-  }
-  \label{fig:improve-reward-generalization}
-\end{figure*}
-
 As discussed in Section~\ref{sec:training-reward-models}, reward models are trained using preference data and subsequently used to optimize LLMs. However, a problem arises: the data distribution during LLM optimization may differ from the distribution of the preference data, making it challenging for the reward model to generalize to unseen input-output pairs. 
 
 A well-known failure mode associated with this problem is commonly referred to as \textit{overoptimization} or \textit{reward hacking}, where the optimization stage improves the reward model score but deteriorates the alignment with true rewards \citep{gao-etal:2023scaling,eisenstein-etal:2023helping}. This mode occurs because the reward model may incorrectly assign high rewards to unseen input-output pairs, leading the LLM to learn and optimize for behaviors that do not truly align with the desired behaviors and objectives. For example, consider a scenario from Section~\ref{sec:policy-gradient} where the reward model is designed to favor informative and accurate outputs for a homework assistant. However, if the reward model fails to generalize to an overly verbose output and assigns a high reward to it (perhaps due to its length or certain keywords), the LLM may learn to prioritize generating a long output, which is not actually more informative but receives a higher reward according to the reward model. Consequently, the LLM becomes misaligned with its true objectives, delivering concise and relevant information, because it optimizes for the incorrect reward signal from the reward model with weak generalization. 
@@ -143,6 +134,14 @@ On the other hand, to improve the generalization of reward models, it is importa
 
 where $\alpha$ is a balancing factor. We further illustrate the parameter freezing and regularization methods in Figure~\ref{fig:improve-reward-generalization}. Note that the regularization term applies only to the LLM. That is, when optimizing this term, only the parameters of the LLM are updated, while the parameters of the reward linear map remain fixed. Therefore, in this equation, the parameter associated with the regularization term is $\theta$, which specifically refers to the parameters of the LLM.
 
+\begin{center}
+  \refstepcounter{figure}
+  \label{fig:improve-reward-generalization}
+  \input{section4/Figures/figure-methods-improving-reward-generalization}
+  \vspace{-0.5em}
+  {\small \textbf{Figure~\thefigure:} When training reward models, we can use parameter freezing and regularization methods to preserve the features of the LLM and thereby improve reward generalization.}
+\end{center}
+
 \begin{figure}[!t]
   \centering
   \input{section4/Figures/figure-generative-reward-model-architecture}

section5/section5.tex

@@ -116,7 +116,7 @@ Training LLMs with RL usually follows a two-phase approach: training a pre-train
 \end{itemize}
 
 
-An interesting issue arises with this design of iterative RL: why is RL aimed at enhancing reasoning capabilities placed at the initial phase rather than at the end? This is partly because the rewards for the reasoning task, which are inherently more stable as described in Section~\ref{sec:lightweight-reward-methods}, are well-suited for large-scale RL compared to traditional reward models. Furthermore, the reasoning capabilities could be applicable across many tasks, allowing for large-scale learning to boost performance in subsequent multi-task learning. We consider that this foundational enhancement in reasoning capabilities may contribute to the observed robust performance across various tasks, even without the use of labeled data in later phases. In this case, large-scale RL could be viewed as an approach to continued pre-training.
+An interesting issue arises with this design of iterative RL: why is RL aimed at enhancing reasoning capabilities placed before the later rejection-sampling and multi-task RL phases rather than postponed to the end? The key reason is that the later phases depend on a model that can already explore meaningful reasoning trajectories. If the model has not first acquired this ability, rejection sampling is likely to draw from low-quality reasoning paths, and multi-task RL has to optimize general alignment and task performance while also trying to repair basic reasoning behavior. Placing reasoning-oriented RL early therefore expands the pool of useful trajectories for subsequent data construction and makes later optimization stages more effective. This ordering is also practical because reasoning rewards, which are inherently more stable as described in Section~\ref{sec:lightweight-reward-methods}, are well-suited for large-scale RL. Furthermore, reasoning capabilities can be applicable across many tasks, allowing the early RL phase to serve as a broad capability-building stage before later task-specific refinement. In this case, large-scale RL could be viewed as an approach to continued pre-training.
 
 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.
 

section7/section7.tex

@@ -125,7 +125,37 @@ where $\mathbf{v}=\mathbf{x}_1-\mathbf{x}_0$ denotes the target velocity field a
 \frac{d\mathbf{x}_t}{dt} = \mathbf{v}_\theta(\mathbf{x}_t,t)
 \end{eqnarray}
 
-Similar to diffusion models, the continuous generation process of flow matching models can also be formulated as an MDP to enable RL optimization. Specifically, by discretizing the continuous flow trajectory, we can consider the generation process  as a sequence of decision steps, where the model determines how to transform the current state toward the target data distribution \citep{liu-etal:flow}. Under this formulation, the state, action, transition, initial state distribution, and reward function are defined as follows. The state at timestep $t$ is defined as:
+Although the continuous generation process of flow matching models can also be formulated as an MDP to enable RL optimization, this formulation faces a unique difficulty. 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 make RL applicable, Flow-GRPO converts the original ODE-based flow matching model into an equivalent stochastic differential equation (SDE) that preserves the marginal distributions while introducing stochasticity into the generation process \citep{liu-etal:flow}. The drift function of the converted SDE can be constructed from the original flow velocity field as:
+\begin{eqnarray}
+\mathbf{v}^{d}_{\theta}(\mathbf{x}_t,t,\mathbf{z}) =
+\mathbf{v}_{\theta}(\mathbf{x}_t,t,\mathbf{z}) +
+\frac{\sigma_t^2}{2t}
+\left[
+\mathbf{x}_t + (1-t)\mathbf{v}_{\theta}(\mathbf{x}_t,t,\mathbf{z})
+\right]
+\end{eqnarray}
+
+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:
+\begin{eqnarray}
+d\mathbf{x}_t = \mathbf{v}^{d}_{\theta}(\mathbf{x}_t,t,\mathbf{z})dt +
+\sigma_t d\mathbf{w}
+\end{eqnarray}
+
+where $d\mathbf{w}$ denotes the Wiener process increment, i.e., the infinitesimal Gaussian noise added along the continuous sampling path.
+
+In practice, we discretize the above SDE for trajectory sampling. Using Euler--Maruyama discretization, each sampling step can be written as:
+\begin{eqnarray}
+\mathbf{x}_{t+\Delta t} = \mathbf{x}_t + \mathbf{v}^{d}_{\theta}(\mathbf{x}_t,t,\mathbf{z})\Delta t +
+\sigma_t\sqrt{\Delta t}\boldsymbol{\epsilon}
+\quad
+\boldsymbol{\epsilon}\sim\mathcal{N}(\mathbf{0},\mathbf{I})
+\end{eqnarray}
+
+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 corresponds to a stochastic Gaussian transition, which makes it possible to compute the transition probabilities required by GRPO.
+
+Specifically, by discretizing the continuous flow trajectory, we can consider the generation process as a sequence of decision steps, where the model determines how to transform the current state toward the target data distribution \citep{liu-etal:flow}. Under this formulation, the state, action, transition, initial state distribution, and reward function are defined as follows. The state at timestep $t$ is defined as:
 \begin{eqnarray}
 s_t=(\mathbf{z},t,\mathbf{x}_t)
 \end{eqnarray}
@@ -140,8 +170,7 @@ Since the flow model deterministically predicts the velocity field, the policy c
 \mathrm{Pr}_{\theta}(a_t|s_t) = \delta(a_t-(\mathbf{x}_t-\Delta t\,\mathbf{v}_{\theta}(\mathbf{x}_t,t,\mathbf{z})))
 \end{eqnarray}
 
-where $\delta(\cdot)$ denotes the Dirac delta distribution, indicating that the next state is deterministically determined by the current state and the learned velocity field.
-However, this deterministic formulation introduces challenges for RL optimization because policy gradient methods typically require stochastic sampling to explore different trajectories. To address this issue, Flow-GRPO converts the original ODE-based flow matching model into an equivalent stochastic differential equation (SDE), which preserves the marginal distribution while introducing sampling diversity \citep{liu-etal:flow}. Under this formulation, the transition probability becomes:
+where $\delta(\cdot)$ denotes the Dirac delta distribution, indicating that the next state is deterministically determined by the current state and the learned velocity field. After the ODE-to-SDE conversion and discretization described above, the transition probability becomes:
 \begin{eqnarray}
 \mathrm{Pr}_\theta(a_t|s_t) = \mathcal{N} (a_t;
 \mu_\theta(\mathbf{x}_t,t,\mathbf{z}),
@@ -163,35 +192,3 @@ Based on these estimated advantages, we formulate the GRPO objective as follows:
 \mathrm{clip}(\frac{\mathrm{Pr}_\theta(a_t^i|s_t^i)}{\mathrm{Pr}_{\theta_{\mathrm{old}}}(a_t^i|s_t^i)},1-\epsilon,1+\epsilon)\hat{A}_i)\right]
 \end{eqnarray}
 
-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:
-\begin{eqnarray}
-d\mathbf{x}_t = \mathbf{v}^d_\theta(\mathbf{x}_t,t,\mathbf{z})dt +
-\sigma_t d\mathbf{w}
-\end{eqnarray}
-
-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:
-\begin{eqnarray}
-\mathbf{v}^{d}_{\theta}(\mathbf{x}_t,t,\mathbf{z}) =
-\mathbf{v}_{\theta}(\mathbf{x}_t,t,\mathbf{z}) +
-\frac{\sigma_t^2}{2t}
-\left[
-\mathbf{x}_t + (1-t)\mathbf{v}_{\theta}(\mathbf{x}_t,t,\mathbf{z})
-\right]
-\end{eqnarray}
-
-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:
-\begin{eqnarray}
-d\mathbf{x}_t = \mathbf{v}^{d}_{\theta}(\mathbf{x}_t,t,\mathbf{z})dt +
-\sigma_t d\mathbf{w}
-\end{eqnarray}
-
-In practice, we discretize the above SDE for trajectory sampling. Using Euler--Maruyama discretization, each sampling step can be written as:
-\begin{eqnarray}
-\mathbf{x}_{t+\Delta t} = \mathbf{x}_t + \mathbf{v}^{d}_{\theta}(\mathbf{x}_t,t,\mathbf{z})\Delta t +
-\sigma_t\sqrt{\Delta t}\boldsymbol{\epsilon}
-\quad
-\boldsymbol{\epsilon}\sim\mathcal{N}(\mathbf{0},\mathbf{I})
-\end{eqnarray}
-
-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.