Reward Normalization

mediumPROReward ModelingLLM 2: Training & Alignment

Implement reward normalization using a running mean and variance.

During RLHF training, reward scores are normalized to have zero mean and unit variance using running statistics. Given a sequence of reward batches, update running mean and variance using exponential moving average (EMA):

$\mu_t = (1 - \alpha) \cdot \mu_{t-1} + \alpha \cdot \bar{r}_t$ $\sigma^2_t = (1 - \alpha) \cdot \sigma^2_{t-1} + \alpha \cdot \text{var}(r_t)$

Normalized reward: $\hat{r} = (r - \mu_t) / \sqrt{\sigma^2_t + \epsilon}$

Input:

Line 1: alpha epsilon (EMA decay, stability constant)
Line 2: N (number of batches)
Next N lines: space-separated reward values for each batch

Output: Normalized rewards of the LAST batch, rounded to 4 decimal places.

Example

Input

0.1 1e-8
2
1.0 2.0 3.0
4.0 5.0 6.0

Output

0.8885 1.9438 2.9991

Reasoning

We initialize the running mean $\mu_0$ and variance $\sigma^2_0$ to 0, and read the EMA decay $\alpha = 0.1$ and stability constant $\epsilon = 1e-8$ .
For the first batch $[1.0, 2.0, 3.0]$ , we calculate the mean $\bar{r}_1 = (1.0 + 2.0 + 3.0) / 3 = 2.0$ and variance $\text{var}(r_1) = (1.0^2 + 2.0^2 + 3.0^2) / 3 - 2.0^2 = 1.0$ , then update the running mean and variance using the EMA formulas: $\mu_1 = (1 - 0.1) \cdot 0 + 0.1 \cdot 2.0 = 0.2$ and $\sigma^2_1 = (1 - 0.1) \cdot 0 + 0.1 \cdot 1.0 = 0.1$ .
For the second batch $[4.0, 5.0, 6.0]$ , we calculate the mean $\bar{r}_2 = (4.0 + 5.0 + 6.0) / 3 = 5.0$ and variance $\text{var}(r_2) = (4.0^2 + 5.0^2 + 6.0^2) / 3 - 5.0^2 = 1.0$ , then update the running mean and variance: $\mu_2 = (1 - 0.1) \cdot 0.2 + 0.1 \cdot 5.0 = 0.52$ and $\sigma^2_2 = (1 - 0.1) \cdot 0.1 + 0.1 \cdot 1.0 = 0.19$ .
We normalize the rewards in the last batch using the updated running mean and variance: $\hat{r} = (r - 0.52) / \sqrt{0.19 + 1e-8}$ , resulting in the normalized rewards $[0.8885, 1.9438, 2.999

Constraints

Initialize μ₀ = 0, σ²₀ = 1
0 < alpha < 1, epsilon = 1e-8
Round to 4 decimal places

Learn this in

Safety & EthicsLLM Study Plan RLHF & AlignmentLLM Study Plan

More from LLM 2: Training & Alignment

RLHF Reward Shapinghard Implicit Reward from DPOmedium PPO Clipped Objectivemedium DPO Lossmedium KL Divergence Penaltyeasy

Concept primerPREVIEW

Background Knowledge

The problem revolves around reward normalization, a crucial concept in Reinforcement Learning from Human Feedback (RLHF). In RLHF, the goal is to train a model to align with human preferences by providing rewards for desired behaviors. However, these rewards can have varying scales and distributions, making it challenging to compare and combine them. Normalization helps to standardize these rewards, ensuring that the model learns from them effectively.

The problem uses exponential moving average (EMA) to update the running mean ( $\mu_t$ ) and variance ( $\sigma^2_t$ ) of the rewards. EMA is a technique for calculating the mean of a set of values over time, giving more weight to recent values. This is useful in tracking changes in the reward distribution over time. The normalized reward ( $\hat{r}$ ) is then calculated by subtracting the running mean and dividing by the square root of the running variance plus a small stability constant ( $\epsilon$ ).

Understanding the mathematical formulas provided is essential to solving this problem. The formulas for updating the running mean and variance using EMA are: $\mu_t = (1 - \alpha) \cdot \mu_{t-1} + \alpha \cdot \bar{r}_t$ $\sigma^2_t = (1 - \alpha) \cdot \sigma^2_{t-1} + \alpha \cdot \text{var}(r_t)$ where $\alpha$ is the EMA decay rate, $\bar{r}_t$ is the mean of the current batch of rewards, and $\text{var}(r_t)$ is the variance of the current batch of rewards.

Algorithm/Approach

The general approach to solving this problem involves:

Initializing the running mean and variance
Iterating over each batch of rewards, updating the running mean and variance using EMA
Calculating the normalized rewards for each batch
Returning the normalized rewards of the last batch

PRO

Continue the full explanation

You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.