ТЕХНИЧЕСКОЕ ЗАДАНИЕ: Реализовать latent reasoning (∇-Reasoner)

1. Цель задачи

Разработать и реализовать механизм латентного рассуждения (latent reasoning) с градиентным спуском в скрытом пространстве модели. Цель — улучшить качество ответов LLM на логические задачи без изменения весов самой модели, используя только дополнительный шаг оптимизации эмбеддинга на этапе инференса. Вы научитесь понимать математику градиентного спуска в непрерывном латентном пространстве, реализовать его на PyTorch и протестировать на синтетических задачах.

Ключевой результат Рабочий код класса LatentReasoner, который применяет несколько шагов градиентного спуска к скрытому представлению промпта, улучшая правильность ответов модели на тестовом наборе из 20 логических задач на 10% и более.

2. Исходные данные

Перед началом необходимо иметь:

Если нет реального инструмента — симулируем:

1. Сгенерировать тестовый набор
   Создайте CSV-файл с 20 задачами вида:
   "У Алисы 3 яблока, у Боба 5. Алиса отдала Бобу 2. Сколько яблок стало у Боба?" → ответ 7.
   Включите задачи на арифметику, логику, факты. Используйте простые шаблоны.

2. Взять базовую модель — скачайте модель через Hugging Face AutoModelForCausalLM. Используйте model.get_input_embeddings() для доступа к скрытому пространству.

3. Подготовить baseline
   Запустите модель без доработок, запишите accuracy на тестовом наборе (считается по точному совпадению последнего числа).

3. Технологический стек

4. Этапы выполнения

Этап 1: Теоретическое обоснование и проектирование (1 час)

Действия

1. Изучите концепцию
   Прочитайте секцию 3.2 из статьи "Training Language Models to Reason with Self-Consistency" (Wang et al., 2022) или "Let's Think Step by Step with Gradient Descent" (аналоги).
   Поймите: мы не меняем веса сети, а оптимизируем эмбеддинги входных токенов прямо во время инференса, используя градиент от loss между сгенерированным ответом и правильным ответом (или от loss, который штрафует неуверенность).

2. Спроектируйте схему ∇-Reasoner
   Нарисуйте блок-схему:

   Prompt tokens → Embedding lookup → [замороженные эмбеддинги]  
   Добавляем learnable residual (только для отдельных токенов)  
   Residual → обновляется через градиенты от лосса  
   → Forward pass через замороженную модель → логиты  
   → Сравнение с целевым ответом → loss → градиент → update residual
   

3. Выберите функцию потерь
   Для точного ответа (закрытая задача) используйте cross-entropy между логитами вероятности правильного токена и one-hot вектором.
   Для ответа без известной цели используйте unsupervised loss: негэнтропия (отрицательная энтропия) логитов, штраф за неуверенность.

Ожидаемый результат этапа
Маркдаун-документ (или раздел в тетрадке) с описанием алгоритма, выбранной loss-функции и планом интеграции.

Этап 2: Базовая реализация градиентного спуска в латентном пространстве (1.5 часа)

Действия

1. Создайте класс LatentReasoner

   class LatentReasoner:
       def __init__(self, model, tokenizer, lr=0.1, steps=20):
           self.model = model
           self.tokenizer = tokenizer
           self.lr = lr
           self.steps = steps
           # Замораживаем все параметры модели
           for param in self.model.parameters():
               param.requires_grad = False
   

2. Реализуйте метод forward_with_hidden
   Напишите функцию, которая:

   • Токенизирует промпт.
   • Получает эмбеддинги (через model.get_input_embeddings()).
   • Создаёт обучаемый тензор residual (shape: batch, seq_len, Вики/Embedding dimension|hidden_dim), заполненный нулями.
   • Добавляет residual к эмбеддингам.
   • Запускает forward-pass модели с output_hidden_states=True.
   • Возвращает логиты, скрытые состояния и входные эмбеддинги с residual.

3. Реализуйте шаг градиентного спуска
   Для одного шага:

   def gradient_step(self, prompt, target_text):
       # forward pass через модель с residual
       logits, hidden = self.forward_with_hidden(prompt)
       # вычисляем loss: cross-entropy между логитами и target
       loss = F.cross_entropy(logits, target_ids)
       loss.backward()  # градиенты придут только к residual
       with torch.no_grad():
           self.residual -= self.lr * self.residual.grad
           self.residual.grad.zero_()
       return loss.item()
   

4. Оттестируйте на одном примере
   Запустите 5 шагов для простой задачи («2+2=?»), выведите loss на каждом шаге. Убедитесь, что loss уменьшается.

Ожидаемый результат этапа
Класс LatentReasoner с работающим градиентным спуском. Скрипт, который за 10 шагов снижает loss на одном примере.

Этап 3: Интеграция с генерацией ответа (1.5 часа)

Действия

1. Модифицируйте метод generate
   В классе LatentReasoner добавьте метод generate_with_reasoning(prompt, max_new_tokens=20):

   • Сначала запустите градиентный спуск на self.steps шагов, используя целевой ответ, если он известен (для обучающих примеров).
     Примечание: В реальном сценарии целевого ответа нет, поэтому используется unsupervised loss (например, максимизация уверенности модели в первом токене).
   • После оптимизации residual, выполните обычную генерацию (через model.generate()) с финальными скрытыми состояниями.

