Training Systems · Local CPULaptop CPULesson 3 of 5 in Phase 18

Phase 18 / Post-training: SFT, LoRA, reward models and RLHF

Learn a reward model from pairwise preferences

Revision v0.7.0 · 2026-10-20

Lesson · about 105 minutes · CPU practice

Represent a preference as a comparison. Fit score differences with a stable likelihood.

Phase study guide: preparation, milestones and project evidence

Before you start

Before this: LoRA: adapt, save and merge low-rank updates
Hardware: CPU

Download notebook

Other downloads

Python worked examples Lesson text

Setup help: Python, notebook, or Cloud TPU VM

Use the tested CPU packages (the recorded environment uses Python 3.14.3). Other Python versions and accelerator environments require separate checks. The optional Cloud TPU specialization explains the target-specific execution boundaries. Logical CPU devices practice partitioning; they do not simulate TPU performance.

Download the repository ZIP or use your existing checkout. Extract it and open a terminal in the folder containing requirements-cpu.txt. The beginner ZIP is separate; these lessons use the full repository folder.

macOS / Linux terminal · repository folder

python3 -m venv .venv
source .venv/bin/activate
python -m pip install -r requirements-cpu.txt

PowerShell · repository folder

py -3 -m venv .venv
.\.venv\Scripts\python.exe -m pip install -r requirements-cpu.txt

After activation, run the complete worked example from that same folder:

macOS / Linux terminal · activated environment

python phases/18-posttraining/03-reward/code/main.py

PowerShell command · repository folder

.\.venv\Scripts\python.exe phases/18-posttraining/03-reward/code/main.py

For an incremental build, create main.py in your own practice folder, append each build block in order, and run python main.py with this environment active. Keep that file separate from the downloaded reference solution.

Download this lesson’s notebook. Open it with a Python kernel using the course environment. Run cells from the top after restarting the kernel; experiments depend on earlier definitions.

For local Jupyter, first install the course requirements using the Local Python steps, then use your environment’s Python to start Jupyter:

macOS / Linux terminal · activated course environment

python -m pip install jupyterlab
python -m jupyter lab

PowerShell command · repository folder

.\.venv\Scripts\python.exe -m pip install jupyterlab
.\.venv\Scripts\python.exe -m jupyter lab

Select the downloaded posttraining-03.ipynb file in JupyterLab. The notebook contains reference solutions; make your prediction and attempt each practice task before opening them. Restart the kernel before changing logical device setup.

Check: The model learns all three declared preferences; common reward offsets preserve probabilities.

Revisit the guided CPU setup · Open 19: TPU Setup. Package instructions are CPU instructions; Windows and accelerator environments need their own validation.

Phase 18.03 Stage 1 of 5 · Understand

1 Understand 2 Build & Run 3 Experiment 4 Exercise & Check 5 Evidence

What you will be able to do

Represent a preference as a comparison
Fit score differences with a stable likelihood
Verify the changed-condition exercise and explain the limits of this fixture.

The problem

You know which of two responses is preferred, but you do not have an absolute score for either. How can a model learn from that comparison, and what does its score mean?

The idea

A pairwise reward model learns to give a preferred item a higher score than its alternative. The loss uses a score difference. That relative objective does not automatically calibrate absolute quality or make preference labels infallible.

A preference compares scores rather than calibrating utility

If chosen and rejected scores are $3$ and $1$ , their margin is $2$ . Adding $10$ to both scores leaves the margin unchanged and therefore leaves a difference-based preference loss unchanged. The absolute offset is not identified by those comparisons.

Conflicting labels can make an intermediate preference probability more appropriate than an extreme one. Inspect disagreement and source quality rather than assuming the model should confidently fit every pair.

The margin plot shows separation on the fixture. To support downstream policy optimization, evaluate generalization and possible exploitation separately. A reward that ranks a few training pairs correctly can still be unreliable on policy-generated outputs.

Before coding: a preference is a comparison

A training example says that one response was preferred to another for the same prompt and rubric. It does not assign a universally meaningful numeric score to either response. Begin by subtracting the rejected score from the chosen score. A positive margin supports the label; a negative margin contradicts it.

