Как устроена Nemotron-4 340b

Компания NVIDIA выпустила одну из самых крупных и качественных open source моделей — Nemotron-4 340B. Сегодня разберём сразу две статьи, в которых рассказывается об её устройстве.

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

Вопреки современным трендам, NVIDIA использовала функцию активации squared ReLU, но при этом отказалась от гейта. Параметризация layer-norm также нестандартна: вместо инициализации единицами применяется инициализация нулями с добавлением единицы на каждом форварде.

В датасетах alignment NVIDIA использовала небольшое количество человеческой разметки. Датасет обучения reward-модели включает в себя всего лишь 10 тысяч пар ответов, размеченных по пяти параметрам: от общей пользы до связности и сложности.

SFT также включал лишь 10 тысяч написанных человеком примеров. При этом суммарно в SFT-датасетах вышло порядка миллиона примеров, из которых почти все были сгенерированы. Интересно, что сгенерированы были и сами промты — для этого NVIDIA использовала Mixtral 8x7B.

По итогу получилось два SFT-датасета и два раунда обучения. Сперва модель обучали Code SFT на 800 тысячах примеров. А затем ту же модель — дообучали General SFT ещё на 200 тысячах примеров.

Далее следовал один раунд DPO и три раунда собственного метода NVIDIA — RPO (Reward-aware Preference Optimization). Для обучения стандартного DPO используется только порядок внутри пары: оптимальная политика максимизирует расстояние между наградой за выбранный и отклонённый ответ. Такой подход может приводить к переобучению, поэтому в RPO авторы предложили аппроксимировать саму разницу между наградами, которые считают Reward-моделью.

После всех этапов обучения и тестирования модель Nemotron показала хорошие результаты на бенчмарках вроде MMLU. Интересно, что модель почти не росла после раунда DPO, но последующие RPO способствовали улучшению.

Nematron опережает большинство других open source-решений и тягается почти на равных с закрытыми. Скажем, в Arena Hard модель NVIDIA показала результат 54,2, что выше чем у Claude 3 Sonnet (46,8) и Qwen2 72b (48,1).

В сравнении, проведённом людьми, у Nemotron практически паритет c GPT-4 по всем задачам, исключая переписывание текста. Здесь модель NVIDIA справлялась лишь в 3,03% случаев. Однако стоит учитывать, что «живые» замеры проводилось на небольшом количестве примеров — их было всего 136.

Рассказывайте в комментариях, что думаете о Nemotron-4?

Разбор подготовил ❣ Руслан Васильев

Душный NLP

DeepSeek-V2 — MoE-модель с технологией MLA

Компания DeepSeek представила MoE-модель DeepSeek-V2 на 236 миллиардов параметров. Сегодня мы разберём статью, которая рассказывает об особенностях этой модели.

Модель состоит из 236 миллиардов параметров, однако на каждый токен активно используется лишь 27 миллиардов. Это позволило значительно увеличить длину контекста до 128 тысяч токенов.

Для сравнения и оценки результатов в DeepSeek использовали dense-модель на 67b. Авторы смогли существенно снизить стоимость обучения и уменьшить размер KV-кэша, что позволило увеличить скорость генерации токенов почти в 8 раз. Эти достижения обеспечили модели преимущество в производительности.

В DeepSeek-V2 внедрён модифицированный attention, который получил название MLA (Multi-Head Latent Attention). Этот механизм позволяет кэшировать уменьшенный вектор с последующим его восстановлением, что значительно снижает объём ресурсов, необходимый для работы модели. Как устроен MLA в сравнении с MHA, GQA и MQA, вы можете увидеть на изображении выше.

При таком подходе нельзя напрямую использовать RoPE (Rotary Position Embedding) — в противном случае матрица RoPE становится неразделимой с остальными операциями, что значительно снижает эффективность вывода. В DeepSeek предложили стратегию так называемого «развязанного RoPE» (decoupled RoPE), в которой дополнительные ключи несут позиционную информацию.

Модель показала хорошие результаты в бенчмарках. В MMLU она добилась показателя в 78,3%, а в тестах на китайском языке, таких как CLUEWSC, C-Eval и CMMLU, DeepSeek-V2 показала лучшие результаты среди всех открытых моделей.

