Chain-of-Thought Reasoning and Reasoning RL

Different types of reasoning

• Chain-of-Thought (CoT) Reasoning with LLMs: Early researches use "srcatchpad" method to break the problem into intermediate steps. Later, other work prompts a strong model to "think step by step", which was found significantly improving.

• Reason with expert iteration : The Self-Taught Reasoner (STaR) [Zelikman et al., 2022] frames reasoning as a bootstrapping loop: a pretrained model first samples diverse chains-of-thought (CoTs), keeps only those that lead to correct answers, and then finetunes on these “expert” traces. Iterating this cycle can improve the LM’s reasoning capabilities and solve rate. STaR demonstrated that this version of expert iteration [Anthony et al., 2017] using automatic, string match–based verification of generated answers can bootstrap reasoning skills without human-written reasoning traces.

• Reasoning RL with verified rewards, o1 and R1: OpenAI o1, Deepseek R1, KIMI1.5, using Policy gradient methods to train on math and code tasks where string matching or unit tests verify correctness.

SFT

Through SFT, we observed that we can improve the performance of our SFT model by filtering out bad examples from the SFT data.

Expert Iteration

Policy gradient

for LM, given $s_t$ as an input or state, $a_t$ as an output under the state, we can view LM as a categorical stochastic policy.

$$a_t\sim {\pi }_{\theta }(·|s_t),{\pi }_{\theta }(a_t|s_t)=[softmax(f_{\theta }(s_t)){]}_{a_t}$$

Trajectory

We call $s_{t+1}=s_t\parallel a_t$ as a trajectory (or episodes or rollouts).

Rewards and Return

A scalar reward $r_t=R(s_t,a_t)$ judges the immediate quality of the action taken at state $s_t$.