The fixture uses three synthetic response feature vectors and three ordered pairs. It teaches comparison likelihood and score identifiability, not human judgment. Keep the prompt/group identity when replacing these vectors with real response representations; comparing unrelated prompts can teach a meaningless ranking.

Represent a preference as a comparison

Keep the prompt identity, chosen/rejected response identities, label source and split. Pairwise ratings are conditional on a shared prompt and rubric; comparing unrelated prompts can teach a spurious score offset. The fixture has three response feature vectors with an explicit ordering.

Fit score differences with a stable likelihood

The probability of preferring a chosen response is the sigmoid of its reward difference. The negative log likelihood is softplus of the negative margin. A tied pair starts at $\log 2$ . A common offset to all rewards cancels, so comparisons alone do not identify an absolute reward origin.

L_{\mathrm{RM}}=\mathbb E\left[\operatorname{softplus}\!\left(-(r_w(x,y^+)-r_w(x,y^-))\right)\right]

Calculate likelihood and its derivative

Let $m=r_\theta(x,y^+)-r_\theta(x,y^-)$ . The probability of the observed preference is $\sigma(m)$ , and its negative log likelihood is $\log(1+e^{-m})$ . At margin zero, probability is one half and loss is $\log2$ . At margin $\log3$ , probability is three quarters and loss is $-\log(3/4)$ . The margin derivative is $\sigma(m)-1$ : confident correct pairs receive a smaller update than uncertain or wrong pairs.

For a linear score, multiply this derivative by the chosen-minus-rejected feature vector and average pairs. That supplies an independent gradient oracle without differentiating the library softplus implementation.

\nabla_w L=\frac1N\sum_i\left(\sigma(m_i)-1\right)(f_i^+-f_i^-)

Disagreement changes the optimum

If the same two responses appear once with each preference direction, the average loss is minimized at a tied score, not at infinite confidence in either label. With unequal vote frequencies, the optimum reflects those frequencies under this likelihood model. Collapsing disputed examples into one apparently certain label discards information.

Group train and evaluation splits by prompt/source. Report agreement, likelihood and calibration under the labeling protocol, and inspect disagreements. A scalar metric can hide systematic preferences that conflict between raters or domains. The next policy optimizer will amplify the score signal, so keep a separate evaluator for qualities the reward model may have missed.

Separate ranking from calibration

Check ordering, held-out pair likelihood and agreement by prompt/source slice. A large positive margin means the model is confident under its likelihood, not that a response is universally good. Separable preferences can drive growing margins; regularization, validation and stopping rules matter even when every training comparison is correct.

Preserve disagreement and data provenance

Real raters may disagree for legitimate reasons. Retain the rating rubric, task, annotator process and uncertainty instead of converting every disagreement into a perfect label. Split prompts and near-duplicate responses by source so a held-out score does not reward memorization. This lab does not collect personal data or human ratings.

Keep the reward model fixed during policy evaluation

The next lesson freezes the learned reward while optimizing a policy. If reward changes mid-comparison, observed improvement may reflect a moving evaluator. Retain a reference policy, independent task metrics and adversarial examples to detect reward exploitation.

Audit both what scores identify and what they do not

Adding one common offset leaves every preference margin unchanged, so pair comparisons cannot identify an absolute zero of reward. Positive scaling preserves ranking but changes probabilities and confidence. The scale also changes the reward-versus-KL tradeoff in policy optimization. Save normalization and model identity when handing scores to the next phase.

A perfectly fitted training ranking can coexist with overconfidence or poor held-out ranking. On separable data the model can keep increasing margins without learning a new ordering. Track weight norms, held-out likelihood and error slices; choose regularization and stopping on development data rather than demanding the smallest possible training loss.

Learn a reward model from pairwise preferences — recorded experiment

Predict before looking: Which final comparison margin should be the largest given the ordering zero above one above two?

Read the figure

The vertical axis is pairwise negative log likelihood in nats. It starts near $0.693=\log2$ , then falls to about $0.087$ . Each point is before an update on the same three synthetic comparisons; it is not a human agreement rate.

The second panel reports final chosen-minus-rejected score margins for the three training comparisons. All are positive, so each ordering agrees with its synthetic label. The zero-versus-two comparison spans the other two and has the largest margin. A common reward offset would leave every bar unchanged; multiplying weights by two would double them without changing ordering. Neither operation is visible in a ranking-only accuracy number.