Разбор подготовил ❣ Сергей Горбунов

Душный NLP

Масштабирование и параметризация

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

Авторы статьи Tensor Programs V предлагают использовать Maximal Update Parametrization (µP) — перенос параметров с маленькой модели на большую без дополнительной настройки.

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

Чтобы достичь стабильности гиперпараметров, нужно правильно масштабировать спектральную норму матриц весов — показатель максимально возможного растяжения или сжатия вектора при его умножении на матрицу. Авторы статьи отмечают, что добиться стабильности можно двумя способами: правильным масштабированием инициализаций и послойных LR, либо напрямую спектральной нормализацией матриц весов и их обновлений в процессе обучения.

Благодаря такому решению масштаб признаков и их изменений на каждом шаге сохраняется при увеличении размера сети — этого оказывается достаточно для стабильности гиперпараметров. В статье A Spectral Condition for Feature Learning предполагается, что обновления весов в градиентном спуске имеют низкий ранг и хорошо согласуются с векторами активаций. Однако предположение о такой согласованности на самом деле выполняется не всегда, поэтому в более свежей статье Scaling Exponents Across Parameterizations and Optimizers авторы предлагают дальнейшее улучшение метода с поправкой на это.

Расскажите в комментариях, что думаете по поводу этих методов!

Разбор подготовил ❣ Дмитрий Лунин

Душный NLP

Branch-Train-MiX — метод создания MoE-моделей

Сегодня рассмотрим статью, в которой предложен метод получения MoE-модели (Mixture-of-Experts) из обычной, dense-модели. Авторы назвали его Branch-Train-MiX (BTX).

Для начала напомним, что такое MoE. Это архитектурный подход, который предполагает использование в трансформенных слоях нескольких FNN-блоков — экспертов. У каждого из них предположительно есть собственная узкая специализация, поэтому над решением поставленной модели задачи они работают не одновременно. Роутер выбирает, какие эксперты лучше подходят для обработки токенов на каждом конкретном этапе.

Во время инференса модели, построенные на архитектуре MoE, задействуют меньше параметров, чем их dense-аналоги — и без сильной потери качества. Кроме того, они прекрасно масштабируются. Поэтому понятно стремление превратить обычную модель в MoE.

В основе идеи авторов лежит метод BTM (Branch-Train-Merge). Суть его заключается в том, чтобы взять несколько одинаковых LLM, параллельно обучить их на разных датасетах, а затем агрегировать предсказания каждой модели во время инференса.

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

Авторы взяли три копии Llama-2 7B и дообучили каждую на своём домене: математика (на тех же данных, что и Llemma), программирование (на тех же данных, что CodeLlama) и общие знания (на «Википедии»). К финальной модели добавили оригинальную Llama-2 7B, не дообученную на чём-то конкретном. Получившуюся MoE-модель файнтюнили на источниках, которые применялись для обучения всех четырёх экспертов. В результате модель не сильно уступает тем же Llemma и CodeLlama в вопросах математики и программирования.

Интересно и то, что BTX-модели обретают интерпретируемость. Авторы показывают, что в эксперта, обученного на определённом домене, попадают токены из сэмпла этого же домена.

Рассказывайте в комментариях, что думаете про BTX!

Разбор подготовил ❣ Александр Пацация

Душный NLP

Self-Rewarding Language Models

В сегодняшней статье — концепция Self-Rewarding Language Models. Они не только генерируют ответы, но и оценивают их качество в процессе обучения.

Традиционные подходы к обучению, такие как обучение с подкреплением на основе обратной связи от людей (RLHF), ограничены качеством и объемом человеческих данных. Авторы статьи предлагают модель, которая сама создает инструкции, генерирует ответы и оценивает их качество, используя технику LLM-as-a-Judge.

Обучение проводится через итерации, каждая из которых состоит из следующих этапов:

1. Генерация ответов на основе текущего датасета. На первой итерации — созданного человеком.
2. Оценка ответов — модель сама оценивает качество ответов по определённым критериям, таким как релевантность, полезность, чёткость и уровень экспертизы.
3. Создание самоинструкций для следующего шага — модель создает новые инструкции на основе исходного набора данных.
4. Следующая итерация — дообученные модели используются для последующих итераций обучения. Это ведёт к улучшению качества как генерации ответов, так и их оценки.

В экспериментах использовали Llama 2-70B, обученную на датасете Open Assistant. Модель прошла через три итерации обучения с самосозданными данными.

— Первичная модель, файнтюненная на IFT (Instruction Fine-Tuning) и EFT (Evaluation Fine-Tuning)-данных с помощью SFT, показала результаты, сопоставимые с базовой моделью.
— Вторая итерация, обученная на AIFT (AI Feedback Training) c помощью DPO, выдавала лучший ответ в 55,5% случаев, а в 32,8% — того же качества.
— Третья итерация в 47,7% случаев выдавала лучший ответ, а в 39,8% — того же качества.

Модель, прошедшая три итерации обучения, продемонстрировала улучшения в тестах AlpacaEval 2.0, превзойдя модели Claude 2, Gemini Pro и GPT-4 0613. Модель третьей итерации показала 20,44% выигрышей в тестах, что значительно выше предыдущих итераций (9,94% для первой и 15,38% для второй).

Разбор подготовил ❣ Валентин Шубин

Душный NLP

ReMax как альтернатива PPO

Сегодняшняя статья — об интересном методе обучения с подкреплением, который называется ReMax. Авторы предлагают его как замену популярному алгоритму Proximal Policy Optimization (PPO).

ReMax основывается на алгоритме обучения с подкреплением, который называется REINFORCE — отсюда и приставка Re. В REINFORCE, в отличие от PPO, нет value-модели. Она выступает в роли бейзлайна для снижения дисперсии оценки градиента и представляет собой среднюю награду, которую наберёт генератор, если будет отвечать на конкретный запрос.

Вместо отдельной value-модели в ReMax предлагают использовать другой бейзлайн — то, сколько награды набирает greedy-генерация обучаемой моделью на запросе. Отсюда окончание Max. Такой бейзлайн тоже отлично подходит с точки зрения теории, и не требует хранения дополнительной модели в памяти — авторы сообщают, что снижение потребления GPU-памяти составляет 46%. А поскольку число моделей, необходимых алгоритму алайнмента, уменьшилось, то уменьшилось и число гиперпараметров, которые нужно подобрать.

Использование ReMax для обучения модели Mistral-7B показало значительные улучшения. Модель достигла 94,78% успеха на leaderboard AlpacaEval и установила новый стандарт для моделей с 7 миллиардами параметров. Эти результаты демонстрируют, что ReMax может стать отличной альтернативой PPO для RLHF-задач, значительно сокращая вычислительные затраты и повышая эффективность обучения крупных языковых моделей.

Разбор подготовил ❣ Павел Темирчев

Душный NLP

О претрейн-датасете LLaMA 3.1

Сегодня расскажем о том, как устроен претрейн-датасет для LLaMA 3.1. Разработчики этой LLM поделились множеством подробностей, поэтому будет интересно.

Одна из важнейших характеристик LLaMA 3.1 — объём данных, на которых она обучалась. Для этой модели использовался претрейн-датасет из 15 триллионов токенов. Это один из самых больших наборов данных для языковых моделей, превосходящий, например, LLaMA 2 с её 2 триллионами токенов. Модель также включает специальный набор данных для длинных контекстов.

Сбор данных для LLaMA 3.1 построен на стандартном пайплайне для обработки текста из HTML. Один из его ключевых аспектов — это фильтрация данных на ранних стадиях. Она включает как дедупликацию данных, так и использование эвристических методов для удаления нежелательных токенов.

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

LLaMA 3.1 демонстрирует хорошие результаты благодаря специальным пайплайнам для математических задач и задач, связанных с программированием. Объем данных по коду и математике в датасете LLaMA 3.1 превосходит весь датасет оригинальной LLaMA.

