Sharing session at Stanford XCS 236 on April 4th 2025 covering RL algorithms in LLM training (with DeepSeek-R1 as an example)

Reinforcement Learning Algorithms in Deep Generative Model

Name: Cheng Shen
Email: shenchg126@gmail.com
Website: https://scaufish.github.io/

RLHF has become a routine in LLM training these days

"To make our models safer, more helpful, and more aligned, we use an existing technique called reinforcement learning from human feedback (RLHF)⁠."

Source: OpenAI: Aligning language models to follow instructions, 01-2022

"We apply preference modeling and reinforcement learning from human feedback (RLHF) to finetune language models to act as helpful and harmless assistants. We find this alignment training improves performance on almost all NLP evaluations, and is fully compatible with training for specialized skills such as python coding and summarization."

Source: Anthropic: Training a Helpful and Harmless Assistant with Reinforcement Learning from Human Feedback, 04-2022

"We introduce our first-generation reasoning models, DeepSeek-R1-Zero and DeepSeek-R1. DeepSeek-R1-Zero, a model trained via large-scale reinforcement learning (RL) without supervised fine-tuning (SFT) as a preliminary step, demonstrated remarkable performance on reasoning."

Source: DeepSeek: DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning, 01-2025

and there've been interesting discussions

Agenda

1. Introduction to RL
2. Introduction to LLM
3. DeepSeek Training Paradigm
4. Q&A

1. Introduction to RL

Take the classic, yet still impressive AlphaGo as an example DeepMind: AlphaGo

Basic Concepts

• State (\(s \in \mathcal{S}\)): The environment's current situation
  • In a game of Go: \(s \in \{null, white, black\}^{361}\) is current state on the board.
• Action (\(a \in \mathcal{A}\)): The agent's possible moves
  • Assuming the agent plays black, \(a \in \{0, ..., 360\}\)
• Reward (\(r \in \mathbb{R}\)): Immediate feedback signal
  • Reward in Go could be quite sparse: -1/+1 depending on final loss/win
  • Other designs: reward for capturing opponent's stones?
  • AlphaGo: sparse reward depending only the final result

Key Functions

• Policy function (\(\pi(a|s)\)): Probability distribution over actions given state
• Value function (\(V^\pi(s) = \mathbb{E}_\pi[\sum_{t=0}^{\infty} \gamma^t r_t | s_0 = s]\)): Expected cumulative reward starting from state \(s\)
  • The expected likelihood to win/lose given the current state and policy
• Q-function (\(Q^\pi(s,a) = \mathbb{E}_\pi[\sum_{t=0}^{\infty} \gamma^t r_t | s_0 = s, a_0 = a]\)): Expected return after taking action \(a\) in state \(s\)
  • Note how value function and Q-function are related:
  • \(V^\pi(s) = \sum_{a \in \mathcal{A}} \pi(a|s) Q^\pi(s,a)\) or
  • \(V^\pi(s) = \mathbb{E}_{a \sim \pi(\cdot|s)}[Q^\pi(s,a)]\)