Connect it to the computation

The initial zero weights tie all response scores. Training increases the chosen-minus-rejected margins. The offset and scaling exercises explain why reward values must be interpreted relative to their model and downstream objective.

Open full-size figure Download data and provenance

The notebook contains the data experiment, plotting code, and recorded figure output. Rerun it to test your changes.

Where to put & run this code: Work in phases/18-posttraining/03-reward/code/main.py (or downloaded exercises/posttraining-03.py / exercises/posttraining-03.ipynb) and execute from the repository root with python3 phases/18-posttraining/03-reward/code/main.py.

1 Define chosen-minus-rejected likelihood

Create main.py in your activated course environment. Paste this block, then run python main.py; function definitions alone print nothing.

main.py · append build step 1

# Step 1 — 1. Define chosen-minus-rejected likelihood: The loss sees feature differences.
import jax
import jax.numpy as jnp
import numpy as np

def preference_loss(w, chosen, rejected):
    # Perform matrix / vector contraction (`@`) to compute `margin`.
    margin = (chosen - rejected) @ w
    return jnp.mean(jax.nn.softplus(-margin))

Code building blocks in this step

x.sum(axis=...) / x.mean(axis=..., keepdims=...) — Reduces values along the named axis (the axis you name is collapsed unless keepdims=True).

The loss sees feature differences. Predict why adding a common score offset later cannot change it.

2 Construct and inspect the three comparisons

Append this block to main.py and run python main.py again. Keep the earlier blocks above it.

main.py · append build step 2

# Step 2 — 2. Construct and inspect the three comparisons: Write the desired order before fitting: response zero above one,...
features = jnp.array([[1.0, 0.0], [0.0, 1.0], [-1.0, -1.0]])
chosen = features[jnp.array([0, 0, 1])]
rejected = features[jnp.array([1, 2, 2])]
w = jnp.zeros(2)
history = []
# Differentiate the objective to obtain `step` via automatic differentiation.
step = jax.jit(
    jax.value_and_grad(lambda w: preference_loss(w, chosen, rejected))
)

Code building blocks in this step

jnp.array(values, dtype=...) — Constructs an immutable device-backed JAX array from Python/NumPy values.
jnp.zeros / jnp.ones(shape, dtype=...) — Allocates a tensor of the given shape initialized with constants.
jax.value_and_grad(loss_fn)(params, ...) — Evaluates both the scalar loss and its gradient PyTree (loss_val, grads) in a single forward+backward pass.
jax.jit(fn) / @jax.jit — Traces fn with abstract shapes and compiles a fused XLA executable cached by input shape and dtype.

Write the desired order before fitting: response zero above one, and one above two. These labels are synthetic.

3 Fit rewards and check their invariances

Append this block to main.py and run python main.py again. Keep the earlier blocks above it.

main.py · append build step 3

# Step 3 — 3. Fit rewards and check their invariances: The host logaddexp oracle checks stable likelihood values.
for _ in range(100):
    value, g = step(w)
    history.append(float(value))
    w = w - 0.1 * g

# Perform matrix / vector contraction (`@`) to compute `rewards`.
rewards = features @ w
assert rewards[0] > rewards[1] > rewards[2]
np.testing.assert_allclose(history[0], np.log(2), atol=1e-6)
assert history[-1] < 0.1
# Convert `margins` to a host NumPy array for inspection or verification.
margins = np.asarray((chosen - rejected) @ w, dtype=np.float64)
# Execute `np.testing.assert_allclose(`.
np.testing.assert_allclose(
    preference_loss(w, chosen, rejected),
    np.mean(np.logaddexp(0, -margins)),
    atol=1e-6,
)
# Only differences are identified: a common score offset leaves preference probabilities unchanged.
np.testing.assert_allclose(
    jax.nn.sigmoid((chosen @ w + 9) - (rejected @ w + 9)),
    jax.nn.sigmoid((chosen - rejected) @ w),
    atol=1e-6,
)
print(
    'Synthetic preference NLL initial/final:',
    history[0],
    history[-1],
    '; rewards:',
    np.asarray(rewards),
)

Code building blocks in this step