Для оценки качества данных использовались различные классификаторы, полученные дистилляцией LLaMA 2. Другие классификаторы отбирали документы, опираясь на вероятность их принадлежности к определённому домену. Один из интересных моментов — использование классификаторов для оценки доменов с целью дальнейшего перевзвешивания различных частей датасета. Например, данные из области искусства и развлечений могут быть уменьшены в весе по сравнению с математикой.

Фильтрация включает несколько уровней дедупликации. Во-первых, стандартная дедупликация с использованием метода MinHash. Есть также локальная дедупликация — метод, при котором данные разбиваются на блоки, и дубликаты удаляются только внутри этих блоков.

Ещё одной важной частью фильтрации данных стала проверка на безопасность. Разработчики LLaMA 3.1 реализовали специальные фильтры для удаления вредоносного или небезопасного контента. Эти фильтры основаны на классификаторах и могут удалять не только нежелательные токены, но и данные с персональной информацией.

LLaMA 3.1 также выделяется своими экспериментами в области Data Mixer — процесса, в котором данные разделяются по доменам, а затем перевзвешиваются для оптимальной работы модели.

Важным этапом оптимизации стал так называемое «сведение». Это метод, при котором на последних этапах обучения модель дообучается на малом наборе данных, включающем специфичные задачи. Это позволяет значительно повысить качество ответов модели на конечных тестах. В ходе экспериментов установили, что даже небольшой объём высококачественных данных может существенно улучшить результаты.

Разбор подготовил ❣ Дмитрий Мокеев

Душный NLP

Архитектура LLaMA 3.1

Продолжаем разбирать LLaMA 3.1. В прошлый раз речь шла о претрейн-датасете, а в этот раз — об архитектуре модели.

Llama 3 использует стандартную архитектуру трансформера, которая не сильно отличается от того, что было в LLaMA и LLaMA 2. Однако отличия есть. Скажем, если в LLaMA 2 Grouped Query Attention (GQA) с восемью ключевыми головами внимания использовались только в моделях на 34B+, то здесь GQA применяется для всех моделей LLaMA 3.1. Это позволило повысить скорость вывода и уменьшить объём данных, необходимых для кеширования во время декодирования.

Ещё одно важное изменение — увеличение контекстного окна до 128 тысяч токенов. Это стало возможным благодаря увеличению гиперпараметра базовой частоты RoPE до 500 тысяч. Такой подход позволяет модели эффективно решать задачи, связанные с большими объёмами текстов. Модель также использует словарь на 128 тысяч токенов.

Разработчики внедрили четырёхмерный параллелизм (4D Parallelism), который включает тензорный, пайплайновый, контекстный и параллелизм данных. Этот подход позволяет значительно улучшить утилизацию ресурсов при обучении на тысячах GPU. Например, для обучения модели с 405 миллиардами параметров использовалось до 16 тысяч GPU, а средняя утилизация вычислительных ресурсов составила около 41%​.

Контекстный параллелизм позволяет разбивать длинные строки на части. В отличие от классических методов, такой параллелизм синхронизирует только ключи и значения в attention-слое, что минимизирует задержки при обработке длинных последовательностей.

Помимо этого, в архитектуре LLaMA 3.1 активно используется FP8-квантизация, которая значительно ускоряет вычисления без значительных потерь в точности. Это позволяет экономить до 50% времени на вычисления по сравнению с традиционными методами, что критично для моделей с миллиардами параметров. FP8-квантизация используется не для всех слоев, потому что она может вызвать ошибки при вычислении в attention-слоях. Подход доказал свою эффективность при решении большинства задач​.

Разбор подготовил ❣ Михаил Хрущев

Душный NLP

Ограничения Instruction Tuning и как их преодолеть

Supervised Full Fine-tuning (SFT) — распространённая практика, но он не лишён недостатков. Авторы сегодняшней статьи задаются вопросом: а может ли LoRA (Low-Rank Adaptation) исправить недочёты?

При использовании Full Fine-tuning возникает две проблемы: у моделей часто возникают сложности с извлечением новых знаний из SFT-датасета, могут участиться галлюцинации. Исследование показало, что модели, обученные с использованием Full Fine-tuning, могут генерировать неверные ответы, если берут слишком много токенов из SFT-датасетов. Эффект особенно заметен, если модель пытается отвечать на вопросы, требующие глубокой экспертизы.

