Module 6. The LLM as a policy: tokens, log-probabilities, and KL
Read the explanations and code in the existing notebooks, or download the files. Code does not run on this site.
Download course files (.zip) 382.1 KBThe archive preserves the source folder structure, both languages, and supporting files.
Module 6. Token probabilities, KL, and a reproducible rollout contract
Main module notebook
Download .ipynb module-06-logprobs-kl.en.ipynb · 19.9 KB
Show notebook text
Module 6. Token probabilities, KL, and a reproducible rollout contract
This notebook assembles the module's mathematical details into one reproducible workflow. The default path needs neither network access nor a GPU: it checks causal shifting, response masks, exact KL, k1/k2/k3 behavior, the actual decoding policy, and stopping semantics. The final optional branch uses a real open model; it is disabled by default and is not presented as an executed course result.
Python
from __future__ import annotations
from pathlib import Path
import os
import sys
import math
import numpy as np
import matplotlib.pyplot as plt
def find_repo_root(start: Path) -> Path:
for candidate in (start, *start.parents):
if (candidate / "modules/module-06/module_6_reference.py").is_file():
return candidate
raise FileNotFoundError("Could not find the RL for LLM repository root")
REPO_ROOT = find_repo_root(Path.cwd().resolve())
MODULE_DIR = REPO_ROOT / "modules/module-06"
sys.path.insert(0, str(MODULE_DIR))
from module_6_reference import (
behavior_logprob,
causal_response_logprobs,
classify_stop_event,
decoding_distribution_with_order,
exact_kl_from_logits,
kl_estimator_moments,
masked_sequence_metrics,
sequence_logprob_ratio,
stable_log_softmax,
validate_rollout_batch,
)
PROFILE = os.getenv("RL_LLM_M6_PROFILE", "smoke").strip().lower()
if PROFILE not in {"smoke", "full", "real-model"}:
raise ValueError("RL_LLM_M6_PROFILE must be smoke, full, or real-model")
RUN_REAL_MODEL = os.getenv("RL_LLM_M6_RUN_REAL_MODEL", "0") == "1" or PROFILE == "real-model"
SEED = 20260805
RNG = np.random.default_rng(SEED)
N_MONTE_CARLO = 20_000 if PROFILE == "smoke" else 250_000
print({"profile": PROFILE, "run_real_model": RUN_REAL_MODEL, "repo": str(REPO_ROOT)})
1. Stable token log-probabilities
Directly evaluating exp(logits) overflows for large values. Subtracting the maximum leaves softmax unchanged while making the computation stable. We then verify normalization and gather the log-probabilities of selected tokens.
Python
large_logits = np.array([[1000.0, 999.0, 997.0], [-1000.0, -1001.0, -1002.0]])
log_probs = stable_log_softmax(large_logits, axis=-1)
probs = np.exp(log_probs)
np.testing.assert_allclose(probs.sum(axis=-1), 1.0, atol=1e-12)
assert np.all(np.isfinite(log_probs))
selected_ids = np.array([0, 2])
selected = log_probs[np.arange(2), selected_ids]
print("selected log-probabilities:", selected)
print("row sums:", probs.sum(axis=-1))
2. Causal shift, response mask, and PPL
Logits at position predict token . The response mask must shift with the target. We construct a small batch in which prompt tokens are excluded from the objective and response tokens are included.
Python
B, L, V = 2, 7, 9
input_ids = np.array([
[0, 1, 2, 4, 5, 6, 8],
[0, 3, 2, 7, 1, 8, 8],
], dtype=np.int64)
response_mask = np.array([
[0, 0, 0, 1, 1, 1, 1],
[0, 0, 0, 1, 1, 1, 0],
], dtype=bool)
model_logits = RNG.normal(0.0, 1.2, size=(B, L, V))
# Make the realized response tokens plausible but not deterministic.
for b in range(B):
for t in range(L - 1):
model_logits[b, t, input_ids[b, t + 1]] += 1.5
selected_lp, shifted_mask = causal_response_logprobs(model_logits, input_ids, response_mask)
metrics = masked_sequence_metrics(selected_lp, shifted_mask)
print("selected shape:", selected_lp.shape)
print("shifted masks:\n", shifted_mask.astype(int))
print("token counts:", metrics.token_count)
print("sum log-probabilities:", np.round(metrics.sum_logprob, 4))
print("perplexities:", np.round(metrics.perplexity, 4))
fig, ax = plt.subplots(figsize=(9.0, 3.6))
for b in range(B):
ax.plot(np.arange(L - 1), selected_lp[b], marker="o", label=f"sequence {b}")
inactive = ~shifted_mask[b]
ax.scatter(np.arange(L - 1)[inactive], selected_lp[b][inactive], marker="x", s=65)
ax.set_xlabel("shifted token position")
ax.set_ylabel("selected log-probability")
ax.set_title("Causal alignment: circles are scored positions; crosses are masked")
ax.grid(True)
ax.legend()
plt.show()
3. Exact KL and sample formulas
On a finite space we can compute both means and exact second moments. We compare nearby distributions with a case containing a rare large ratio. This makes clear why non-negative k3 samples do not guarantee low variance.
Python
near_q = np.array([0.40, 0.30, 0.20, 0.10])
near_p = np.array([0.35, 0.30, 0.22, 0.13])
tail_q = np.array([0.999, 0.001])
tail_p = np.array([0.9, 0.1])
near = kl_estimator_moments(near_p, near_q)
tail = kl_estimator_moments(tail_p, tail_q)
print("near-policy moments:", {k: round(v, 6) for k, v in near.items()})
print("heavy-ratio-tail moments:", {k: round(v, 6) for k, v in tail.items()})
# Cross-check the near-policy exact value through logits.
near_exact_from_logits = float(exact_kl_from_logits(np.log(near_q), np.log(near_p)))
assert abs(near_exact_from_logits - near["exact_kl"]) < 1e-12
labels = ["k1", "k2", "k3"]
x = np.arange(3)
width = 0.34
fig, ax = plt.subplots(figsize=(8.5, 4.6))
ax.bar(x - width / 2, [near[f"{k}_std"] for k in labels], width, label="near policies")
ax.bar(x + width / 2, [tail[f"{k}_std"] for k in labels], width, label="rare large p/q")
ax.set_yscale("log")
ax.set_xticks(x, labels)
ax.set_ylabel("exact standard deviation")
ax.set_title("Estimator variance depends on the ratio tail")
ax.grid(True, axis="y", which="both")
ax.legend()
plt.show()
4. Monte Carlo: what a finite log reveals
Exact moments provide a reference, but a real log is finite. We sample from and observe how the means of k1, k2, and k3 approach their expectations. k2 is treated as a local formula, not as a universally unbiased KL estimator.
Python
def estimator_samples(p: np.ndarray, q: np.ndarray, n: int, rng: np.random.Generator):
ids = rng.choice(len(q), size=n, p=q)
log_r = np.log(p[ids]) - np.log(q[ids])
return {
"k1": -log_r,
"k2": 0.5 * log_r**2,
"k3": (p[ids] / q[ids] - 1.0) - log_r,
}
samples = estimator_samples(tail_p, tail_q, N_MONTE_CARLO, RNG)
checkpoints = np.unique(np.geomspace(20, N_MONTE_CARLO, 90).astype(int))
fig, ax = plt.subplots(figsize=(9.0, 4.6))
for name, values in samples.items():
running = np.cumsum(values)[checkpoints - 1] / checkpoints
ax.plot(checkpoints, running, label=name)
ax.axhline(tail["exact_kl"], linestyle="--", linewidth=1.4, label="exact KL")
ax.set_xscale("log")
ax.set_xlabel("number of behavior samples")
ax.set_ylabel("running sample mean")
ax.set_title("Finite-sample behavior under a rare ratio tail")
ax.grid(True)
ax.legend()
plt.show()
print("final means:", {name: float(np.mean(values)) for name, values in samples.items()})
5. The behavior policy is the full decoding pipeline
Identical model weights do not imply identical action distributions when temperature, truncation, or processor order changes. In a concrete example, reordering top-k and top-p changes the support itself.
Python
order_logits = np.array([1.720, 3.523, 1.987])
order_a = decoding_distribution_with_order(
order_logits, temperature=1.7, top_k=2, top_p=0.7,
order=("temperature", "top_k", "top_p"),
)
order_b = decoding_distribution_with_order(
order_logits, temperature=1.7, top_k=2, top_p=0.7,
order=("temperature", "top_p", "top_k"),
)
print("temperature → top-k → top-p:", np.round(order_a, 6))
print("temperature → top-p → top-k:", np.round(order_b, 6))
assert np.count_nonzero(order_a) == 1 and np.count_nonzero(order_b) == 2
chosen = 2
print("chosen token logprob under A:", behavior_logprob(order_a, chosen))
print("chosen token logprob under B:", behavior_logprob(order_b, chosen))
ratio = sequence_logprob_ratio(
np.log(np.array([0.55, 0.42, 0.61])),
np.log(np.array([0.50, 0.45, 0.30])),
)
print("sequence probability ratio:", ratio)
fig, ax = plt.subplots(figsize=(7.8, 4.0))
positions = np.arange(len(order_logits))
ax.bar(positions - 0.18, order_a, 0.36, label="T→k→p")
ax.bar(positions + 0.18, order_b, 0.36, label="T→p→k")
ax.set_xticks(positions, [f"token {i}" for i in positions])
ax.set_ylabel("behavior probability")
ax.set_title("Processor order can change support")
ax.legend(); ax.grid(True, axis="y")
plt.show()
6. Stop reason and rollout-batch validation
stop_reason records the engine event. Boolean flags encode task-specific semantics. An unknown reason must not be silently converted into a terminal event, while a true terminal event may legitimately coincide with an external boundary.
Python
stop_examples = [
classify_stop_event("eos", terminal_reasons=("eos",), truncation_reasons=("length", "timeout")),
classify_stop_event("length", terminal_reasons=("eos",), truncation_reasons=("length", "timeout")),
classify_stop_event("success_and_timeout", terminal_reasons=("success_and_timeout",), truncation_reasons=("success_and_timeout",)),
classify_stop_event("server_cancelled", terminal_reasons=("eos",), truncation_reasons=("length", "timeout")),
]
for item in stop_examples:
print(item)
batch = {
"response_ids": np.array([[5, 6, 2], [7, 8, 0], [4, 2, 0]]),
"response_mask": np.array([[1, 1, 1], [1, 1, 0], [1, 1, 0]]),
"old_logprobs": np.array([[-0.2, -0.4, -0.1], [-0.3, -0.5, np.nan], [-0.6, -0.2, np.nan]]),
"ref_logprobs": np.array([[-0.25, -0.45, -0.12], [-0.35, -0.55, np.nan], [-0.5, -0.3, np.nan]]),
"sequence_rewards": np.array([1.0, 0.2, 1.0]),
"terminated": np.array([True, False, True]),
"truncated": np.array([False, True, True]),
"stop_reason": np.array(["eos", "length", "success_and_timeout"], dtype=object),
}
summary = validate_rollout_batch(batch)
print(summary)
assert summary.batch_size == 3 and summary.active_tokens == 7
assert summary.terminated_sequences == 2 and summary.truncated_sequences == 2
7. Optional real-model branch
The branch below uses Qwen/Qwen3-0.6B with Transformers. It runs only when RL_LLM_M6_RUN_REAL_MODEL=1 or the real-model profile is selected. The code compares transition scores from generate() with teacher forcing. With a direct Transformers call, the stop reason is partly inferred from tokens and length; that is not equivalent to an explicit reason field from a production engine.
Python
def teacher_forced_generated_logprobs(model, sequences, prompt_length: int):
import torch
with torch.no_grad():
logits = model(sequences).logits[:, :-1, :]
log_probs = torch.log_softmax(logits.float(), dim=-1)
targets = sequences[:, 1:]
selected = log_probs.gather(-1, targets.unsqueeze(-1)).squeeze(-1)
# generated token at absolute position prompt_length is predicted by logits prompt_length-1
return selected[:, prompt_length - 1:]
def inferred_transformers_stop_reason(generated_ids, eos_ids, max_new_tokens: int):
eos_set = {int(x) for x in eos_ids}
if generated_ids and int(generated_ids[-1]) in eos_set:
return {"stop_reason": "eos_inferred", "terminated": True, "truncated": False, "inferred": True}
if len(generated_ids) >= max_new_tokens:
return {"stop_reason": "length_inferred", "terminated": False, "truncated": True, "inferred": True}
return {"stop_reason": "unknown_inferred", "terminated": False, "truncated": False, "inferred": True}
print("Real-model helpers are defined; no model has been downloaded or executed yet.")
Running the real model
When the branch is disabled, the cell exits without importing Transformers. The complete default notebook can therefore be executed and validated offline.
Python
REAL_MODEL_RESULT = None
if RUN_REAL_MODEL:
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
model_id = "Qwen/Qwen3-0.6B"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
model_id,
torch_dtype="auto",
device_map="auto",
)
prompt = "Explain in one sentence why a sampled token log-ratio may be negative while KL is non-negative."
encoded = tokenizer(prompt, return_tensors="pt").to(model.device)
max_new_tokens = 32
with torch.no_grad():
generated = model.generate(
**encoded,
max_new_tokens=max_new_tokens,
do_sample=True,
temperature=0.8,
top_p=0.9,
return_dict_in_generate=True,
output_scores=True,
)
transition = model.compute_transition_scores(
generated.sequences,
generated.scores,
getattr(generated, "beam_indices", None),
normalize_logits=True,
)
prompt_length = encoded["input_ids"].shape[1]
teacher = teacher_forced_generated_logprobs(model, generated.sequences, prompt_length)
generated_count = transition.shape[1]
teacher = teacher[:, :generated_count]
max_delta = float(torch.max(torch.abs(transition.float() - teacher.float())).cpu())
new_ids = generated.sequences[0, prompt_length:].tolist()
eos_ids = tokenizer.eos_token_id
eos_ids = [eos_ids] if isinstance(eos_ids, int) else list(eos_ids or [])
stop = inferred_transformers_stop_reason(new_ids, eos_ids, max_new_tokens)
REAL_MODEL_RESULT = {
"model_id": model_id,
"generated_tokens": len(new_ids),
"max_transition_teacher_delta": max_delta,
**stop,
"text": tokenizer.decode(generated.sequences[0], skip_special_tokens=True),
}
print(REAL_MODEL_RESULT)
else:
print("Optional real-model branch skipped. Set RL_LLM_M6_RUN_REAL_MODEL=1 to run it explicitly.")
8. Checks and conclusions
The default path must pass independently of the optional branch. The final statements describe synthetic computations in this notebook and are not presented as an external benchmark.
Python
assert np.allclose(np.exp(log_probs).sum(axis=-1), 1.0)
assert np.isclose(near["k1_mean"], near["exact_kl"], atol=1e-14)
assert np.isclose(near["k3_mean"], near["exact_kl"], atol=1e-14)
assert tail["k3_std"] > 10 * tail["k1_std"]
assert np.count_nonzero(order_a) != np.count_nonzero(order_b)
assert stop_examples[-1].terminated is False and stop_examples[-1].truncated is False
COURSE_RESULT = {
"synthetic_only": not RUN_REAL_MODEL,
"causal_batch_perplexity": np.asarray(metrics.perplexity).round(6).tolist(),
"near_exact_kl": near["exact_kl"],
"heavy_tail_exact_kl": tail["exact_kl"],
"heavy_tail_k3_std": tail["k3_std"],
"rollout_active_tokens": summary.active_tokens,
"real_model_result": REAL_MODEL_RESULT,
}
print(COURSE_RESULT)