x.sum(axis=...) / x.mean(axis=..., keepdims=...) — Reduces values along the named axis (the axis you name is collapsed unless keepdims=True).
jnp.allclose(actual, expected, rtol=..., atol=...) — Checks that two arrays match elementwise within floating-point tolerance.

The host logaddexp oracle checks stable likelihood values. The score-offset test establishes non-identifiability, not calibration.

View the cumulative build in main.py

These blocks append in order to the same file.

main.py · cumulative Python build

# Step 1 — 1. Define chosen-minus-rejected likelihood: The loss sees feature differences.
import jax
import jax.numpy as jnp
import numpy as np

def preference_loss(w, chosen, rejected):
    # Perform matrix / vector contraction (`@`) to compute `margin`.
    margin = (chosen - rejected) @ w
    return jnp.mean(jax.nn.softplus(-margin))

# Step 2 — 2. Construct and inspect the three comparisons: Write the desired order before fitting: response zero above one,...
features = jnp.array([[1.0, 0.0], [0.0, 1.0], [-1.0, -1.0]])
chosen = features[jnp.array([0, 0, 1])]
rejected = features[jnp.array([1, 2, 2])]
w = jnp.zeros(2)
history = []
# Differentiate the objective to obtain `step` via automatic differentiation.
step = jax.jit(
    jax.value_and_grad(lambda w: preference_loss(w, chosen, rejected))
)

# Step 3 — 3. Fit rewards and check their invariances: The host logaddexp oracle checks stable likelihood values.
for _ in range(100):
    value, g = step(w)
    history.append(float(value))
    w = w - 0.1 * g

# Perform matrix / vector contraction (`@`) to compute `rewards`.
rewards = features @ w
assert rewards[0] > rewards[1] > rewards[2]
np.testing.assert_allclose(history[0], np.log(2), atol=1e-6)
assert history[-1] < 0.1
# Convert `margins` to a host NumPy array for inspection or verification.
margins = np.asarray((chosen - rejected) @ w, dtype=np.float64)
# Execute `np.testing.assert_allclose(`.
np.testing.assert_allclose(
    preference_loss(w, chosen, rejected),
    np.mean(np.logaddexp(0, -margins)),
    atol=1e-6,
)
# Only differences are identified: a common score offset leaves preference probabilities unchanged.
np.testing.assert_allclose(
    jax.nn.sigmoid((chosen @ w + 9) - (rejected @ w + 9)),
    jax.nn.sigmoid((chosen - rejected) @ w),
    atol=1e-6,
)
print(
    'Synthetic preference NLL initial/final:',
    history[0],
    history[-1],
    '; rewards:',
    np.asarray(rewards),
)

Run the example

python3 phases/18-posttraining/03-reward/code/main.py

About this command

This shortcut uses an activated macOS/Linux course environment in the repository folder. For Windows, use the explicit interpreter command in “Setup help” above. The script includes worked examples and reference solutions.

Inspect the complete assembled example

main.py · complete Python example

# Step 1 — 1. Define chosen-minus-rejected likelihood: The loss sees feature differences.
import jax
import jax.numpy as jnp
import numpy as np

def preference_loss(w, chosen, rejected):
    # Perform matrix / vector contraction (`@`) to compute `margin`.
    margin = (chosen - rejected) @ w
    return jnp.mean(jax.nn.softplus(-margin))

# Step 2 — 2. Construct and inspect the three comparisons: Write the desired order before fitting: response zero above one,...
features = jnp.array([[1.0, 0.0], [0.0, 1.0], [-1.0, -1.0]])
chosen = features[jnp.array([0, 0, 1])]
rejected = features[jnp.array([1, 2, 2])]
w = jnp.zeros(2)
history = []
# Differentiate the objective to obtain `step` via automatic differentiation.
step = jax.jit(
    jax.value_and_grad(lambda w: preference_loss(w, chosen, rejected))
)

# Step 3 — 3. Fit rewards and check their invariances: The host logaddexp oracle checks stable likelihood values.
for _ in range(100):
    value, g = step(w)
    history.append(float(value))
    w = w - 0.1 * g