Например, на вопрос «Какие основные работы Эйнштейн сделал после того, как в 1915 году открыл Общую теорию относительности?» модель начинала выдавать не соответствующие действительности ответы — скажем, о «квантовой теории атома трития».

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

Первые несколько процентов токенов, сгенерированных LoRA-моделью, могут быть изменены (по сравнению с ответом предобученной модели), чтобы правильно начать ответ. Но большая часть предложения остаётся такой же, как у предобученной модели. Это позволяет уменьшить количество галлюцинаций. Эксперименты показали, что LoRA даёт более точные ответы.

LoRA эффективен даже при малом объёме датасета. Например, модель с LoRA, обученная на наборе данных из 1000 инструкций, может превосходить модели с SFT на датасетах по срезам фактологичености и полезности, содержащих 52 000 или даже 326 000 инструкций. В экспериментах использовались различные открытые и домен-специфичные датасеты, включая MedInstruct и Alpaca. Модели с LoRA демонстрировали лучшее соответствие фактам и были менее подвержены галлюцинациям.

Разбор подготовил ❣ Алексей Шимко

Душный NLP

Алаймент LlaMA 3.1

Возвращаемся к LlaMA 3.1 и продолжаем разбираться, как она устроена. В этот раз речь пойдёт об алайменте модели.

По сравнению с LLaMA 2 у третьей версии изменилась разметка пар. Помимо стандартных chosen и rejected добавилась ещё метка edited. Она ставится в тех случаях, когда победивший объект не слишком хорош и его переписывают. Ответы оцениваются по семибалльной шкале.

SFT происходит в шесть раундов. Если в LLaMA 2 использовался PPO, то в LlaMA 3 — DPO. Разработчики отмечают, что это связано с тем, что PPO требует больше вычислительных ресурсов, а качество выходит хуже.

Ещё одно важное отличие — это специализация. На претрейне модель доучивают для решения специальных задач. Потом делают отдельный алаймент, полученную специализированную модель используют для генерации новых обучающих данных стадии алайнмента, а также мержат веса нескольких специализированных модель в единую модель.

Reward-модель обучается над претрейном. Margin term, который был в Llama 2, в третьей версии отсутствует, так как, по словам разработчиков, он не даёт никакого прироста в качестве. Как и в DPO, оставляют только те ответы, которые помечены как «сильно лучше» и «лучше». Кроме того, в reward-модели есть отдельные награды для полезности и безопасности.

За один раунд SFT 405B-модель суммарно проходит 576 тысяч сэмплов. В DPO используют сэмплы от моделей с последних раундов (а в reward-модели — все). Служебные токены, такие как EOS или токены для вызовов функций, маскируют для стабилизации обучения. Кроме того, к DPO добавляют NLL (CE) с коэффициентом 0,2. Это нужно, чтобы повысить вероятность chosen-ответов. Промпты для обучения пишут люди, а ответы — модели. На один промпт выбирают лучший ответ из 10-30 поступивших.

В LlaMA 3.1 есть четыре уровня фильтрации данных:

1. Rule-based — удаляет дисклеймеры, смайлики и восклицания;

2. Quality — качественными считаются 25% ответов с наибольшим скором. Кроме того, используется LLM-as-judge. Ответы оцениваются по трём критериям для обычных данных и двум — для кодинга. Ответ считается качественным, если все критерии выполнены. Сэмпл попадает в обучение, если хотя бы один из методов показал, что ответ качественный;

3. Difficulty — оценивается по числу интентов в запросе: чем их больше, тем сложнее запрос. Также модель оценивает сложность по трёхбальной шкале;

4. SemDedup — этот метод используется для удаления похожих данных, при отборе отдается предпочтение семплам с максимальным значением quality * difficulty.

Алаймент для каждой из функциональных возможностей (Capabilities) LLaMA 3.1 имеет свои особенности. Например, в коде есть много синтетических данных, используется execution feedback и перевод на редкие языки программирования. А для математики берут тексты из претрейна и уже к ним генерируют инстракты.

