Как работает ∇-Reasoner (градиентный спуск в пространстве токенов на этапе теста)?

Краткий тезис

∇-Reasoner — это метод улучшения качества рассуждений LLM на этапе инференса, при котором скрытые представления токенов оптимизируются с помощью градиентного спуска, направляемого reward model. В отличие от обычного авторегрессивного декодирования, ∇-Reasoner выполняет 1–5 итераций backpropagation через всю модель, обновляя эмбеддинги или скрытые состояния так, чтобы максимизировать ожидаемую награду. Это позволяет LLM «размышлять» над ответом, исправляя неверные шаги, ценой 3–5-кратного увеличения времени генерации, но давая прирост точности на 10–20% на сложных задачах (математика, логика).

1. Термин: ∇-Reasoner

∇-Reasoner (произносится «набла-ризонер») — это техника, впервые предложенная в ICLR 2026, которая объединяет идеи из дифференцируемого программирования и обучения с подкреплением. Название происходит от символа ∇ (набла), обозначающего градиент. Основная идея: на этапе теста (инференса) мы не просто генерируем токены слева направо, а итеративно улучшаем текущий «черновик» ответа, пропуская градиенты через модель.

Ключевые компоненты:

• Reward model: небольшая нейросеть, которая оценивает качество текущего ответа (например, правильность математического решения, согласованность фактов).
• Оптимизируемое пространство: вместо дискретных токенов мы работаем с непрерывными векторными представлениями — эмбеддингами токенов или скрытыми состояниями.
• Backpropagation: обратное распространение ошибки от reward model через LLM до этих представлений, аналогично обучению, но без изменения весов модели.

2. Мотивация: почему нужен градиентный спуск на инференсе

Стандартное авторегрессивное декодирование имеет фундаментальный недостаток: решение принимается одношагово, без возможности вернуться и исправить ошибку. Даже такие техники, как Beam Search или self-consistency, лишь перебирают несколько траекторий, но не оптимизируют сами представления. ∇-Reasoner предлагает непрерывную оптимизацию:

• Ошибка в середине длинной цепочки рассуждений может быть исправлена путём градиентного смещения эмбеддингов последующих токенов.
• Reward model даёт количественную оценку всего ответа, а градиентный спуск находит направление, в котором нужно изменить скрытые представления, чтобы увеличить награду.
• Это напоминает метод 「замороженной модели с обучаемыми входными данными」 (adversarial reprogramming), но для рассуждений.

3. Архитектура: reward model и оптимизируемые представления

Система состоит из трёх частей:

Обычно оптимизируют эмбеддинги токенов на каждом шаге генерации. Это позволяет напрямую влиять на то, какой токен будет выбран на следующем шаге, без изменения весов модели.

4. Как работает градиентный спуск в дискретном пространстве токенов

Проблема: пространство токенов дискретно (выбор из словаря), а градиентный спуск требует непрерывного пространства. ∇-Reasoner использует relaxation:

• Вместо one-hot вектора токена мы используем непрерывный soft-эмбеддинг — взвешенную сумму всех эмбеддингов словаря (аналогично Gumbel-Softmax или continuous relaxation в нейронном машинном переводе).
• Во время forward pass модели мы передаём этот непрерывный эмбеддинг, а модель вычисляет скрытые состояния и логарифмы вероятностей.
• Reward model оценивает выход (например, как вероятность правильного ответа). Градиент reward по отношению к мягкому эмбеддингу вычисляется и обновляет его.
• После 1–5 итераций мы «застываем» (harden) — выбираем наиболее вероятный токен из словаря (argmax) и переходим к следующей позиции.

Существует также вариант прямого обновления скрытых состояний без relaxation: мы храним для каждого токена его скрытое состояние из последнего слоя и обновляем его градиентом от reward. Но на выходе модели это всё равно должно отображаться в logits, поэтому relaxation остаётся необходимостью.

5. Алгоритм шаг за шагом

Пусть у нас есть запрос q. Мы генерируем ответ длины T. Процесс для одной итерации:

1. Прямой проход (forward): Для каждого шага t от 1 до T:
   • Получаем текущий непрерывный эмбеддинг e_t (изначально — эмбеддинг первого токена, для остальных — мягкий эмбеддинг от предыдущей итерации).
   • Вычисляем скрытые состояния через LLM, получаем logits l_t.
   • Выбираем следующий токен (для начальной итерации можно использовать argmax или sampling).
2. Оценка награды Пропускаем генерируемый ответ (или его части) через reward model. Получаем скалярную награду R.
3. Обратный проход (backward): Вычисляем градиент ∇R по отношению к каждому e_t.
4. Обновление Делаем шаг градиентного спуска: e_t := e_t + α * ∇_{e_t} R (обычно с небольшим learning rate ≈ 0.01–0.1).
5. Повтор Возвращаемся к шагу 1 с обновлёнными e_t. Через 1–5 итераций фиксируем токены (выбираем argmax из logits, соответствующих обновлённым эмбеддингам).