# Perform matrix / vector contraction (`@`) to compute `rewards`.
rewards = features @ w
assert rewards[0] > rewards[1] > rewards[2]
np.testing.assert_allclose(history[0], np.log(2), atol=1e-6)
assert history[-1] < 0.1
# Convert `margins` to a host NumPy array for inspection or verification.
margins = np.asarray((chosen - rejected) @ w, dtype=np.float64)
# Execute `np.testing.assert_allclose(`.
np.testing.assert_allclose(
    preference_loss(w, chosen, rejected),
    np.mean(np.logaddexp(0, -margins)),
    atol=1e-6,
)
# Only differences are identified: a common score offset leaves preference probabilities unchanged.
np.testing.assert_allclose(
    jax.nn.sigmoid((chosen @ w + 9) - (rejected @ w + 9)),
    jax.nn.sigmoid((chosen - rejected) @ w),
    atol=1e-6,
)
print(
    'Synthetic preference NLL initial/final:',
    history[0],
    history[-1],
    '; rewards:',
    np.asarray(rewards),
)

Expected result: The model learns all three declared preferences; common reward offsets preserve probabilities.

Recorded reference execution

Reference output captured from executing the complete lesson script (code/main.py) and its assertions in the course CPU environment.

2026-10-20T10:52:15.398245+00:00 · Python 3.14.3 · JAX 0.9.2 · 1 CPU device

Synthetic preference NLL initial/final: 0.6931471824645996 0.08669456839561462 ; rewards: [ 1.9811294  0.327278  -2.3084073]
Synthetic preference NLL initial/final: 0.6931471824645996 0.08669456839561462 ; rewards: [ 1.9811294  0.327278  -2.3084073]
Forward/reversed preference NLL: 0.08596369624137878 2.94565486907959
Balanced conflicting labels prefer a zero margin.
Tied baseline and preference directions verified.
Ranking unchanged; confidence and downstream reward scale change.
Independent reward gradients agree in both label directions.
PASS: posttraining-03

Download execution record · The notebook includes executed cells and figure outputs. Compare your run with these values; timings depend on the machine.

Reverse the comparison

Reverse the comparison · Python experiment

# Experiment — Reverse the comparison: This checks label direction.
forward = float(preference_loss(w, chosen, rejected))
reverse = float(preference_loss(w, rejected, chosen))
assert reverse > forward
print('Forward/reversed preference NLL:', forward, reverse)

Expected: The reversed-label loss is larger.

This checks label direction. A decreasing incorrectly signed loss can train a consistent but reversed preference model.

Fit conflicting labels conceptually

Fit conflicting labels conceptually · Python experiment

# Experiment — Fit conflicting labels conceptually: Disagreement is not fixed by flipping labels until training...
conflict = lambda margin: 0.5 * (
    jax.nn.softplus(-margin) + jax.nn.softplus(margin)
)
np.testing.assert_allclose(conflict(0.0), np.log(2), atol=1e-6)
# Differentiate the objective to obtain gradients ``.
np.testing.assert_allclose(jax.grad(conflict)(0.0), 0.0, atol=1e-7)
assert float(conflict(2.0)) > float(conflict(0.0))
# Execute `np.testing.assert_allclose(conflict(2.0), conflict(-2.0))`
np.testing.assert_allclose(conflict(2.0), conflict(-2.0))
print('Balanced conflicting labels prefer a zero margin.')

Expected: Zero margin minimizes the balanced contradictory-pair objective.

Disagreement is not fixed by flipping labels until training accuracy is high. The likelihood must represent the information actually present.

Hands-On Exercise · Modify & Verify Run in the declared lesson environment

Make it yours

Where to put your code: Append your solution at the bottom of phases/18-posttraining/03-reward/code/main.py (or inside the Make it yours cell of exercises/posttraining-03.ipynb) and run python3 phases/18-posttraining/03-reward/code/main.py to verify your prediction before opening the reference solution.

Calculate the initial loss for tied reward scores independently and verify the trained ordering.

Show a writing guide and starter scaffold

How to write this exercise code: step-by-step recipe & starter scaffold

1. Key functions & syntax to use:

jnp.zeros / jnp.ones(shape, dtype=...) — Allocates a tensor of the given shape initialized with constants.
jnp.allclose(actual, expected, rtol=..., atol=...) — Checks that two arrays match elementwise within floating-point tolerance.

2. Step-by-step implementation plan:

Set up the changed input or configuration parameters requested in the prompt.
Run the target function or transformation and assert that the output matches your prediction.