Что касается фактологичности, то разработчики не добавляют новых знаний поверх претрейна. Для этого модель обучают отвечать только на те вопросы, на которые она может выдать ответ, согласованный с документом из претрейна. А для чувствительных тем, по которым в датасете много некорректной информации, используют ручную разметку.

Разбор подготовил ❣ Алексей Зотов

Душный NLP

TDPO — потокенный DPO или просто регуляризация?

Авторы сегодняшней статьи предлагают метод потокенного Direct Preference Optimization (DPO), который на бумаге должен исправить некоторые проблемы оффлайн-обучения с подкреплением. Но на деле все оказывается не так просто.

DPO — метод обучения, не полагающийся на reward-модель. Здесь применяют датасет с размеченными парами запросов и ответов, чтобы натренировать генератор на контрастный лосс.

Проблема в том, что в случае с DPO мы работаем с вероятностями последовательностей целиком. Метод ограниченно контролирует поведение модели на уровне отдельных токенов. Это приводит к тому, что модель может ошибочно сильно повышать или понижать вероятность отдельных токенов значительно после совершенных ошибок.

Эту проблему можно нивелировать, если сделать DPO потокенным. Авторы статьи пытаются добиться этого.

Для начала они предлагают ввести необычное ограничение — сделать так, чтобы сумма наград всех токенов-продолжений для произвольного префикса была равна 0. Это довольно сильное допущение: например, если мы решаем задачу копирования какого-то куска текста, то будем сильно штрафовать модель за любое отклонение. Как результат — награда за правильный токен окажется очень большой. В этом случае, если при выборе между длинной и короткой строкой, модель будет склоняться к длинной строке.

Такое ограничение позволило авторам в их расчётах лосса избавиться от нормировочной константы вероятностного распределения. Чтобы ее вычислить, нужно суммировать награду по всем возможным ответам, а это невозможно, поэтому от константы при расчётах избавляются. В DPO нормировочная константа одинакова для победившего и проигравшего ответов, поэтому она сокращается в лоссе, но авторы статьи сделали это несколько иначе.

Из их математической модели выводится функция, которая очень похожа на DPO. Но в отличие от DPO, авторы вычитают из неё разницу между SeqKL проигравшего и победившего ответа. Этот метод, названный Token-level Direct Preference Optimization (TDPO), обеспечил незначительное улучшение по сравнению с обычным DPO. На датасете Anthropic HH точность увеличилась всего на 0,65%.

Далее авторы предлагают умножить на дополнительный коэффициент разницу SeqKL и не пропускать градиенты для победившего варианта. Это можно трактовать так: при росте SeqKL проигравшего ответа всегда увеличивается лосс, в то время, как при росте SeqKL победившего — лосс уменьшается. Получается, что добавка к DPO, после остановки градиента для её части, по сути работает, как регуляризация.

С ней метод получил название TDPO2 и он действительно неплохо улучшает показатели. На том же Anthropic HH прирост по сравнению с DPO составил уже не 0,65%, а 7,9%.

Авторы действительно предложили лучшее решение. Но возникает вопрос: насколько здесь велик вклад выведенной математической модели. По факту, авторы сильно меняют основные моменты в этой модели, а то, что остается, очень похоже на простую потокенную регуляризацию. Но её идея не нова: часто к DPO добавляют negative log likelihood loss — например, при DPO-обучении Llama 3.1, — что тоже является вариантом потокенной регуляризации. Мы склоняемся к тому, что научный вклад этой статьи невелик, а ключевые выводы — ошибочны.

Разбор подготовил ❣ Михаил Хрущев

Душный NLP

WARM — метод улучшения reward-моделей

Сегодняшняя статья — о методе усреднения весов reward-модели для устранения проблем, связанных с RL-обучением. Но для начала напомним, как работает Reward-модель.

На вход она принимает промпты и ответы, а на выход выдаёт скаляры. По ним возможно ранжировать ответы от лучшего к худшему. Всё это делается с помощью обучения на минимизацию лосса, который вытекает из модели Брэдли-Терри. Как правило, reward-модели обучаются на датасете из преференсных данных — то есть таких, в которых ответы уже размечены асессорами или другой моделью.

