CS336, лекция 17 и итоги курса: policy gradient и GRPO в коде

Как устроен этот перевод

Лекция — исполняемая программа lecture_17.py. Текст переведён по порядку, код приведён с русскими комментариями, значения переменных — из трассы выполнения var/traces/lecture_17.json, итоги экспериментов — из логов var/policy_gradient_*.txt в репозитории. Графики обучения не копируем, а описываем числами из логов. Пояснения, которых нет в оригинале, помечены «Примечание переводчика». Видео — в плейлисте 2025 года, лекция 17.

Как читать вместе с курсом

• Это «механика» к теории лекций 15–16. Нужны понятия MDP, политики и преимущества из модуля 6.0 и урок 6.3.1 «Основы RL для языковых моделей».
• Игрушечная модель здесь работает на CPU за секунды — весь код можно запустить у себя. Затем переходите к практикуму 6.4.5 «GRPO своими руками» и к заданию 5 CS336 (путеводитель): функции compute_deltas и compute_loss отсюда — прямые прототипы того, что там проверяют тесты.
• В конце страницы — итоги всего курса CS336 и что было после этой лекции.

---

Лекция 17. Выравнивание — RL (policy gradient)

Прошлая лекция: обзор RL с проверяемыми наградами (policy gradient). Эта лекция: подробно о механике policy gradient (например, GRPO).

Постановка RL для языковых моделей

• Состояние $s$: промпт + уже сгенерированная часть ответа.
• Действие $a$: сгенерировать следующий токен.
• Награда $R$: насколько хорош ответ; мы сосредоточимся на:
  • наградах за результат (outcome rewards), которые зависят от всего ответа;
  • проверяемых наградах (verifiable rewards), вычисление которых детерминировано;
  • понятия дисконтирования и бутстрэппинга здесь менее применимы.
  • Пример: «… Therefore, the answer is 3 miles.» («…Следовательно, ответ — 3 мили».)
• Вероятности переходов $T(s' \mid s, a)$: детерминированы, $s' = s + a$:
  • можно планировать и тратить вычисления во время ответа (в отличие от робототехники);
  • состояния по сути «придуманы» (в отличие от робототехники), поэтому много гибкости.
• Политика $\pi(a \mid s)$: просто языковая модель (дообученная).
• Роллаут / эпизод / траектория: $s \to a \to \dots \to a \to a \to R$.
• Цель: максимизировать ожидаемую награду $\mathbb{E}[R]$ (ожидание берётся по промптам $s$ и токенам ответа $a$).

Policy gradient

Для простоты записи пусть $a$ обозначает весь ответ целиком. Мы хотим максимизировать ожидаемую награду по политике $\pi$:

$$\mathbb{E}[R] = \int p(s)\, \pi(a \mid s)\, R(s, a).$$

Очевидный шаг — взять градиент:

$$\nabla \mathbb{E}[R] = \int p(s)\, \nabla \pi(a \mid s)\, R(s, a) = \int p(s)\, \pi(a \mid s)\, \nabla \log \pi(a \mid s)\, R(s, a) = \mathbb{E}\left[\nabla \log \pi(a \mid s)\, R(s, a)\right].$$

Примечание переводчика: откуда логарифм

Это «трюк с логарифмической производной»: $\nabla \pi = \pi \cdot \nabla \log \pi$, потому что $\nabla \log \pi = \frac{\nabla \pi}{\pi}$. Он превращает интеграл с $\nabla \pi$ (который нельзя оценить сэмплированием) в ожидание по $\pi$ (которое можно: сгенерировали ответ — посчитали $\nabla \log \pi \cdot R$). Подробно — урок 6.3.1.

Наивный policy gradient:

• сэмплируем промпт $s$, сэмплируем ответ $a \sim \pi(a \mid s)$;
• обновляем параметры по $\nabla \log \pi(a \mid s)\, R(s, a)$ (то же, что SFT, но с весом $R(s, a)$).

Постановка: $R(s, a) \in \{0, 1\}$ — правилен ли ответ.

• Наивный policy gradient обновляет только правильные ответы.
• Как SFT, но датасет меняется со временем вместе с политикой.

Трудность: высокий шум/дисперсия. В этой постановке награды разреженные (немногие ответы получают 1, большинство — 0). Для контраста: в RLHF модели награды (обученные на попарных предпочтениях) более непрерывны.

Базовые линии

Напомним: $\nabla \mathbb{E}[R] = \mathbb{E}[\nabla \log \pi(a \mid s)\, R(s, a)]$. Величина $\nabla \log \pi(a \mid s)\, R(s, a)$ — несмещённая оценка $\nabla \mathbb{E}[R]$, но, может быть, есть другие, с меньшей дисперсией…

Пример: два состояния.

• $s_1$: $a_1$ → награда 11, $a_2$ → награда 9;
• $s_2$: $a_1$ → награда 0, $a_2$ → награда 2.

Мы не хотим $s_1 \to a_2$ (награда 9), потому что $a_1$ лучше, и хотим $s_2 \to a_2$ (награда 2), но 9 > 2.

Идея: максимизировать награду с базовой линией: $\mathbb{E}[R - b(s)]$. Это просто $\mathbb{E}[R]$, сдвинутое на константу $\mathbb{E}[b(s)]$, которая не зависит от политики $\pi$. Обновляемся по $\nabla \log \pi(a \mid s)\,(R(s, a) - b(s))$.

Какую $b(s)$ взять? Пример с двумя состояниями: считаем распределение по $(s, a)$ равномерным и $|\nabla \pi(a \mid s)| = 1$.

naive_variance = torch.std(torch.tensor([11., 9, 0, 2]))  # → 5.323
# Определим базовую линию b(s1) = 10, b(s2) = 1
baseline_variance = torch.std(torch.tensor([11. - 10, 9 - 10, 0 - 1, 2 - 1]))  # → 1.155

Дисперсия (здесь — стандартное отклонение) уменьшилась с 5,323 до 1,155.

Оптимальная $b^*(s) = \frac{\mathbb{E}[(\nabla \pi(a \mid s))^2 R \mid s]}{\mathbb{E}[(\nabla \pi(a \mid s))^2 \mid s]}$ (для однопараметрических моделей). Её трудно вычислить… поэтому эвристика — использовать среднюю награду: $b(s) = \mathbb{E}[R \mid s]$. Её тоже трудно вычислить, и её приходится оценивать.

Функции преимущества

Такой выбор $b(s)$ связан с функциями преимущества.

• $V(s) = \mathbb{E}[R \mid s]$ — ожидаемая награда из состояния $s$.
• $Q(s, a) = \mathbb{E}[R \mid s, a]$ — ожидаемая награда из состояния $s$ при действии $a$.

(Замечание: $Q$ и $R$ здесь совпадают, потому что мы считаем, что $a$ содержит все действия, а награда — за результат.)

Определение (преимущество): $A(s, a) = Q(s, a) - V(s)$. Интуиция: насколько действие $a$ лучше ожидаемого из состояния $s$.

Если $b(s) = \mathbb{E}[R \mid s]$, то награда с базовой линией совпадает с преимуществом: $\mathbb{E}[R - b(s)] = A(s, a)$.

В общем случае:

• идеал: $\mathbb{E}[\nabla \log \pi(a \mid s)\, R(s, a)]$;
• оценка: $\nabla \log \pi(a \mid s)\, \delta$.

Есть несколько вариантов $\delta$, как мы увидим дальше. Конспект CS224R о policy gradient.

Разбор обучения: GRPO

Group Relative Policy Optimization (GRPOGRPO (групповая относительная оптимизация политики)EN Group Relative Policy Optimization (GRPO)PPO без критика: на промпт генерируется группа из G ответов, преимущество — (r − среднее) / std по группе. DeepSeekMath, 2024.) (статья DeepSeekMath):

• упрощение PPO, убирающее критика (функцию ценности);
• использует групповую структуру в постановке LM (несколько ответов на промпт), которая даёт естественную базовую линию $b(s)$.

[Алгоритм GRPO из статьи: для каждого шага сэмплируем батч вопросов, для каждого — $G$ ответов старой политикой, считаем награды и групповые преимущества, делаем несколько шагов оптимизации целевой функции GRPO.]

Простая задача

Задача: отсортировать $n$ чисел. Промпт: $n$ чисел (например, [1, 0, 2]). Ответ: $n$ чисел ([0, 1, 2]). Награда должна отражать, насколько ответ близок к отсортированному.

Определим награду, равную числу позиций, где ответ совпадает с эталоном:

def sort_distance_reward(prompt: list[int], response: list[int]) -> float:
    """
    Насколько response близок к ground_truth = sorted(prompt).
    Считаем число позиций, где ответ совпадает с эталоном.
    """
    assert len(prompt) == len(response)
    ground_truth = sorted(prompt)
    return sum(1 for x, y in zip(response, ground_truth) if x == y)

sort_distance_reward([3, 1, 0, 2], [0, 1, 2, 3])  # → 4
sort_distance_reward([3, 1, 0, 2], [7, 2, 2, 5])  # → 1
sort_distance_reward([3, 1, 0, 2], [0, 3, 1, 2])  # → 1

Определим альтернативную награду, дающую больше частичных баллов:

def sort_inclusion_ordering_reward(prompt: list[int], response: list[int]) -> float:
    """
    Насколько response близок к ground_truth = sorted(prompt).
    """
    assert len(prompt) == len(response)
    # Один балл за каждый токен промпта, который встречается в ответе
    inclusion_reward = sum(1 for x in prompt if x in response)
    # Один балл за каждую соседнюю пару в ответе, стоящую по порядку
    ordering_reward = sum(1 for x, y in zip(response, response[1:]) if x <= y)
    return inclusion_reward + ordering_reward

sort_inclusion_ordering_reward([3, 1, 0, 2], [0, 1, 2, 3])  # → 7  (4 + 3)
sort_inclusion_ordering_reward([3, 1, 0, 2], [7, 2, 2, 5])  # → 3
sort_inclusion_ordering_reward([3, 1, 0, 2], [0, 3, 1, 2])  # → 6

Заметьте: вторая функция награды даёт третьему ответу больше баллов, чем первая.

Примечание переводчика

Это маленький пример формирования награды (reward shaping): обе награды максимальны на правильном ответе, но вторая «плотнее» — даёт сигнал и почти правильным ответам. С плотной наградой RL учится легче, но растёт риск взлома наградывзлом награды (reward hacking)EN reward hackingПолитика находит способ получить высокую награду, не делая того, что задумали люди.: например, ответ [0, 0, 0] для промпта [0, 1, 2] получит 1 + 2 = 3 балла, ни разу не отсортировав.

Простая модель

Определим простую модель, отображающую промпты в ответы:

• длины промпта и ответа фиксированы;
• позиционная информация — через отдельные параметры для каждой позиции;
• каждая позиция ответа декодируется независимо (не авторегрессионно).

class Model(nn.Module):
    def __init__(self, vocab_size: int, embedding_dim: int, prompt_length: int, response_length: int):
        super().__init__()
        self.embedding_dim = embedding_dim
        self.embedding = nn.Embedding(vocab_size, embedding_dim)
        # Для каждой позиции — своя матрица кодирования и своя матрица декодирования
        self.encode_weights = nn.Parameter(torch.randn(prompt_length, embedding_dim, embedding_dim) / math.sqrt(embedding_dim))
        self.decode_weights = nn.Parameter(torch.randn(response_length, embedding_dim, embedding_dim) / math.sqrt(embedding_dim))

    def forward(self, prompts: torch.Tensor) -> torch.Tensor:
        """
        Args:
            prompts: int[batch pos]
        Returns:
            logits: float[batch pos vocab]
        """
        # Эмбеддинги промптов
        embeddings = self.embedding(prompts)   # [batch pos dim]
        # Преобразуем матрицей своей позиции промпта и сворачиваем в один вектор
        encoded = einsum(embeddings, self.encode_weights, "batch pos dim1, pos dim1 dim2 -> batch dim2")
        # Превращаем в один вектор на каждую позицию ответа
        decoded = einsum(encoded, self.decode_weights, "batch dim2, pos dim2 dim1 -> batch pos dim1")
        # Переводим в логиты (вход и выход делят эмбеддинги)
        logits = einsum(decoded, self.embedding.weight, "batch pos dim1, vocab dim1 -> batch pos vocab")
        return logits

model = Model(vocab_size=3, embedding_dim=10, prompt_length=3, response_length=3)

Берём промпт $s$ и генерируем ответы $a$:

def generate_responses(prompts: torch.Tensor, model: Model, num_responses: int) -> torch.Tensor:
    """
    Args:
        prompts (int[batch pos])
    Returns:
        сгенерированные ответы: int[batch trial pos]
    Пример (batch_size = 3, prompt_length = 3, num_responses = 2, response_length = 4)
    p1 p1 p1 r1 r1 r1 r1
             r2 r2 r2 r2
    p2 p2 p2 r3 r3 r3 r3
             r4 r4 r4 r4
    p3 p3 p3 r5 r5 r5 r5
             r6 r6 r6 r6
    """
    logits = model(prompts)  # [batch pos vocab]
    batch_size = prompts.shape[0]
    # Сэмплируем num_responses ответов (независимо) для каждой пары [batch pos]
    flattened_logits = rearrange(logits, "batch pos vocab -> (batch pos) vocab")
    flattened_responses = torch.multinomial(softmax(flattened_logits, dim=-1), num_samples=num_responses, replacement=True)  # [batch pos trial]
    responses = rearrange(flattened_responses, "(batch pos) trial -> batch trial pos", batch=batch_size)
    return responses

prompts = torch.tensor([[1, 0, 2]])  # [batch pos]
torch.manual_seed(10)
responses = generate_responses(prompts=prompts, model=model, num_responses=5)
# → [[[1, 0, 2], [1, 0, 2], [1, 0, 1], [1, 0, 2], [1, 0, 2]]]

Считаем награды $R$ этих ответов:

def compute_reward(prompts: torch.Tensor, responses: torch.Tensor, reward_fn: Callable[[list[int], list[int]], float]) -> torch.Tensor:
    """
    Args:
        prompts (int[batch pos])
        responses (int[batch trial pos])
    Returns:
        rewards (float[batch trial])
    """
    batch_size, num_responses, _ = responses.shape
    rewards = torch.empty(batch_size, num_responses, dtype=torch.float32)
    for i in range(batch_size):
        for j in range(num_responses):
            rewards[i, j] = reward_fn(prompts[i, :], responses[i, j, :])
    return rewards

rewards = compute_reward(prompts=prompts, responses=responses, reward_fn=sort_inclusion_ordering_reward)
# → [[4.0, 4.0, 3.0, 4.0, 4.0]]

Считаем дельты $\delta$ по наградам $R$ (для обновлений):

def compute_deltas(rewards: torch.Tensor, mode: str) -> torch.Tensor:
    """
    Args:
        rewards (float[batch trial])
    Returns:
        deltas (float[batch trial]) — похожие на преимущество величины для обновлений
    """
    if mode == "rewards":
        return rewards
    if mode == "centered_rewards":
        # Среднее по всем ответам (trial) на каждый промпт (batch)
        mean_rewards = rewards.mean(dim=-1, keepdim=True)
        centered_rewards = rewards - mean_rewards
        return centered_rewards
    if mode == "normalized_rewards":
        mean_rewards = rewards.mean(dim=-1, keepdim=True)
        std_rewards = rewards.std(dim=-1, keepdim=True)
        centered_rewards = rewards - mean_rewards
        normalized_rewards = centered_rewards / (std_rewards + 1e-5)
        return normalized_rewards
    if mode == "max_rewards":
        # Обнуляем все награды, кроме максимальной в каждом батче
        max_rewards = rewards.max(dim=-1, keepdim=True)[0]
        max_rewards = torch.where(rewards == max_rewards, rewards, torch.zeros_like(rewards))
        return max_rewards
    raise ValueError(f"Unknown mode: {mode}")

compute_deltas(rewards=rewards, mode="rewards")             # → [[4.0, 4.0, 3.0, 4.0, 4.0]]
compute_deltas(rewards=rewards, mode="centered_rewards")    # среднее 3.8 → [[0.2, 0.2, -0.8, 0.2, 0.2]]
compute_deltas(rewards=rewards, mode="normalized_rewards")  # std 0.447 → [[0.447, 0.447, -1.789, 0.447, 0.447]]
compute_deltas(rewards=rewards, mode="max_rewards")         # → [[4.0, 4.0, 0.0, 4.0, 4.0]]

Считаем логарифмы вероятностей этих ответов:

def compute_log_probs(prompts: torch.Tensor, responses: torch.Tensor, model: Model) -> torch.Tensor:
    """
    Args:
        prompts (int[batch pos])
        responses (int[batch trial pos])
    Returns:
        log_probs (float[batch trial pos]) по модели
    """
    # Логарифмы вероятностей ответов по модели
    logits = model(prompts)  # [batch pos vocab]
    log_probs = F.log_softmax(logits, dim=-1)  # [batch pos vocab]
    # Размножаем, чтобы выровнять с ответами
    num_responses = responses.shape[1]
    log_probs = repeat(log_probs, "batch pos vocab -> batch trial pos vocab", trial=num_responses)  # [batch trial pos vocab]
    # Выбираем из log_probs элементы по индексам ответов
    log_probs = log_probs.gather(dim=-1, index=responses.unsqueeze(-1)).squeeze(-1)  # [batch trial pos]
    return log_probs

log_probs = compute_log_probs(prompts=prompts, responses=responses, model=model)
# → для ответа [1, 0, 2]: [-1.3e-05, -8.1e-06, -0.119]  (модель уже почти уверена в нём)

Считаем потери, по которым будем обновлять параметры модели:

def compute_loss(log_probs: torch.Tensor, deltas: torch.Tensor, mode: str, old_log_probs: torch.Tensor | None = None) -> torch.Tensor:
    if mode == "naive":
        return -einsum(log_probs, deltas, "batch trial pos, batch trial -> batch trial pos").mean()
    if mode == "unclipped":
        ratios = torch.exp(log_probs - old_log_probs)  # [batch trial]
        return -einsum(ratios, deltas, "batch trial pos, batch trial -> batch trial pos").mean()
    if mode == "clipped":
        epsilon = 0.01
        unclipped_ratios = torch.exp(log_probs - old_log_probs)  # [batch trial]
        unclipped = einsum(unclipped_ratios, deltas, "batch trial pos, batch trial -> batch trial pos")
        clipped_ratios = torch.clamp(unclipped_ratios, min=1 - epsilon, max=1 + epsilon)
        clipped = einsum(clipped_ratios, deltas, "batch trial pos, batch trial -> batch trial pos")
        return -torch.minimum(unclipped, clipped).mean()
    raise ValueError(f"Unknown mode: {mode}")

loss = compute_loss(log_probs=log_probs, deltas=deltas, mode="naive")  # → 0.127

Примечание переводчика

В этой точке лекции deltas — последнее вычисленное значение, то есть режим max_rewards. Знак минус стоит потому, что оптимизатор минимизирует потери, а мы хотим максимизировать $\mathbb{E}[\log \pi \cdot \delta]$.

Замораживаем параметры

Мотивация: в GRPO встречаются отношения $p(a \mid s) / p_{old}(a \mid s)$. При оптимизации важно заморозить $p_{old}$ и не дифференцировать через него.

w = torch.tensor(2., requires_grad=True)
p = torch.nn.Sigmoid()(w)
p_old = torch.nn.Sigmoid()(w)
ratio = p / p_old
ratio.backward()
grad = w.grad  # → 0.0  (отношение тождественно равно 1, градиент пропал!)

# Как правильно:
w = torch.tensor(2., requires_grad=True)
p = torch.nn.Sigmoid()(w)
with torch.no_grad():  # Важно: считаем p_old константой!
    p_old = torch.nn.Sigmoid()(w)
ratio = p / p_old
ratio.backward()
grad = w.grad  # → 0.1192

Примечание переводчика: откуда 0,1192

При замороженном $p_{old}$ градиент отношения равен $\sigma'(w) / \sigma(w) = 1 - \sigma(w)$. При $w = 2$: $\sigma(2) \approx 0{,}8808$, и $1 - 0{,}8808 = 0{,}1192$. Это частая ошибка в реализациях GRPO и PPO: без no_grad (или .detach()) обучение «идёт», но градиент отношения нулевой.

Дальше в лекции — потери с отношением вероятностей (притворимся, что old_model — старый чекпоинт):

old_model = Model(vocab_size=3, embedding_dim=10, prompt_length=3, response_length=3)  # Притворимся, что это старый чекпоинт
old_log_probs = compute_log_probs(prompts=prompts, responses=responses, model=old_model)
loss = compute_loss(log_probs=log_probs, deltas=deltas, mode="unclipped", old_log_probs=old_log_probs)  # → -639.9
loss = compute_loss(log_probs=log_probs, deltas=deltas, mode="clipped", old_log_probs=old_log_probs)    # → -3.22

Примечание переводчика

Здесь «старая» модель — это просто другая случайная модель, и отношения вероятностей огромны. Без обрезки потери −640 означают гигантский шаг; обрезка до $[1 - 0{,}01;\ 1 + 0{,}01]$ ограничивает его. Именно поэтому PPO и GRPO обрезают отношения.

Штраф KL

Иногда для регуляризации модели используют явный штраф KL. Это полезно, если вы хотите «влить» новую способность в модель через RL, но не хотите, чтобы она забыла исходные способности.

$$\mathrm{KL}(p \,\|\, q) = \mathbb{E}_{x \sim p}\left[\log\frac{p(x)}{q(x)}\right] = \mathbb{E}_{x \sim p}\left[-\log\frac{q(x)}{p(x)}\right] = \mathbb{E}_{x \sim p}\left[\frac{q(x)}{p(x)} - \log\frac{q(x)}{p(x)} - 1\right],$$

потому что $\mathbb{E}_{x \sim p}\left[\frac{q(x)}{p(x)}\right] = 1$.

def compute_kl_penalty(log_probs: torch.Tensor, ref_log_probs: torch.Tensor) -> torch.Tensor:
    """
    Оценка KL(model | ref_model), где модели заданы:
        log_probs [batch trial pos vocab]
        ref_log_probs [batch trial pos vocab]
    Используем оценку:
        KL(p || q) = E_p[q/p - log(q/p) - 1]
    """
    return (torch.exp(ref_log_probs - log_probs) - (ref_log_probs - log_probs) - 1).sum(dim=-1).mean()

kl_penalty = compute_kl_penalty(log_probs=log_probs, ref_log_probs=old_log_probs)  # → 8.25

Примечание переводчика: почему именно так

Каждое слагаемое $\frac{q}{p} - \log\frac{q}{p} - 1$ неотрицательно (функция $t - \log t - 1 \ge 0$ с минимумом 0 при $t = 1$), поэтому оценка никогда не бывает отрицательной даже по одному сэмплу, в отличие от «наивной» $\log\frac{p}{q}$. Это та же оценка «k3», что в формуле GRPO из лекции 16 (слайд 30).

Итоги:

• генерируем ответы;
• считаем награды $R$ и $\delta$ (награды, центрированные, нормированные, максимальные);
• считаем логарифмы вероятностей ответов;
• считаем потери по логарифмам вероятностей и $\delta$ (наивные, без обрезки, с обрезкой).

Алгоритм целиком

Теперь определим алгоритм GRPO.

def run_policy_gradient(num_epochs: int = 100,
                        num_steps_per_epoch: int = 10,
                        compute_ref_model_period: int = 10,
                        num_responses: int = 10,
                        deltas_mode: str = "rewards",
                        loss_mode: str = "naive",
                        kl_penalty: float = 0.0,
                        reward_fn: Callable[[list[int], list[int]], float] = sort_inclusion_ordering_reward,
                        use_cache: bool = False) -> tuple[str, str]:
    """Обучаем модель методом policy gradient.
    Возвращаем:
    - путь к картинке с кривой обучения;
    - путь к файлу лога.
    """
    torch.manual_seed(5)

    image_path = f"var/policy_gradient_{deltas_mode}_{loss_mode}.png"
    log_path = f"var/policy_gradient_{deltas_mode}_{loss_mode}.txt"

    # Уже запускали — берём из кэша
    if use_cache and os.path.exists(image_path) and os.path.exists(log_path):
        return image_path, log_path

    # Данные
    prompts = torch.tensor([[1, 0, 2], [3, 2, 4], [1, 2, 3]])
    vocab_size = prompts.max() + 1
    prompt_length = response_length = prompts.shape[1]

    model = Model(vocab_size=vocab_size, embedding_dim=10, prompt_length=prompt_length, response_length=response_length)
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

    records = []
    ref_log_probs = None
    ref_model = None
    old_log_probs = None

    if use_cache:
        out = open(log_path, "w")
    else:
        out = sys.stdout

    for epoch in tqdm(range(num_epochs), desc="epoch"):
        # Если есть штраф KL, нужна эталонная модель (замораживаем её раз в несколько эпох)
        if kl_penalty != 0:
            if epoch % compute_ref_model_period == 0:
                ref_model = model.clone()

        # Сэмплируем ответы и оцениваем их награды
        responses = generate_responses(prompts=prompts, model=model, num_responses=num_responses)  # [batch trial pos]
        rewards = compute_reward(prompts=prompts, responses=responses, reward_fn=reward_fn)  # [batch trial]
        deltas = compute_deltas(rewards=rewards, mode=deltas_mode)  # [batch trial]

        if kl_penalty != 0:  # Считаем по эталонной модели
            with torch.no_grad():
                ref_log_probs = compute_log_probs(prompts=prompts, responses=responses, model=ref_model)  # [batch trial]

        if loss_mode != "naive":  # Считаем по текущей модели (но замораживаем на время внутренних шагов)
            with torch.no_grad():
                old_log_probs = compute_log_probs(prompts=prompts, responses=responses, model=model)  # [batch trial]

        # Делаем несколько шагов на этих ответах
        for step in range(num_steps_per_epoch):
            log_probs = compute_log_probs(prompts=prompts, responses=responses, model=model)  # [batch trial]
            loss = compute_loss(log_probs=log_probs, deltas=deltas, mode=loss_mode, old_log_probs=old_log_probs)
            if kl_penalty != 0:
                loss += kl_penalty * compute_kl_penalty(log_probs=log_probs, ref_log_probs=ref_log_probs)

            # Печатаем информацию
            print_information(epoch=epoch, step=step, loss=loss, prompts=prompts, rewards=rewards, responses=responses, log_probs=log_probs, deltas=deltas, out=out)
            global_step = epoch * num_steps_per_epoch + step
            records.append({"epoch": epoch, "step": global_step, "loss": loss.item(), "mean_reward": rewards.mean().item()})

            # Обратное распространение и обновление параметров
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

    if use_cache:
        out.close()

    if use_cache:
        # Рисуем шаг против потерь и награды на двух графиках
        ...  # (код matplotlib: два графика — «Train Loss» и «Mean Reward» по шагам)

    return image_path, log_path

Примечание переводчика: две мелочи в коде лекции

model.clone() у nn.Module нет — ветка с kl_penalty != 0 в таком виде упадёт (в лекции её не запускали). Используйте copy.deepcopy(model). Функции print_information и tstr печатают награды, дельты и логарифмы вероятностей каждого ответа; их код тривиален и есть в оригинале.

Эксперименты

Обучим несколько моделей (100 эпох × 10 шагов, 10 ответов на промпт, награда sort_inclusion_ordering_reward, потери naive).

Начнём с обновлений по «сырым» наградам (deltas_mode="rewards"). [График: потери падают, средняя награда растёт с 2,33 до 2,67 и застывает.] Просмотрев вывод, вы увидите, что к концу мы так толком и не научились сортировать (и это ещё на обучающей выборке).

Попробуем центрированные награды (deltas_mode="centered_rewards"). [График: средняя награда растёт с 2,33 до 3,33.] Похоже, это помогает, потому что:

• неоптимальные награды получают отрицательное обновление градиента;
• если у всех ответов на данный промпт одинаковая награда, мы не обновляемся.

В целом лучше, но мы всё ещё застреваем в локальных оптимумах.

Наконец, нормируем на стандартное отклонение (deltas_mode="normalized_rewards"). [График: средняя награда тоже доходит до 3,33.] Разница небольшая; и действительно, варианты вроде Dr. GRPO не делают этой нормировки, чтобы избежать смещения по длине (здесь это не проблема — у всех ответов одинаковая длина) (статья).

В целом, как видите, обучение с подкреплением нетривиально, и легко застрять в субоптимальных состояниях. Гиперпараметры, вероятно, можно настроить лучше…

Примечание переводчика: что видно в логах

К концу обучения с центрированными наградами модель на каждый промпт выдаёт один и тот же ответ все 10 раз: для [1, 0, 2] — [2, 0, 1] (награда 4 из 5 возможных), для [3, 2, 4] — [0, 0, 1] (награда 2), для [1, 2, 3] — [2, 3, 1] (награда 4). Все ответы в группе одинаковы, поэтому центрированные дельты равны нулю — обучение остановилось. Это то самое схлопывание мод из лекции 15: политика стала детерминированной раньше, чем нашла правильный ответ. Лечат это энтропийным бонусом, бо́льшим числом ответов в группе, температурой сэмплирования и штрафом KL.

Итоги

• Обучение с подкреплением — ключ к тому, чтобы превзойти способности человека.
• Если вы можете это измерить, вы можете это оптимизировать.
• Схема policy gradient концептуально ясна, нужны лишь базовые линии, чтобы уменьшить дисперсию.
• Системы RL гораздо сложнее претрейна (нагрузки инференса, управление несколькими моделями).

Две последние лекции:

• Цзюньян Линь (Junyang Lin, Qwen) — техотчёт Qwen3;
• Майк Льюис (Mike Lewis, Llama) — статья Llama 3.

---

Итоги курса CS336

О чём этот раздел

Это раздел переводчика, а не часть лекции 17: короткая сводка всего курса по нашим переводам. Гостевые лекции 18–19 (весна 2025) в репозиторий не выложены — есть только видео в плейлисте, поэтому их перевода нет.

За 17 лекций CS336 проходит весь путь создания языковой модели. Если свести каждую часть к одной мысли:

Часть курса	Лекции	Главная мысль	Задание	Уроки нашего курса
Основы	1–2	Всё решает эффективность: считайте байты и FLOPs	1	4.2, 4.6, 3.3
Архитектура	3–4	Современный блок почти у всех одинаков; MoE — разреженность ради дешёвых параметров	1	4.5, 5.2, 5.3
Системы	5–8	Вычисления дёшевы, перемещение данных дорого — внутри GPU и между GPU	2	5.4, 5.7
Масштабирование	9, 11	Решения для большой модели принимают по маленьким	3	5.1
Инференс и оценка	10, 12	Генерация упирается в память; единственно верной оценки нет	—	7.1–7.3
Данные	13–14	Данные не падают с неба: фильтрация и дедупликация решают	4	5.1.2
Пост-трейнинг	15–17	SFT извлекает, RL оптимизирует то, что можно измерить	5	6.1–6.4

Что изменилось в курсе весной 2026 года (параллелизм в коде, лекция 17 о мультимодальных моделях) — в обзоре курса.

Что сделать руками

1. Запустите лекцию у себя. Скопируйте функции этой страницы (они самодостаточны, нужны torch, einops, tqdm) и повторите три эксперимента на CPU. Совпадёт ли итоговая средняя награда с 2,67 и 3,33 при torch.manual_seed(5)?
2. Вылечите схлопывание. Добавьте к потерям энтропийный бонус $-0{,}01 \cdot H(\pi)$ (энтропия по словарю в каждой позиции) или увеличьте num_responses до 32. Удастся ли поднять среднюю награду выше 3,33?
3. Штраф KL. Исправьте model.clone() на copy.deepcopy(model) и запустите с kl_penalty=0.1, loss_mode="clipped". Как штраф влияет на скорость роста награды?
4. Перейдите к настоящей LM. Сделайте практикум 6.4.5: там та же схема (группа ответов → награды → нормировка → обрезанные потери), но политика — маленькая Qwen, а задача — арифметика.

Словарик

EN	RU
state, action, reward, transition	состояние, действие, награда, переход
policy, rollout (episode, trajectory)	политика, роллаут (эпизод, траектория)
outcome reward, verifiable reward	награда за результат, проверяемая награда
policy gradient, log-derivative trick	градиент по политике, трюк с логарифмической производной
baseline, variance reduction	базовая линия, снижение дисперсии
value function V(s), Q-function Q(s, a), advantage A(s, a)	функция ценности, Q-функция, преимущество
centered / normalized rewards	центрированные / нормированные награды
probability ratio, clipping	отношение вероятностей, обрезка
KL penalty, reference model	штраф KL, эталонная модель
reward shaping, reward hacking	формирование награды, взлом награды
mode collapse, entropy bonus	схлопывание мод, энтропийный бонус