2. Добавьте unsupervised loss
   Если правильный ответ неизвестен (для теста), определите loss как:

   probs = F.softmax(logits, dim=-1)
   entropy = -(probs * torch.log(probs + 1e-8)).sum(-1)
   loss = entropy.mean()  # минимизируем энтропию → модель более уверена
   

3. Протестируйте на одном примере из тестового набора
   Сравните ответ до и после ∇-спуска с помощью generate_with_reasoning vs model.generate(). Запишите визуально, улучшился ли ответ.

4. Добавьте визуализацию
   Выведите изменения residual на нескольких шагах (среднее по dim, max изменение).

Ожидаемый результат этапа
Рабочая генерация с латентным reasoning. Вывод сравнения ответов для одного примера.

Этап 4: Полноценное тестирование на наборе задач (1.5 часа)

Действия

1. Запуск baseline
   Для всех 20 задач из тестового набора получите ответ обычной модели (без ∇). Запишите accuracy (точное совпадение с правильным ответом).

2. Запуск ∇-Reasoner
   Для каждой задачи:

   • Используйте генерацию с reasoning (с unsupervised loss, т.к. правильный ответ неизвестен).
   • Количество шагов = 20, learning rate = 0.1.
   • Запишите ответ.
   • Сравните с правильным ответом.

3. Соберите статистику
   Создайте таблицу:

   Задача	Baseline answer	∇-answer	Correct?	Baseline correct?
   ...	...	...	...	...

   Рассчитайте:

   • Baseline accuracy
   • ∇-accuracy
   • относительное улучшение.

4. Экспериментируйте с гиперпараметрами
   Попробуйте разные learning rates (0.01, 0.1, 0.5), количество шагов (10, 30). Зафиксируйте влияние на accuracy и время.

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

Этап 5: Анализ и выявление edge cases (1 час)

Действия

1. Исследуйте случаи, где ∇-Reasoner не помог
   Для задач, где accuracy не улучшилась, выведите:

   • График loss по шагам
   • Изменение residual (норма L2)
   • Логиты до и после

2. Проверьте градиентный взрыв/затухание
   Вычислите норму градиентов на каждом шаге. Если норма > 100 или < 1e-5, примените gradient clipping (torch.nn.utils.clip_grad_norm_).

3. Документируйте математику
   В отдельном разделе пропишите формулу градиентного спуска в латентном пространстве: [ h_{t+1} = h_t - \eta \nabla_{h} \mathcal{L}(f_\theta(h_t), y) ] Объясните, почему градиенты не проходят в веса модели (они заморожены) — используйте detach() в нужных местах.

4. Сравните с типичным fine-tuning
   Напишите 2-3 предложения о различиях: fine-tuning меняет веса (параметры дорогие), latent reasoning меняет вход (параметры – только residual, всего 1seq_len768 ≈ 30k параметров для 128 токенов).

Ожидаемый результат этапа
Отчёт о найденных edge cases, графики, математическая нотация в README.

5. Критерии приемки (Definition of Done)

• Реализован класс LatentReasoner с методами __init__, gradient_step, generate_with_reasoning.
• Код работает на одной из лёгких LLM (TinyLlama, Qwen 0.5B) без модификации весов.
• Функция потерь реализована: для обучения с учителем (cross-entropy) и без учителя (минимизация энтропии).
• На тестовом наборе из 20 задач ∇-Reasoner показывает улучшение accuracy не менее чем на 10% относительно baseline.
• Код содержит визуализацию изменения loss и residual на первых 5 шагах (математически корректно).
• Написаны модульные тесты (pytest) для проверки:
  • Градиенты проходят только в residual, веса модели не обновляются.
  • Loss уменьшается с каждым шагом.
• Есть документация (README.md) с описанием алгоритма и математикой.
• Гиперпараметры (lr, steps) вынесены в конфиг и легко изменяются.

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

Основной артефакт

• Файл latent_reasoner.py (или .ipynb) с реализацией класса LatentReasoner.
• Тестовый набор задач в test_tasks.csv с правильными ответами.
• README.md с подробным описанием: цель, математика, результаты на тестовом наборе, инструкция по запуску.

Дополнительные артефакты

• Графики loss и нормы residual (PNG/PDF).
• Таблица сравнения accuracy для разных гиперпараметров.

Содержание README (обязательные разделы):

• Формулировка задачи
• Теоретическая основа (формулы, блок-схема)
• Описание кода и ключевых моментов
• Результаты экспериментов (таблица)
• Выводы и возможные улучшения

7. Возможные сложности и их решение

8. Бюджет времени (оценка)

Примечание: Время может увеличиться, если потребуется больше экспериментов с гиперпараметрами.

9. Связанные вопросы из базы знаний

10. Чек-лист самопроверки

• Я понимаю, почему градиенты не проходят в веса замороженной модели, и создаю residual как nn.Parameter.
• Я реализовал как supervised, так и unsupervised loss для flexibility.
• Код корректно обрабатывает batch_size=1 (можно расширить на batch, но не обязательно).
• Я проверил, что loss монотонно (или в среднем) уменьшается на этапе оптимизации.
• Я сравнил хотя бы 3 гиперпараметра (lr, steps) и задокументировал результаты.
• Функция generate_with_reasoning работает детерминированно при одном seed.
• README содержит раздел с математикой, чтобы любой инженер воспроизвёл без подсказок.