Псевдокод на Python:

def generate_with_nabla_reasoner(model, reward_model, query, num_iter=3, lr=0.05):
    tokens = model.tokenize(query)
    embeddings = model.get_embeddings(tokens)  # [seq_len, d_model]
    # инициализируем soft-эмбеддинги для первых K токенов ответа (или все сразу)
    answer_prefix = model.generate_tokens(query, max_new_tokens=10)  # первый черновик
    soft_embs = model.get_embeddings(answer_prefix).detach().clone().requires_grad_(True)

    for i in range(num_iter):
        # прямой проход: передаём soft_embs как вход
        outputs = model(inputs_embeds=soft_embs.unsqueeze(0))
        logits = outputs.logits[0]  # [seq_len, vocab]
        # получаем вероятности токенов (gumbel-softmax не нужен, используем softmax)
        probs = torch.softmax(logits, dim=-1)
        # выбираем наиболее вероятные токены для оценки
        pred_tokens = probs.argmax(dim=-1)
        reward = reward_model(pred_tokens, query)  # скаляр

        # обратный проход
        reward.backward()
        with torch.no_grad():
            soft_embs += lr * soft_embs.grad
            soft_embs.grad.zero_()

    # финализация: выбираем argmax
    final_tokens = model.generate(inputs_embeds=soft_embs.unsqueeze(0))
    return model.decode(final_tokens)

6. Пример: математическая задача

Задача «Решите уравнение x^2 - 5x + 6 = 0».

Обычная LLM может сгенерировать: «x = 2 и x = 3» (правильно) или «x = 2.5, 3.5» (ошибка).

∇-Reasoner с reward model, обученной проверять правильность корней (подстановка), начнёт с черновика «x = 2.5, 3.5». Reward будет низким (0.2). Градиентный спуск сместит эмбеддинги токенов, чтобы следующие шаги дали logits, приближающиеся к правильным токенам (числа 2 и 3). После 3 итераций reward станет 0.95, и ответ превратится в верный.

7. Компромисс качество/скорость

∇-Reasoner особенно эффективен, когда reward model может дать чёткий сигнал (например, проверка математического равенства). Для задач с субъективной оценкой (суммаризация, креативное письмо) улучшение меньше.

8. Сравнение с другими методами улучшения рассуждений

∇-Reasoner занимает нишу между быстрым CoT и медленным MCTS: даёт значительное улучшение при умеренном замедлении.

9. Ограничения и направления улучшений

Основные ограничения

• Зависимость от качества reward model: плохая reward может привести к неверной оптимизации.
• Необходимость непрерывной релаксации: soft-эмбеддинги могут быть далеки от реальных токенов, и после hardening качество может упасть.
• Градиенты через всю LLM (пусть и замороженную) требуют много памяти (нужно хранить активации для backprop).
• Применимость только к задачам с вычислимым reward (математика, программирование, фактчекинг).

Возможные улучшения

• Использовать стохастический градиент и Gumbel-Softmax для лучшего приближения дискретного выбора.
• Оптимизировать не все позиции, а только «критические» точки (например, где uncertainty высока).
• Применять early stopping на основе прироста reward.
• Комбинировать с chain-of-thought: дать модели CoT-подсказку, а затем уточнить её ∇-Reasoner.

10. Пет-проект для закрепления

Задача Реализовать упрощённый ∇-Reasoner для задачи «решение квадратного уравнения».

Инструменты PyTorch, transformers (GPT-2), sklearn (для reward model).

Шаги:

1. Выберите маленькую языковую модель (GPT-2 small, 124M).
2. Создайте датасет квадратных уравнений с корнями (1000 примеров).
3. Обучите reward model: простой MLP, который принимает эмбеддинги сгенерированного ответа (усреднённый last hidden state) и предсказывает, является ли ответ правильным (бинарная классификация).
4. Напишите функцию инференса с ∇-Reasoner:
   • Генерируйте начальный ответ (10–20 токенов) через обычную генерацию.
   • Переведите токены в непрерывные эмбеддинги (soft-эмбеддинги через Gumbel-Softmax).
   • Для 5 итераций: forward через GPT-2, получить reward от reward model, backward, обновить эмбеддинги.
   • Финализируйте: замените soft-эмбеддинги на argmax.
5. Сравните точность на тестовом наборе с обычной генерацией (CoT) и с 3, 5 итерациями ∇-Reasoner.

Ожидаемый результат

• Прирост точности на 10–20% на сложных уравнениях (с дробными корнями).
• Время генерации увеличится в 3–5 раз (измерьте).
• Вы увидите, что reward model даёт больший градиент для неверных чисел, смещая их в сторону верных.

11. Связь с другими вопросами

12. Навигация

Навигация

• Предыдущий: 637
• Следующий: 639
• Индекс: 00. Индекс разборов