@@ -9,7 +9,7 @@ In this subsection, we introduce the fundamentals of supervised fine-tuning for
Although pre-trained LLMs possess a vast amount of general knowledge, they are typically limited to tasks related to language modeling. Our ultimate goal is to enable them to perform various specific tasks, such as summarization, question answering, and machine translation.
One straightforward method to achieve this goal is supervised fine-tuning (SFT), in which the pre-trained LLM is further trained on a dataset comprising task-specific input (i.e., instruction+user input)\footnote{In this paper, the \textit{instruction} represents the description of a specific task, e.g., ``Summarize the following article''; the \textit{user input} represents the extra content to the instruction, e.g., ``Article: In recent years, solar energy has seen a significant increase in the amount of energy available to the public. recent years, solar energy has seen unprecedented growth, becoming the fastest-growing ...''} paired with their expected outputs \citep{ouyang:2022training,wei-etal:2022finetuned}.
Specifically, let $\mathbf{x}=x_0...x_m$ be an input (e.g., instruction + user input) and $\mathbf{y}=y_1...y_n$ be the corresponding output. In SFT, we aim to maximize the probability of the output $\mathbf{y}$ given the input $\mathbf{x}$. Consider an LLM with pre-trained parameters $\hat{\theta}$. The fine-tuning objective can then be formulated as:
Specifically, let $\mathbf{x}=x_1...x_m$ be an input (e.g., instruction + user input) and $\mathbf{y}=y_1...y_n$ be the corresponding output. In SFT, we aim to maximize the probability of the output $\mathbf{y}$ given the input $\mathbf{x}$. Consider an LLM with pre-trained parameters $\hat{\theta}$. The fine-tuning objective can then be formulated as:
\INPUT the optimized reward model $r_{\phi}(\cdot)$, the reference model $\pi_{\theta_{\mathrm{ref}}}(\cdot)$, the initialized value model $V_{\omega}(\cdot)$, the policy model $\pi_{\theta}(\cdot)$, the input-only dataset $S_{x}$;
\INPUT the optimized reward model $R_{\phi}(\cdot)$, the reference model $\pi_{\theta_{\mathrm{ref}}}(\cdot)$, the initialized value model $V_{\omega}(\cdot)$, the policy model $\pi_{\theta}(\cdot)$, the input-only dataset $\mathcal{S}_{x}$;
\State sample an output from $\pi_{\theta_\mathrm{old}}$: $\mathbf{y}\sim\pi_{\theta_\mathrm{old}}(\cdot|\mathbf{x})$
\State compute the reward for $[\mathbf{x},\mathbf{y}]$: $r_{\phi}(\mathbf{x},\mathbf{y})$
\State compute the reward for $[\mathbf{x},\mathbf{y}]$: $R_{\phi}(\mathbf{x},\mathbf{y})$
\State\%\%\% roll out
\For{ppo\_epoch\_idx=1 to $\mathrm{PPO\_EPOCH}$}
\State\%\%\%\textit{compute policy model loss and value model loss through advantages and returns}
...
...
@@ -15,7 +15,7 @@
\For{t=1 to T}
\State compute the penalty through Eq. (\ref{eq:penalty}): $\mathrm{Penalty}_{t}=\log\mathrm{Pr}_{\theta}(y_t|\mathbf{x},\mathbf{y}_{<t})-\log\mathrm{Pr}_{\theta_{\mathrm{ref}}}(y_t|\mathbf{x},\mathbf{y}_{<t})$
\State\hspace{0.5cm} predict the next time step value $V_{\omega}(\mathbf{x},\mathbf{y}_{<t+1},y_{t+1})$
\State\textbf{end if}
\State\%\%\%\textit{compute policy model loss}
\State compute the advantage $A(\mathbf{x},\mathbf{y}_{<t},y_{t})$ through Eq. (\ref{eq:advantage}): $A(\mathbf{x},\mathbf{y}_{<t},y_{t})=r_{t}+\gamma V_{\omega}(\mathbf{x},\mathbf{y}_{<t+1},y_{t+1})-V_{\omega}(\mathbf{x},\mathbf{y}_{<t},y_{t})$
\State compute the advantage $A(\mathbf{x},\mathbf{y}_{<t},y_{t})$: $A(\mathbf{x},\mathbf{y}_{<t},y_{t})=r_{t}+\gamma V_{\omega}(\mathbf{x},\mathbf{y}_{<t+1},y_{t+1})-V_{\omega}(\mathbf{x},\mathbf{y}_{<t},y_{t})$
\State compute the policy model loss through Eq. (\ref{eq:ppo-loss}) and add it to $\mathrm{loss}_{p}$: $\mathrm{loss}_{p}=\mathrm{loss}_{p}+\mathrm{Clip}\Big(\frac{\pi_{\theta}(y_t|\mathbf{x},\mathbf{y}_{<t})}{\pi_{\theta_{\mathrm{old}}}(y_t|\mathbf{x},\mathbf{y}_{<t})}\Big) A(\mathbf{x},\mathbf{y}_{<t},y_t)$
\State\%\%\%\textit{compute value model loss}
\State compute the return $\mathrm{Return}_{t}$: $\mathrm{Return}_{t}=A(\mathbf{x},\mathbf{y}_{<t},y_t)+V_{\omega}(\mathbf{x},\mathbf{y}_{<t},y_{t})$
@@ -173,7 +173,7 @@ In practice, there are many ways to implement the value function. One simple app
V_{t}& = &\mathbf{h}_{y_{t}}\mathbf{W}_{v}
\end{eqnarray}
where $\mathbf{h}_{y_{t}}$ is a $d$-dimensional vector, and $\mathbf{W}{v}$ is a $d \times1$ linear mapping matrix. This value model is typically trained using the Temporal Difference (TD) error. This training method relies on the principle that the difference between the values at adjacent time steps should correspond to the reward received at the former time step, i.e., $V_t - V_{t+1}= r_t$. Suppose the value model is parameterized by $\omega$. Given a sequence $\mathrm{seq}_{\mathbf{x}, \mathbf{y}}$, the loss function is given by
where $\mathbf{h}_{y_{t}}$ is a $d$-dimensional vector, and $\mathbf{W}_{v}$ is a $d \times1$ linear mapping matrix. This value model is typically trained using the Temporal Difference (TD) error. This training method relies on the principle that the difference between the values at adjacent time steps should correspond to the reward received at the former time step, i.e., $V_t - V_{t+1}= r_t$. Suppose the value model is parameterized by $\omega$. Given a sequence $\mathrm{seq}_{\mathbf{x}, \mathbf{y}}$, the loss function is given by
where $\gamma\in[0,1]$ is the discount factor that adjusts the importance of future rewards. When $\gamma$ is set to less than 1, it signifies that early rewards are considered more important than future rewards. This basic idea is also applied in other fields. For example, in LLMs, it has been demonstrated that early token generation plays a crucial role, as it can influence the style and accuracy of the entire output \citep{wang-and-zhou:2024chain}.
...
...
@@ -227,6 +228,7 @@ Addressing this issue involves considering trust regions in optimization \citep{
At the time step $t$, we can obtain the penalty as
@@ -249,7 +251,7 @@ A further improvement to TRPO involves addressing the gradient variance problem
In Eq. (\ref{eq:importance-sampling-with-penalty}), the range of the ratio function is from $[0, +\infty)$, which could lead to high gradient variance in the learning process. To mitigate this problem, clipping is commonly employed to limit the magnitude of importance weights, thereby preventing excessively large updates and promoting stability in the learning process. A clipped version can be given by
where $\mathbf{W}{r}$ is a $d \times1$ linear mapping matrix, and $\phi$ represents the parameters of the reward model, which includes both the parameters of the LLM and the $\mathbf{W}{r}$. This architecture of the reward model is illustrated in Figure \ref{fig:reward-model}.
where $\mathbf{W}_{r}$ is a $d \times1$ linear mapping matrix, and $\phi$ represents the parameters of the reward model, which includes both the parameters of the LLM and the $\mathbf{W}_{r}$. This architecture of the reward model is illustrated in Figure \ref{fig:reward-model}.
\begin{figure*}[!t]
\centering
...
...
@@ -331,7 +333,7 @@ where $(\mathbf{x},\mathbf{y}_a,\mathbf{y}_b)$ is drawn from a human-annotated d
\hat{\phi}& = &\argmin_{\phi}\mathcal{L}_r(\phi)
\end{eqnarray}
Since the reward model itself is also an LLM, we can directly reuse the Transformer training procedure to optimize the reward model. The difference from training a standard LLM is that we only need to replace the cross-entropy loss with the pairwise comparison loss as illustrated in Figure \ref{fig:training-reward-models}. After the training of the reward model, we can apply the trained reward model $r_{\hat{\phi}}(\cdot)$ to supervise the target LLM for alignment.
Since the reward model itself is also an LLM, we can directly reuse the Transformer training procedure to optimize the reward model. The difference from training a standard LLM is that we only need to replace the cross-entropy loss with the pairwise comparison loss as illustrated in Figure \ref{fig:training-reward-models}. After the training of the reward model, we can apply the trained reward model $R_{\hat{\phi}}(\cdot)$ to supervise the target LLM for alignment.
It is worth noting that although we train the reward model to perform pairwise ranking, we apply it to score each input-output pair independently during the alignment process. The pairwise ranking objective ensures that the reward model is sensitive to subtle differences between outputs, but we rely on the continuous scores produced by the reward model to guide the optimization of the LLM. An advantage of this approach is that we can choose from or combine various ranking loss functions and still apply the resulting reward models in the same way as we have done in Section \ref{sec:improved-reward-generalization}. However, a challenge arises with this method: the reward model can provide only sparse rewards; that is, it offers a delayed reward rather than an intermediate one. Hence, in this case, we can obtain
@@ -28,7 +28,7 @@ As illustrated in Figure \ref{fig:bon-sampling}, in the BoN sampling, we first s
where the $\mathrm{argTopN}$ operation returns the top-$N$ outputs that maximize the function $\Pr(\mathbf{y}|\mathbf{x})$. These outputs can be sampled in various ways, depending on the search algorithm used by the model (e.g., nucleus sampling or beam search). Once the $N$-best output candidates are sampled, the reward model is used to evaluate and select the best one:
This process not only identifies the output with the highest reward but also makes it a direct evaluation of the effectiveness of the reward model. For example, the agreement between the reward model and human judgments can be assessed by judging whether $\mathbf{y}_{\mathrm{best}}$ matches the best output selected by humans. Given its speed and cost-efficiency compared to more complex evaluation methods such as those used in PPO, this approach is widely used for evaluating the performance of reward models \citep{rafailov:2023direct,gao-etal:2023scaling}. Nevertheless, in scenarios without a reward model, we can also employ the majority vote approach to determine the final output of the reasoning task. The basic steps involve first identifying the answer that most reasoning paths converge upon and then considering any of the outputs leading to this answer as the final output.
...
...
@@ -46,7 +46,7 @@ We can collect or generate reasoning paths corresponding to problems from existi
Given a set of step-level annotated reasoning paths and corresponding inputs, we can train a reward model to provide a reward for each step in the reasoning process. The reward model can be treated as a classification model. So its architecture can be an LLM with a Softmax layer stacked on top, akin to the architecture depicted in Figure \ref{fig:reward-model}. Here, consider an reasoning path including $n_{s}$ steps, represented as $\mathbf{y}=\{\bar{\mathbf{y}}_{1},\cdots ,\bar{\mathbf{y}}_{n_{s}}\}$. At each step $k$, the process reward model takes both the problem description, denoted by $\mathbf{x}$, and the reasoning steps generated so far, denoted by $\bar{\mathbf{y}}$, as inputs. It then outputs a probability distribution over the set of labels ${\text{\textit{correct}}, \text{\textit{incorrect}}}$, or ${\text{\textit{correct}}, \text{\textit{incorrect}}, \text{\textit{neutral}}}$, to evaluate the reasoning at that point. This model can trained by a casual classification loss, e.g., Sigmoid Cross-entropy. Once trained, the reward model can be used to evaluate reasoning paths by assessing the correctness of each step. A simple method to use log-probabilities of classification to define the reward of each reasoning step, for example, the reward of the $k$-th reasoning step can be given by
where $\Pr_{\phi}(\textit{correct}|\mathbf{x},\bar{\mathbf{y}}_{\le k})$ denotes the probability of the \textit{correct} label generated by the reward model. The reward score $R_{\phi}(\mathbf{x},\mathbf{y})$ can then be used to select the best step while generating a reasoning path. Additionally, as discussed in Section \ref{sec:generative-reward-models}, there is the option to train a generative process reward model that further improves this performance on evaluating the reasoning step, also called generative verifier in the literature \citep{zhang-etal:2024generative}. Note that in practice, the process reward model serves not only to provide rewards for test-time scaling but also to train the model using RL, e.g., the rewards from this model can be employed as shaping rewards.
...
...
@@ -149,7 +149,7 @@ Compared with conventional offline distillation, OPD adopts a simple but importa
This implementation of OPD is relatively simple, as illustrated in Figure~\ref{fig:opd-implementation}. Given an input $\mathbf{x}$, we first sample an output $\mathbf{y}=\{y_1,\ldots,y_n\}$ from the current student model $\mathrm{Pr}_{\theta}(\cdot|\mathbf{x})$. For each token position $i$, the prefix $(\mathbf{x},\mathbf{y}_{<i})$ denotes a state actually visited by the student model. We then feed the same state into the teacher model $\mathrm{Pr}_{\mathrm{tea}}$ and obtain its probability distribution over the next token. The student model is optimized to match this teacher distribution. This OPD objective can defined as:
