ТЕХНИЧЕСКОЕ ЗАДАНИЕ: Реализовать verifier-guided decoding

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

Разработать механизм verifier-guided decoding (VGD) для генерации рассуждений (chain-of-thought) малой LLM. Небольшой verifier (например, на 0.5B параметров) на каждом шаге (токене или группе токенов) оценивает вероятность того, что текущий путь приведёт к правильному ответу. На основе оценки decoding либо прерывается, либо происходит откат (backtracking) к предыдущим шагам. Ключевая идея — отсекать неверные ветви рассуждения рано, экономя тестовое время (test-time compute) и повышая точность финального ответа.

Ключевой результат Функциональный пайплайн VGD, который для набора математических или логических задач (например, GSM8K или PrOntoQA) показывает улучшение accuracy на ≥10% при фиксированном бюджете токенов (или сокращении числа токенов на ≥30% при той же accuracy) по сравнению с обычным жадным декодированием.

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

Если нет реальных моделей (например, нет GPU) — симулируем:

1. Заменить LLM на OpenAI API (gpt-3.5-turbo-instruct), но verifier — локальная малая модель (через Transformers) на CPU.
2. Подготовить скрипт, который делает вызовы к API и локально оценивает шаги.
3. Ограничить размер тестового датасета (50 примеров) для ускорения итераций.

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

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

Этап 1: Настройка окружения и загрузка данных (оценка: 30 мин)

Действия

1. Создать виртуальное окружение и установить зависимости:

   python -m venv vgd_env
   source vgd_env/bin/activate
   pip install transformers torch datasets accelerate wandb vllm
   

2. Загрузить датасет GSM8K (тестовый сплит):

   from datasets import load_dataset
   dataset = load_dataset("gsm8k", "main", split="test")
   questions = dataset["question"]
   answers = dataset["answer"]   # содержит финальное число с пояснением
   

3. Выбрать 50 примеров для отладки и 200 для финальной оценки.
4. Загрузить базовую LLM (генератор) и verifier:

   from transformers import AutoModelForCausalLM, AutoTokenizer
   gen_tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3")
   gen_model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3", device_map="auto")
   
   vrf_tokenizer = AutoTokenizer.from_pretrained("HuggingFaceTB/SmolLM-360M")
   vrf_model = AutoModelForSequenceClassification.from_pretrained("HuggingFaceTB/SmolLM-360M", num_labels=2).to("cuda")
   

   Примечание: для verifier нужно добавить классификационную голову (binary: правильный/неправильный шаг). Можно использовать предобученный allenai/longformer-base-4096 с head, но здесь используем SmolLM и дообучим на синтетических данных (этап 2).

Ожидаемый результат этапа Окружение готово, данные загружены, модели загружены в память. Проверочный вызов генерации выдаёт первые 5 шагов.

Этап 2: Подготовка verifier (обучение/настройка) (оценка: 2 часа)

Действия

1. Сгенерировать синтетические данные для обучения verifier:
   • На 200 примерах из датасета (train split) сгенерировать 3-5 разных путей рассуждения через обычное жадное декодирование (с разными температурами).
   • Каждый путь разбить на шаги (по предложениям, разделённым \n).
   • Для каждого шага пометить label=1 если итоговый ответ правильный, label=0 если неправильный. (Правильность определяется по совпадению финального числа с эталоном).
2. Дообучить verifier (sequence classification):
   • Использовать токенизатор SmolLM с padding=True.
   • Каждый шаг (текст) подаётся как вход, метка 0/1.
   • Обучить 1 эпоху, batch_size=8, lr=2e-5.

   from transformers import Trainer, TrainingArguments
   training_args = TrainingArguments(
       output_dir="./verifier_model",
       per_device_train_batch_size=8,
       num_train_epochs=1,
       evaluation_strategy="no",
       save_total_limit=1,
   )
   trainer = Trainer(model=vrf_model, args=training_args, train_dataset=training_dataset)
   trainer.train()
   

3. Оценить качество verifier на валидационной выборке (accuracy, precision, recall).
   • Цель: accuracy ≥0.8 на бинарной классификации шагов.

Ожидаемый результат этапа Обученная модель verifier, способная оценить любой шаг рассуждения и дать вероятность правильности p ∈ [0,1]. Сохранена в ./verifier_model.

Этап 3: Реализация verifier-guided decoding (оценка: 3 часа)

Действия

1. Разработать функцию generate_with_verifier
   • На вход: question (str), generator, verifier, tokenizers, параметры (beam_width, max_steps, threshold_early_stop).
   • Процесс:
     1. Инициализировать очередь путей (каждый путь — список токенов и текущая вероятность).
     2. На каждом шаге генерации (каждый новый токен) расширять все пути на top-k (k=beam_width) вариантов с помощью топ-вероятностей генератора.
     3. После генерации полного шага (определяем по разделителю, например \n или точка с пробелом) передать текст шага verifier.
     4. Если verifier даёт p < threshold_early_stop (например, 0.3) — отсечь путь (удалить из очереди).
     5. Если p > threshold_good (0.8) — оставить путь с повышенным приоритетом (умножить score на 1.2).
     6. Продолжать, пока не будет сгенерирован финальный ответ (символ = или The answer is), либо не кончится лимит шагов.
   • Вернуть лучший финальный ответ (по суммарной вероятности или по последней оценке verifier).
2. Реализовать beam search с отсечениями
   • Использовать model.generate с num_beams=... и early_stopping=True, но дополнить кастомным логированием через logits_processor или переопределить _beam_search (сложно). Альтернатива: написать свой простой beam search с нуля (для понимания) на <=200 строк.
3. Интегрировать verifier как callback
   • Создать класс VerifierCallback с методом __call__(input_ids, step_text), возвращающим решение (continue/stop/backtrack).
4. Опционально: реализовать backtracking (откат на 1-2 шага) при очень низкой оценке verifier (p < 0.1).

Ожидаемый результат этапа Скрипт vgd_decode.py, который для одного вопроса возвращает финальный ответ и трассировку (список шагов с их оценками). Можно протестировать на 5-10 примерах.

Этап 4: Эксперимент и сравнение с baseline (оценка: 1.5 часа)

Действия

1. Запустить обычное greedy decoding (без verifier) на тестовых 200 примерах:
   • Параметры: max_new_tokens=512, do_sample=False, temperature=0.
   • Зафиксировать accuracy и среднее число токенов на один ответ.
2. Запустить verifier-guided decoding на тех же примерах:
   • Параметры: beam_width=5, max_steps=10, threshold_early_stop=0.3.
   • Зафиксировать те же метрики + доля прерванных путей (отношение числа прерванных шагов к общему числу шагов).
3. Сравнить результаты в таблице (экспорт в CSV или wandb).
4. Провести ablation: изменить threshold (0.1, 0.5, 0.7) и beam_width (3, 5, 10).

Ожидаемый результат этапа Чёткая сводка метрик, подтверждающая улучшение accuracy и/или снижение токенов при использовании VGD.

Этап 5: Документирование и упаковка (оценка: 1 час)

Действия

1. Написать README.md с описанием архитектуры VGD, инструкцией по запуску, примерами.
2. Зафиксировать конфигурации лучшего эксперимента (threshold, beam_width) в файле config.yaml.
3. Создать ноутбук demo.ipynb с визуализацией
   • Пример одного вопроса с отображением шагов, оценок verifier и точек прерывания.
4. Выложить код на GitHub (или другую платформу) и добавить ссылку.

Ожидаемый результат этапа Репозиторий с чистым кодом, инструкциями и результатами.

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

• Код verifier-guided decoding работает end-to-end и генерирует ответ для заданного вопроса.
• Verifier обучен и даёт accuracy ≥0.8 на валидационном наборе шагов.
• Эксперимент на 200 примерах GSM8K показал улучшение accuracy на ≥10% по сравнению с greedy decoding.
• Среднее число токенов на ответ снижено на ≥30% при той же accuracy.
• Доля прерванных путей (неправильных шагов, отсечённых verifier) зафиксирована и составляет ≥30% (чтобы verifier действительно отсекал).
• Код структурирован: отдельные модули для verifier, beam search, экспериментов.
• README включает инструкцию по воспроизведению и пример вывода.
• Конфигурация (параметры декодирования, пути к моделям) вынесена в YAML.

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

• Основной артефакт Python-скрипт vgd.py с классом VGDGenerator и функцией run_experiment.
• Содержание реализация beam search с verifier callback, отсечение неверных путей, экспорт метрик в wandb/CSV.
• Дополнительно веса обученного verifier (./verifier_model/), конфиг config.yaml, файл с результатами results.csv, ноутбук-демо.
• Метрики: accuracy (финальный ответ), среднее число токенов, histogram оценок verifier по шагам.

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

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

Примечание Для первого раза можно упростить — пропустить Этап 2 (использовать предобученный verifier-классификатор, например, roberta-large-mnli с правилами), тогда общее время ~5 ч.

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

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

• Я понимаю, как verifier получает на вход шаг рассуждения и выдаёт вероятность корректности.
• Я реализовал beam search с возможностью прерывания путей по оценке verifier.
• Я обучил verifier на синтетических данных и убедился в его качестве (accuracy ≥0.8).
• Я сравнил VGD с greedy decoding на репрезентативной выборке и получил улучшение accuracy или сокращение числа токенов.
• Я задокументировал конфигурацию эксперимента и результаты так, чтобы их можно было воспроизвести.