Starter code scaffold · copy into your file & fill in the TODOs

# Exercise solution: Calculate the initial loss for tied reward scores independently and...
np.testing.assert_allclose(
    preference_loss(jnp.zeros(2), chosen, rejected), -np.log(0.5), atol=1e-6
)
assert np.all(np.asarray((chosen - rejected) @ w)  # TODO: complete assertion check
print('Tied baseline and preference directions verified.')

Reveal the reference solution

Reference solution

# Exercise solution: Calculate the initial loss for tied reward scores independently and...
np.testing.assert_allclose(
    preference_loss(jnp.zeros(2), chosen, rejected), -np.log(0.5), atol=1e-6
)
assert np.all(np.asarray((chosen - rejected) @ w) > 0)
print('Tied baseline and preference directions verified.')

Show that feature scale can change scores without changing ranking

Transfer

Multiply the trained reward weights by two and compare ordering and preference probabilities.

Show a writing guide and starter scaffold

How to write this practice solution: step-by-step recipe & starter scaffold

1. Key functions & syntax to use:

ranking(...) — Call ranking with your updated parameters or inputs from this lesson's workspace.
contraction(...) — Call contraction with your updated parameters or inputs from this lesson's workspace.
assert condition — Verify that the observed output shape, status, or numerical value satisfies the contract.

2. Step-by-step implementation plan:

Perform matrix / vector contraction (@) to compute scaled_rewards.
Execute np.testing.assert_array_equal(np.argsort(rewards), np.argsor

Starter code scaffold · copy into your file & fill in the TODOs

# Show that feature scale can change scores without changing ranking (Transfer): A policy’s reward-versus-KL tradeoff depends on reward scale.
# Perform matrix / vector contraction (`@`) to compute `scaled_rewards`.
scaled_rewards = ...  # TODO: compute scaled_rewards
# Execute `np.testing.assert_array_equal(np.argsort(rewards), np.argsor`
np.testing.assert_array_equal(np.argsort(rewards), np.argsort(scaled_rewards))
assert float(preference_loss(2 * w, chosen, rejected))  # TODO: complete assertion check
    preference_loss(w, chosen, rejected)
)
print('Ranking unchanged; confidence and downstream reward scale change.')

Show a hint

Positive scaling preserves ordering but sharpens probabilities.

Reveal reference solution and reasoning

Show that feature scale can change scores without changing ranking · Python solution

# Show that feature scale can change scores without changing ranking (Transfer): A policy’s reward-versus-KL tradeoff depends on reward scale.
# Perform matrix / vector contraction (`@`) to compute `scaled_rewards`.
scaled_rewards = features @ (2 * w)
# Execute `np.testing.assert_array_equal(np.argsort(rewards), np.argsor`
np.testing.assert_array_equal(np.argsort(rewards), np.argsort(scaled_rewards))
assert float(preference_loss(2 * w, chosen, rejected)) < float(
    preference_loss(w, chosen, rejected)
)
print('Ranking unchanged; confidence and downstream reward scale change.')

A policy’s reward-versus-KL tradeoff depends on reward scale. Rank agreement alone is distinct from an equivalent RL objective.

Derive the linear reward gradient

Challenge

At a fresh weight vector, calculate the mean feature difference weighted by probability-minus-one and compare with autodiff. Repeat after reversing labels.

Show a writing guide and starter scaffold

How to write this practice solution: step-by-step recipe & starter scaffold

1. Key functions & syntax to use:

jnp.array(values, dtype=...) — Constructs an immutable device-backed JAX array from Python/NumPy values.
x.sum(axis=...) / x.mean(axis=..., keepdims=...) — Reduces values along the named axis (the axis you name is collapsed unless keepdims=True).
jnp.allclose(actual, expected, rtol=..., atol=...) — Checks that two arrays match elementwise within floating-point tolerance.
jax.grad(loss_fn)(params, ...) — Transforms a scalar-output function into a function returning the gradient PyTree with the same structure as params.

2. Step-by-step implementation plan:

Convert difference to a host NumPy array for inspection or verification.
Perform matrix / vector contraction (@) to compute margin_host.
Reduce along axis=0 to compute analytic.
Differentiate the objective to obtain gradients .

Starter code scaffold · copy into your file & fill in the TODOs

# Derive the linear reward gradient (Challenge): Checking both directions catches a sign error that a single...
probe_w = jnp.array(...)  # TODO: compute probe_w
for positive, negative in [(chosen, rejected), (rejected, chosen)]:
    # Convert `difference` to a host NumPy array for inspection or verification.
    difference = np.asarray(...)  # TODO: compute difference
    # Perform matrix / vector contraction (`@`) to compute `margin_host`.
    margin_host = ...  # TODO: compute margin_host
    # Reduce along axis=0 to compute `analytic`.
    analytic = np.mean(...)  # TODO: compute analytic
        (1 / (1 + np.exp(-margin_host)) - 1)[:, None] * difference, axis=0
    )
    # Differentiate the objective to obtain gradients ``.
    np.testing.assert_allclose(
        jax.grad(preference_loss)(probe_w, positive, negative),
        analytic,
        atol = ...  # TODO: compute atol
    )