Есть ряд проблем, с которыми можно столкнуться во время обучения reward-модели. Во-первых, разметка может оказаться достаточно шумной — например, при расхождениях в оценках одного и того же ответа разными асессорами. Кроме того, в некоторых случаях политика может генерировать OOD-ответы для выбранной reward-модели.

Наконец, возможно и такое, что reward-модель выучится на какой-то черте данных — например, особенностях оформления. При этом на файнтюнинге модель научится генерировать те ответы, которые будут давать высокий скор именно из-за этой особенности, а не из-за качества самих ответов. Скажем, будет отдавать приоритет хорошо оформленным, а не правильным ответам.

Существует несколько методов, призванных справится с вышеописанными проблемами. Например, можно обучить много абсолютно разных reward-моделей и усреднить их логиты. Этот метод называется prediction ensembling (ENS), а его главный недостаток заключается в необходимости инферить сразу несколько моделей, что не очень экономично в условиях файнтюнинга.

Авторы статьи, в свою очередь, предлагают обучать reward-модель с помощью одного датасета с преференсными данными, но с разными гиперпараметрами, а также с разных чекпоинтов SFT-обучения. В результате получается несколько моделей с одинаковой архитектурой. Их веса следует усреднить в одну модель — Weight Average Reward-Model (WARM), которая поступает как reward-функция в RL. Проведенный авторами анализ показал, что WARM — это аппроксимация ENS.

Почему это должно работать? Известно, что существует линейная связь в моделях, обученных из одного претрейна. Она позволяет усреднять веса, не теряя при этом в качестве. Однако это справедливо только для одного претрейна.

Проверки c использованием датасета TL;DR summarization показали, что WARM запоминает меньше испорченных или некорректных данных разметки в датасете, чем ENS. То же самое касается работы с OOD-примернами. Однако на «чистом» фрагменте датасета, где разметка без ошибок, ENS выдаёт лучшие результаты.

Авторы заявляют, что преимущество их метода заключается в использовании всего одной модели в ходе файнтюнинга — это позволяет экономить время и вычислительные ресурсы. Кроме того, WARM решает некоторые проблемы, связанные с «грязными» данными. Однако есть и ограничения. Например, необходимость обучаться из одного претрейна и невозможность использовать разные архитектуры.

Разбор подготовил ❣ Илья Черемушкин

Душный NLP

ReDrafter — быстрый метод спекулятивного декодирования

Сегодняшняя статья посвящена Recurrent Drafter (ReDrafter) — новому подходу к спекулятивному декодированию. Авторы заявляют, что он позволяет значительно увеличить скорость моделей.

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

В ReDrafter в качестве черновой модели используется RNN. Благодаря рекуррентной структуре, учитывается зависимость между черновыми токенами. Кроме того, RNN, помимо последнего токена, также видит и скрытое состояние из последнего слоя LLM, что даёт ей возможность лучше следовать генерациям основной модели.

С помощью алгоритма динамического древа внимания (dynamic tree attention algorithm), в сгенерированных RNN кандидатах убираются повторяющиеся префиксы. Таким образом, в основную модель попадает меньше вариантов, а значит затраты вычислительных ресурсов становятся меньше. Далее основная модель оценивает предложенные варианты (при этом, сразу несколько), выбирает лучший и процесс повторяется снова.

RNN обучается с помощью дистилляции из основной модели. Это позволяет RNN предсказывать токены с вероятностями, максимально приближенными к ожиданиям LLM. Таким образом, «черновая» модель реже предлагает токены, которые будут отклонены.

Авторы отмечают, что использование дистилляции дает лучшие результаты, чем обучение «черновой» модели на исходном датасете, на котором тренировали LLM. Так, скорость генерации и число принятых токенов за один шаг на Vicuna 7B выросли примерно на 10%.

ReDrafter показал лучшие результаты по сравнению с методами Medusa и EAGLE на бенчмарках MT-Bench и AlpacaEval. При этом, по сравнению с жадной генерацией, генерация семплированием показывает ещё большее ускорение, чего обычно не показывают другие методы спекулятивного декодирования.

Разбор подготовил ❣ Алексей Гликин

Душный NLP