$$r_T=R(s_T,a_T):=\{\begin{array}{ll}1& \text{if the trajectory }{s}_T\parallel a_T\text{ matches the ground-truth according to our reward function}\\ 0& \text{otherwise.}\end{array}$$

finite-horizon undiscounted returns:

$$R(\tau ):=\sum _{t=0}^T{r}_t$$

and infinite-horizon discounted returns:

$$R(\tau ):=\sum _{t=0}^{\mathrm{\infty }}{\gamma }^t{r}_t,0<\gamma <1.$$

In our case, we will use the undiscounted formulation since episodes have a natural termination point (end-of-text or max generation length). The objective of the agent is to maximize the expected return:

$$J(\theta )={\mathbb{E}}_{\tau \sim {\pi }_{\theta }}[R(\tau )],$$

leading to the optimization problem:

$${\theta }^{\ast }=\arg \underset{\theta }{max}J(\theta ).$$

In one word, 最大化策略的回报期望

Vanilla Policy Gradient

Next, let us attempt to learn policy parameters $\theta$ with gradient ascent on the expected return:

$${\theta }_{k+1}={\theta }_k+\alpha {\mathrm{\nabla }}_{\theta }J({\theta }_k).$$

The core identity that we will use to do this is the REINFORCE policy gradient, shown below:

$${\mathrm{\nabla }}_{\theta }J({\pi }_{\theta })={\mathbb{E}}_{\tau \sim {\pi }_{\theta }}[\sum _{t=0}^T{\mathrm{\nabla }}_{\theta }\log {\pi }_{\theta }(a_t|s_t)R(\tau )].$$

Deriving the policy gradient

How did we get this equation? For completeness, we will give a derivation of this identity below. We will make use of a few identities.

1. The probability of a trajectory is given by

   $$P(\tau |\theta )={\rho }_0(s_0)\prod _{t=0}^{T}P(s_{t+1}|s_t,a_t){\pi }_{\theta }(a_t|s_t)$$

   Therefore, the log-probability of a trajectory is:

   $$\log P(\tau |\theta )=\log {\rho }_0(s_0)+\sum _{t=0}^T[\log P(s_{t+1}|s_t,a_t)+\log {\pi }_{\theta }(a_t|s_t)]$$

2. The log-derivative trick:

   $${\mathrm{\nabla }}_{\theta }P=P{\mathrm{\nabla }}_{\theta }\log P$$

3. The environment terms are constant in $\theta$. ${\rho }_0,P(\cdot |\cdot )$ and $R(\tau )$ do not depend on the policy parameters, so

   $${\mathrm{\nabla }}_{\theta }{\rho }_0={\mathrm{\nabla }}_{\theta }P={\mathrm{\nabla }}_{\theta }R(\tau )=0$$

Applying the facts above:

$$\begin{array}{rl}{\mathrm{\nabla }}_{\theta }J(\theta )& ={\mathrm{\nabla }}_{\theta }{\mathbb{E}}_{\tau \sim {\pi }_{\theta }}[R(\tau )]\\ & ={\mathrm{\nabla }}_{\theta }\sum _{\tau }P(\tau |\theta )R(\tau )\\ & =\sum _{\tau }{\mathrm{\nabla }}_{\theta }P(\tau |\theta )R(\tau )\\ & =\sum _{\tau }P(\tau |\theta ){\mathrm{\nabla }}_{\theta }\log P(\tau |\theta )R(\tau )\text{(Log-derivative trick)}\\ & ={\mathbb{E}}_{\tau \sim {\pi }_{\theta }}[{\mathrm{\nabla }}_{\theta }\log P(\tau |\theta )R(\tau )]\\ & ={\mathbb{E}}_{\tau \sim {\pi }_{\theta }}[\sum _{t=0}^T{\mathrm{\nabla }}_{\theta }\log {\pi }_{\theta }(a_t|s_t)R(\tau )]\end{array}$$

Intuitively, this gradient will increase the log probability of every action in a trajectory that has high return, and decrease them otherwise.

Sample estimate of the gradient. Given a batch of $N$ rollouts $\mathcal{D}=\{{\tau }^{(i)}{\}}_{i=1}^N$ collected by sampling a starting state $s_0^{(i)}\sim {\rho }_0(s_0)$ and then running the policy ${\pi }_{\theta }$ in the environment, we form an unbiased estimator of the gradient as:

$$\hat{g}=\frac{1}{N}\sum _{i=1}^N\sum _{t=0}^{T_i}{\mathrm{\nabla }}_{\theta }\log {\pi }_{\theta }(a_t^{(i)}|s_t^{(i)})R({\tau }^{(i)})$$

where $T_i$ is the length of the $i$-th trajectory, which may vary across samples. This vector is used in the gradient-ascent update: $\theta \leftarrow \theta +\alpha \hat{g}$.

Policy Gradient Baseline

当我们引入了Vanilla Policy Gradient（REINFORCE）基础策略梯度算法后，我们开始考虑它的不足。一个很大的问题是， 策略梯度不稳定，方差（variance）非常大，这会导致收敛缓慢。

为了减小方差，一个常用的技巧是引入 Baseline (基准) $b(s_t)$。

引入 Baseline 的策略梯度

带 Baseline 的策略梯度公式如下：

$$B={\mathbb{E}}_{\tau \sim {\pi }_{\theta }}[\sum _{t=0}^T{\mathrm{\nabla }}_{\theta }\log {\pi }_{\theta }(a_t|s_t)(R(\tau )-b(s_t))]$$

一个合理的 Baseline 选择是状态值函数 (On-policy value function) $V^{\pi }(s)={\mathbb{E}}_{\tau \sim {\pi }_{\theta }}[R(\tau )|s_t=s]$。直观上，$(R(\tau )-V^{\pi }(s_t))$ 衡量了实际观测到的回报比预期的好多少。

无偏性证明

只要 Baseline $b(s_t)$ 仅依赖于状态 $s_t$ 而不依赖于具体的动作 $a_t$，引入它就不会产生偏差。我们可以通过展开期望来证明这一点：

$$B={\mathbb{E}}_{\tau \sim {\pi }_{\theta }}[\sum _{t=0}^T{\mathrm{\nabla }}_{\theta }\log {\pi }_{\theta }(a_t|s_t)R(\tau )]-{\mathbb{E}}_{\tau \sim {\pi }_{\theta }}[\sum _{t=0}^T{\mathrm{\nabla }}_{\theta }\log {\pi }_{\theta }(a_t|s_t)b(s_t)]$$

我们要证明减号后面的那一项等于 0。根据全期望公式，我们可以把那一项重写为：

$${\mathbb{E}}_{\tau \sim {\pi }_{\theta }}[\sum _{t=0}^T{\mathrm{\nabla }}_{\theta }\log {\pi }_{\theta }(a_t|s_t)b(s_t)]=\sum _{t=0}^T{\mathbb{E}}_{s_t}[b(s_t){\mathbb{E}}_{a_t\sim {\pi }_{\theta }(\cdot |s_t)}[{\mathrm{\nabla }}_{\theta }\log {\pi }_{\theta }(a_t|s_t)]]$$

这里最关键的数学恒等式是：Score Function 的期望为 0。

Score Function 期望为 0 的证明

对于任何概率分布 $P_{\theta }(x)$，关于其参数 $\theta$ 的记分函数 (Score Function) ${\mathrm{\nabla }}_{\theta }\log P_{\theta }(x)$ 在原分布下的期望恒等于 0：

$${\mathbb{E}}_{x\sim P_{\theta }}[{\mathrm{\nabla }}_{\theta }\log P_{\theta }(x)]=0$$

证明过程如下：

$$\begin{array}{rl}{\mathbb{E}}_{x\sim P_{\theta }}[{\mathrm{\nabla }}_{\theta }\log P_{\theta }(x)]& =\int P_{\theta }(x){\mathrm{\nabla }}_{\theta }\log P_{\theta }(x)dx\\ & =\int P_{\theta }(x)\frac{{\mathrm{\nabla }}_{\theta }{P}_{\theta }(x)}{P_{\theta }(x)}dx\\ & =\int {\mathrm{\nabla }}_{\theta }{P}_{\theta }(x)dx\\ & ={\mathrm{\nabla }}_{\theta }\int P_{\theta }(x)dx\\ & ={\mathrm{\nabla }}_{\theta }(1)\\ & =0\end{array}$$

由于 ${\mathbb{E}}_{a_t\sim {\pi }_{\theta }(\cdot |s_t)}[{\mathrm{\nabla }}_{\theta }\log {\pi }_{\theta }(a_t|s_t)]=0$，整个基准项的期望确实为 0。 这意味着 $B={\mathrm{\nabla }}_{\theta }J(\theta )$，即引入基准后，梯度的期望值保持不变（无偏），但方差可以显著降低。

pg_loss

pg_loss is not a loss in the canonical sense—it’s not meaningful to report pg_loss on the train or validation set as an evaluation metric, and a good validation pg_loss doesn’t indicate that our model is generalizing well. Instead, it is a surrogate objective function whose gradient is the policy gradient estimator.

$$\text{pg\_loss}(\theta )=\frac{1}{N}\sum _{i=1}^N\sum _{t=0}^{T_i}\log {\pi }_{\theta }(a_t^{(i)}|s_t^{(i)})(R({\tau }^{(i)})-b(s_t^{(i)}))$$

When doing RL, you should always log and report train and validation rewards. These are the “meaningful” evaluation metrics and what we are attempting to optimize with policy gradient method

Off-Policy Policy Gradient

之前的vanilla policy以及添加了baseline的方法都是典型的on-policy方法，即每次更新策略参数$\theta$后，都需要根据policy model新的 $\theta$ 重新收集一批新的数据来计算梯度。这导致了数据利用率低的问题。为了解决这个问题，我们可以使用off-policy方法.

较为典型的有 PPO 和 GRPO ，他们都是通过收集旧policy的 rollouts来更新新policy的参数。

$${\hat{g}}_{off\mathrm{\_}policy}=\frac{1}{N}\sum _{i=1}^N\sum _{t=0}^{T_i}{\displaystyle \frac{{\pi }_{\theta }(a_t^{(i)}|s_t^{(i)})}{{\pi }_{{\theta }_{old}}(a_t^{(i)}|s_t^{(i)})}}{\mathrm{\nabla }}_{\theta }\log {\pi }_{\theta }(a_t^{(i)}|s_t^{(i)})R({\tau }^{(i)})$$

当${\pi }_{\theta }$ 和 ${\pi }_{{\theta }_{old}}$相差不大时，上面的权重系数是合理的

Group Realtive Plocy Optimization(GRPO)

Baseline设置：对于一个问题，我们采样多条rollouts，计算它们的group-normalized reward. For a question $q$ and group outputs ${o^{(i)}}_{i=1}^G\sim {\pi }_{\theta }(\cdot |q)$, let $r^{(i)}=R(q,o^{(i)})$.

$$A=r_{norm}^{(i)}=\frac{r^{(i)}-mea{n}_r^{(j)}}{st{d}_j{r}^{(j)}+ϵ}$$

对于同一个output $i$中的所有token, 他们的$A^{(i)}$是相同的

GRPO-Clip Objective

Let us first write out the full GRPO-Clip objective, and then we can build some intuition on what the clipping does:

$$J_{\text{GRPO-Clip}}(\theta )={\mathbb{E}}_{q\sim \mathcal{D},\{o^{(i)}{\}}_{i=1}^G\sim {\pi }_{\theta }(\cdot |q)}\left[\frac{1}{G}\sum _{i=1}^G\frac{1}{|o^{(i)}|}\sum _{t=1}^{|o^{(i)}|}\underset{\text{per-token objective}}{\underbrace{min(\frac{{\pi }_{\theta }(o_t^{(i)}|q,o_{<t}^{(i)})}{{\pi }_{{\theta }_{\text{old}}}(o_t^{(i)}|q,o_{<t}^{(i)})}{A}^{(i)},\text{clip}(\frac{{\pi }_{\theta }(o_t^{(i)}|q,o_{<t}^{(i)})}{{\pi }_{{\theta }_{\text{old}}}(o_t^{(i)}|q,o_{<t}^{(i)})},1-ϵ,1+ϵ)A^{(i)})}}\right]$$

The hyperparameter $ϵ>0$ controls how much the policy can change. To see this, we can rewrite the per-token objective in a more intuitive way following Achiam [2018a, b]. Define the function

$$g(ϵ,A^{(i)})=\{\begin{array}{ll}(1+ϵ)A^{(i)}& \text{if }{A}^{(i)}\ge 0\\ (1-ϵ)A^{(i)}& \text{if }{A}^{(i)}<0\end{array}$$

We can rewrite the per-token objective as

$$\text{per-token objective}=min(\frac{{\pi }_{\theta }(o_t^{(i)}|q,o_{<t}^{(i)})}{{\pi }_{{\theta }_{\text{old}}}(o_t^{(i)}|q,o_{<t}^{(i)})}{A}^{(i)},g(ϵ,A^{(i)}))$$

We can now reason by cases. When the advantage $A^{(i)}$ is positive, the per-token objective simplifies to

$$\text{per-token objective}=min(\frac{{\pi }_{\theta }(o_t^{(i)}|q,o_{<t}^{(i)})}{{\pi }_{{\theta }_{\text{old}}}(o_t^{(i)}|q,o_{<t}^{(i)})},1+ϵ)A^{(i)}$$

Since $A^{(i)}>0$, the objective goes up if the action $o_t^{(i)}$ becomes more likely under ${\pi }_{\theta }$, i.e., if ${\pi }_{\theta }(o_t^{(i)}|q,o_{<t}^{(i)})$ increases. The clipping with $min$ limits how much the objective can increase: once ${\pi }_{\theta }(o_t^{(i)}|q,o_{<t}^{(i)})>(1+ϵ){\pi }_{{\theta }_{\text{old}}}(o_t^{(i)}|q,o_{<t}^{(i)})$, this per-token objective hits its maximum value of $(1+ϵ)A^{(i)}$. So, the policy ${\pi }_{\theta }$ is not incentivized to go very far from the old policy ${\pi }_{{\theta }_{\text{old}}}$.

Analogously, when the advantage $A^{(i)}$ is negative, the model tries to drive down ${\pi }_{\theta }(o_t^{(i)}|q,o_{<t}^{(i)})$, but is not incentivized to decrease it below $(1-ϵ){\pi }_{{\theta }_{\text{old}}}(o_t^{(i)}|q,o_{<t}^{(i)})$ (refer to Achiam [2018b] for the full argument).