print('Independent reward gradients agree in both label directions.')

Show a hint

Use NumPy sigmoid values for the independent calculation; reversing both feature order and probabilities is necessary.

Reveal reference solution and reasoning

Derive the linear reward gradient · Python solution

# Derive the linear reward gradient (Challenge): Checking both directions catches a sign error that a single...
probe_w = jnp.array([0.2, -0.4])
for positive, negative in [(chosen, rejected), (rejected, chosen)]:
    # Convert `difference` to a host NumPy array for inspection or verification.
    difference = np.asarray(positive - negative, dtype=np.float64)
    # Perform matrix / vector contraction (`@`) to compute `margin_host`.
    margin_host = difference @ np.asarray(probe_w)
    # Reduce along axis=0 to compute `analytic`.
    analytic = np.mean(
        (1 / (1 + np.exp(-margin_host)) - 1)[:, None] * difference, axis=0
    )
    # Differentiate the objective to obtain gradients ``.
    np.testing.assert_allclose(
        jax.grad(preference_loss)(probe_w, positive, negative),
        analytic,
        atol=1e-6,
    )
print('Independent reward gradients agree in both label directions.')

Checking both directions catches a sign error that a single decreasing training curve might conceal.

Diagnose the result

If margins have the wrong sign, inspect chosen/rejected order. If training ranking is perfect but held-out agreement fails, inspect prompt leakage, rubric consistency and reward-model overfitting.

Carry forward

Pairwise preference loss depends only on score margins $r_\theta(x,y^+)-r_\theta(x,y^-)$ , starting at $\log 2$ for tied scores and weighting feature differences by $\sigma(m)-1$ .
Common additive reward offsets cancel in pairwise comparisons while positive weight rescaling sharpens confidence and alters the downstream reward-versus-KL balance.

Keep your evidence

Download evidence template

Keep synthetic label provenance, stable preference-loss and gradient references, offset/scale tests and the learned ordering.

Save your modified code, its output, and an explanation of your prediction. Checkpoint answers and exercise evidence are tracked separately.

Saved in this browser. Export or back up your portfolio anytime in My learning.

Save work for this lesson

Verified on local CPU. Stage to a Cloud TPU VM via Cloud TPU setup to record accelerator execution receipts (TODO: add automated Cloud TPU CI execution receipts).

Primary references

InstructGPT: reward modeling from preferences ↗

Learn a reward model from pairwise preferences

What you will be able to do

The problem

The idea

A preference compares scores rather than calibrating utility

Before coding: a preference is a comparison

Represent a preference as a comparison

Fit score differences with a stable likelihood

Calculate likelihood and its derivative

Disagreement changes the optimum

Separate ranking from calibration

Preserve disagreement and data provenance

Keep the reward model fixed during policy evaluation

Audit both what scores identify and what they do not

Learn a reward model from pairwise preferences — recorded experiment

Read the figure

Connect it to the computation

1 Define chosen-minus-rejected likelihood

2 Construct and inspect the three comparisons

3 Fit rewards and check their invariances

Run the example

Recorded reference execution

Reverse the comparison

Fit conflicting labels conceptually

Make it yours

Show that feature scale can change scores without changing ranking

Derive the linear reward gradient

Check your understanding

Diagnose the result

Carry forward

Keep your evidence

Primary references