• Reward function (\(r(s,a,s')\)): Immediate reward for transition from \(s\) to \(s'\) via action \(a\)

Objective

• Maximize expected cumulative reward: \(J(\pi) = \mathbb{E}_\pi[\sum_{t=0}^{\infty} \gamma^t r_t]\)
• Find optimal policy: \(\pi^* = \arg\max_\pi J(\pi)\)

Approaches

• Value-based: Learn \(V(s)\) or \(Q(s,a)\) (e.g., Q-learning, DQN)
• Policy-based: Directly optimize \(\pi(a|s)\) (e.g., REINFORCE, PPO)
• Actor-Critic: Combine both approaches (e.g. PPO on advantage)

Policy Gradients and PPO

Policy gradient methods learn a parameterized policy that can select actions without consulting a value function.

The policy gradient theorem gives us the gradient of this objective (Lilian's blog):

Where:

• \(\theta\) represents the policy parameters
• \(\tau\) is a trajectory (sequence of states and actions)
• \(p_\theta(\tau)\) is the probability of trajectory \(\tau\) under policy \(\pi_\theta\)
• \(A(s_t, a_t)\) is the advantage function at state \(s_t\) and action \(a_t\)
• \(\log \pi_\theta(a_t|s_t)\) is the log probability of taking action \(a_t\) in state \(s_t\) under policy \(\pi_\theta\)

Where \(A(s_t, a_t)\) is the advantage function, which measures how much better taking action \(a_t\) in state \(s_t\) is compared to the average action according to the current policy:

Key algorithms and concepts:

REINFORCE (Monte Carlo Policy Gradient)

• The most basic policy gradient algorithm
• Updates policy parameters using the gradient: \(\theta \leftarrow \theta + \alpha \nabla_\theta \log \pi_\theta(a_t|s_t) \cdot G_t\)
• Where \(G_t\) is the return from time step \(t\)
• High variance in gradient estimates due to randomness in trajectories

Proximal Policy Optimization (PPO)

• Addresses the problem of large policy updates destabilizing training
• Uses a clipped surrogate objective with advantage function:

The full PPO loss function combines three components:

where:

• \(L^{CLIP}(\theta)\) is the clipped surrogate objective described above
• \(L^{VF}(\theta)\) is the value function loss, typically a squared-error loss:
  \[L^{VF}(\theta) = (V_\theta(s_t) - V_t^{target})^2\]
• \(S[\pi_\theta](s)\) is an entropy bonus to encourage exploration:
  \[S[\pi_\theta](s) = -\sum_a \pi_\theta(a|s) \log \pi_\theta(a|s)\]
• \(c_1\) and \(c_2\) are coefficients that balance the importance of each term

This combined objective encourages both policy improvement and accurate value estimation while maintaining sufficient exploration.

where:

• \(r_t(\theta) = \frac{\pi_\theta(a_t|s_t)}{\pi_{\theta_{old}}(a_t|s_t)}\) is the probability ratio
• \(A_t\) is the advantage estimate at timestep \(t\)
• The clipping prevents the ratio from moving too far from 1, limiting the policy update size

Comparisions

Key advantages of policy gradient methods:

• Can learn stochastic policies, which may be optimal in partially observable environments
• Can handle continuous action spaces naturally
• Can learn policies with arbitrary action distributions
• More stable learning in many environments compared to value-based methods

Challenges:

• High variance in gradient estimates
• Sample inefficiency
• Sensitivity to hyperparameters
• Potential for premature convergence to suboptimal policies

Try them out!

• Gymnasium
• Huggingface RL Algorighms
• ...

# python3.9, gym=0.26.2, stable_baselines3=2.6.0
import gym
from stable_baselines3 import PPO
from stable_baselines3.common.evaluation import evaluate_policy

# Create the CartPole environment
env = gym.make("CartPole-v1")

# Instantiate the PPO model with an MLP policy
model = PPO("MlpPolicy", env, verbose=1)

# Train the model for 10,000 timesteps
model.learn(total_timesteps=10000)

# Evaluate the trained model
mean_reward, std_reward = evaluate_policy(model, env, n_eval_episodes=10)
print(f"Mean reward: {mean_reward:.2f} +/- {std_reward:.2f}")

# Save the trained model
model.save("ppo_CartPole")

# To demonstrate loading, delete the existing model and load it back
del model
model = PPO.load("ppo_CartPole", env=env)

# Enjoy the trained agent
obs, _ = env.reset()
for _ in range(1000):
    action, _ = model.predict(obs, deterministic=True)
    obs, reward, done, truncated, info = env.step(action)
    env.render()
    if done or truncated:
        obs, _ = env.reset()

env.close()

2. Introduction to LLM

2.1 Autoregressive generative model

Autoregressive language models generate text by predicting one token at a time, with each prediction conditioned on all previously generated tokens. This sequential generation process can be formalized as:

Where:

• \(p(x)\) is the probability of generating the entire sequence \(x\)
• \(x_t\) is the token at position \(t\)
• \(x_{<t}\) represents all tokens before position \(t\)
• \(T\) is the total sequence length
• Note how this is analogous to policy function in RL: \(\pi(a|s)\)

Key principles:

• Conditional probability: Each token is predicted based on the context of all previous tokens
• Maximum likelihood training: Models are trained to maximize the log probability of predicting the correct next token
• Teacher forcing: During training, ground truth tokens are used as context regardless of the model's predictions
• Attention mechanisms: Modern autoregressive models like GPT use self-attention to capture dependencies between tokens at different positions
• Causal masking: To maintain the autoregressive property, attention is masked to prevent looking at future tokens

The training objective is typically to minimize the negative log-likelihood:

Where \(\theta\) represents the model parameters.

2.2 Transformer and Decoder-only model

Generative Pretrained Transformer visualization by Jay Alammar

2.3 RLHF in LLM

• Current LLM training usually contains following stages: pretraining, supervised fine-tuning (SFT), Reinforcement Learning with Human Feedback / Preference Optimization
• Used for alignment, and improvements where directly labeled data is hard to obtain.
• Human annotators are better at picking the better answer, than writing a good answer -- so is the reward model

with PPO (Proximal Policy Optimization)

1. Collect outputs and human preferences on input \((x, y_1, y_2, b)\), where \(b\) is the human selection. Initialize reward model \(r\) as \(\pi_{SFT}\), and initialize reference model \(\pi_{ref}\) as \(\pi_{SFT}\).
2. Train reward model \(r\) based on human preference (Bradley-Terry Loss). \(r\) is usually initialized with a LLM after SFT \(\pi_{SFT}\).
3. Optimize policy \(\pi\) over modified reward \(R = r(x, y) + \beta \mathbf{D}_{KL}(\pi_\theta(y|x) | \pi_{ref}(y|x))\)
4. Collect more human feedback on samples generated by \(\pi_{ref}\), repeat step 2-3.

Simplification with DPO (Direct Preference Optimization)

• Proposed specifically for LLM RLHF stage.
• No need to train a separate reward model, as the reward model can be rewritten in the form with policy.
• Leverages offline generated samples and human labels.

3. DeepSeek Training Paradigm

3.1 DeepSeek-R1 and DeepSeek-R1-Zero visualization by Harry Chan

• DeepSeek-V3 Base: Trained on 14.8T tokens
  • Good old next-token prediction
  • Fill-in-Middle (FIM) strategy on 10% of data
• DeepSeek-V3:
  • Supervised Fine-tuning (SFT): 1.5M instances of instruction
  • Reinforcement Leraning (GRPO)
• Reward model is more result-oriented
  • Leverage rule-based reward for coding and math questions
  • Reward model trained with human feedback on non-reasoning tasks

3.2 GRPO

GRPO Paper

Simplified loss function compared to PPO

PPO (simplified for gradient analysis)

DPO

GRPO

3.3 RL for reasoning

Performance

• Aha Moment, and longer thinking process
• Mixture of different languages Anthropic's research about the languages in LLM's head.
  • alignment (to enforce consistent language) results in a slight degradation in the model's performance
• Observed reward hacking on neural reward model -- DeepSeek-R1-Zero is trained purely with rule-based reward
• Failure with MCTS
  • Inspired by AlphaGo
  • Echos the question -- what is a unit action in LLM, token? sequence? or block?
• Dislike of few-shot prompting

4. Q&A