AI[ex]Time

GLoRe: When, Where, and How to Improve LLM Reasoning via Global and Local Refinements Статья про базовые подходы к критике моделей для решения multi-step reasoning задач. Здесь в качестве простого бенчмарка рассматривался GSM8K. Для контекста, как может выглядеть задачка оттуда: Мистер Санчес выяснил, что 40% учеников его 5-го класса получили итоговую оценку ниже B. Сколько учеников получили итоговую оценку B и выше, если у него 60 учеников в 5-м классе? То есть это довольно простые задачи на вычисления, но требующие нескольких операций. Авторы рассматривают 3 подхода: — Outcome Reward Model (ORM). По собранному датасету из множества запусков оцениваем вероятность, что данное решение правильное. В терминах ML это просто кросс-энтропия на 2 класса решено/нет. Это самый простой способ и для своего юзкейса работает хорошо, например, если надо выбрать одно из 5 решений. Собирать данные здесь тоже легко, так как мы заранее знаем правильный ответ. — Process Reward Model (PRM). Учим предсказывать корректность каждого действия. В таком случае на каждом шаге размышления модели мы сможем оценить его правильность. Но в такой постановке нам нужно разметить каждое действие по отдельности, что может быть дорого и медленно, если использовать человеческую разметку. — Stepwise Outcome Reward Model (SORM). Это собственно первая новизна в статье. Для каждого шага учимся предсказывать V-функцию оптимальной стратегии (обозначается V*). Что это значит более простыми словами: V*(S) = 1, если из этой позиции можно решить задачу и V*(S) = 0, если нельзя. Здесь есть нюанс: на самом деле оптимальная стратегия зачастую может решить задачу из любого состояния, для gsm8k уж точно, но разметка вида V*(S) = 1 везде нам полезной не будет. Так как мотивация исходит из того, чтобы далее критиком дополнить основную LLM, нам нужно уметь отличать ее хорошие действия от плохих. Поэтому в качестве аппроксимации V* авторы берут Best-of-K генераций основной модели, то есть таргет для состояния = 1, если хотя бы одно решение из него оказалось правильным, и 0 иначе. Таким образом, мы учим что-то вроде “вероятность, что какое-то решение из этого состояния будет правильным”. Все это очень сильно перекликается с недавними идеями из статьи про Advantage Verifiers, про которую возможно в другой раз. Сравнения с PRM для критики рассуждений нет, но утверждают, что SORM работает лучше ORM, правда только для промежуточных шагов. Для финального ранжирования по всем траекториям ORM выигрывает. Вторая новизна заключается в том, чтобы использовать SORM для прокачки ризонинга основного генератора, про это — во второй части.

На выходных ездил в Париж на хакатон Alan X Mistral, посвященный healthcare, где Nebius выступал одним из спонсоров и давал всем желающим 1xH100 на эксперименты. Выступал в роли судьи и делал доклад. Несмотря на то, что конкретная тема снизила разнообразие идей и решений (все вертелось так или иначе вокруг медицинских помощников), некоторые команды показали интересные подходы. Кто-то запилил голосового помощника, которому можно по видео что-нибудь показать; кто-то столкнулся с тем, что полезные материалы находятся только в книжных версиях, нашел лекции по этим книжкам на youtube, по транскрипциям выцепили полезные куски и на них уже строили SFT. В общем, было круто общаться, знакомиться и обсуждать подходы. Многие спонсоры там же и предлагали пойти к ним на собесы. Очень классная история для студентов, ищущих стажировки и начальные позиции. Рассказывал я про outcome/process supervision в LLM агентах, совсем простыми словами, чтобы ребята могли успеть что-то запихнуть в решения к себе. Так как в последнее время на работе занимаюсь различными видами guidance в агентах, думаю сделать какой-то обзор подходов и перспективных идей, которые за этим стоят, в основном все вокруг offline reinforcement learning. Если у вас есть на примете интересные статьи на тему, накидайте плиз те, про которые вам бы хотелось увидеть пост.

Kaggle соревнование lmsys chatbot arena, часть 2. Технические подходы. В продолжение разбора соревнования обещал написать вторую часть с техническим обзором. Время пришло. Задачу можно было решать двумя вариантами: добавить голову и учить модель на задачу классификации или же оставить предсказание следующего токена и напрямую предсказывать токен-метку. Разницы особой нет, но во втором случае можно использовать много разных фреймворков с оптимизациями обучения и инференса по типу unsloth. Бейзлайн выглядит так: 1. Берем llama3-8B или gemma2-9B 2. Учим лору, вставляя адаптеры во все линейные слои 3. Инференсим квантизованную модель в int4/8 без мерджа весов адаптеров Улучшить решение можно было несколькими способами: 1. Pseudo-labeling. берем какой-нибудь lmsys-1M-dataset, составляем пары ответов на один промпт и размечаем llama3.1_405B. Были попытки и с нуля генерировать синтетические данные, но докидывали они значительно меньше, все-таки распределение данных в таком случае сильно отличается от целевого. 2. External Datasets. Просто докидываем больше данных в post pre-train. Важно, что не в финальный fine-tune, тк на последнем шаге лучше использовать только данные из соревнования. Много интересных датасетов можно найти в RLHFlow. Авторы так же в свое время писали неплохую статью про RLHF. 3. Ensembling. Пришлось пробовать много разных моделей: MistralNemo, Llama3/3.1, Phi, Yi, Qwen, Gemma и тд. Лучше всего заработала gemma2-it, причем с большим отрывом по сравнению с другими моделями. На втором месте Llama3 (интересно, что 3.1 не докидывала). Удивительно, но модели от Mistral вообще не могли справиться с задачей. Если добавить всякие оптимизации во время инференса (dynamic batch size, dataset length sorting), где-то пожертвовать длиной контекста, то можно было уместить на 2xT4 инференс gemma + llama за 9 часов. Gemma работала значительно дольше, в частности, из-за огромного словаря. 4. Inference tricks. Всякие мелкие, но важные детали. Например, если мы используем ансамбль, то в одну модель лучше отправлять question-responseA-responseB, а в другую ответы поменять местами, чтобы добавить больше разнообразия. Важно также выставить truncation left side, чтобы жертвовать токенами из начала — они меньше влияет на предикт модели. Кто-то лез совсем в детали и выключал logit soft-capping в gemma, писали, что докидывает пару тысячных на лб — типичный кегл 😋 Кстати, если я не ошибаюсь, это первое соревнование, в котором завели инференс 33B моделей: vllm + квантизация AWQ + Tensor Parallel. 5. И напоследок прием, который зарешал больше всех — Distillation. Парень с таким подходом и взял первое место. Логика следующая: 1. Бьем весь трейн на 5 фолдов. 2. Тренируем на фолдах Llama3-70B и Qwen2-72B и размечаем весь датасет их предиктами. 3. Опять же на фолдах дистиллируем предикты больших моделей в gemma2, используя самый простой KL loss. Учим только LoRA адаптеры и в итоге получаем 5 моделей. 4. Усредняем веса всех адаптеров и получаем с помощью такого model merging финальную модель. 5. На все про все — А100 80G * 8 + ZeRO2 Часть 1 про лик в соревновании

Хорошие слайды от одного из сотрудников DeepMind на тему LLM Reasoning, RL, Verifiers — в общем всего, что занимает умы многих последние месяцы 😅 Вкратце шаг за шагом проходят по многим популярным методам: — ReST: генерируем синтетику, фильтруем правильные примеры (для этого у задачи должен быть правильный ответ, который нам известен) и докидываем в тренировку. — Добавляем Verifier: во время инференса генерируем множество вариантов и выбираем один с помощью эвристики/модели. - В качестве эвристики отлично работает Majority Voting, то есть берем генерацию с самым частым финальным ответом среди кандидатов. - Модель же можно обучать разными способами: от бинарной классификации на правильный/неправильный ответ до DPO на генерации из ReST. — Verifier дает нам возможность разменивать качество на объем вычислений: общая тенденция такая, что генерируя больше кандидатов и выбирая из них ответ, можно сильно растить метрики. Отличная статья на тему scaling test-time compute: https://arxiv.org/abs/2408.03314v1 — Переходим к Generative Verifier — очень популярный в последнее время метод обучения критиков. Заключается он в том, что модель мы учим так же в режиме next token prediction и можем добавить Chain of Thought, от чего качество становится еще выше. — Еще есть всякие дополнительные топики по типу Cost-Matched Sampling, Reinforced In Context Learning (что, кстати, позволяет делать интересные штуки по типу “дообучаться на лету” под поведение пользователя)

Подоспела запись вебинара от OpenAI и Alistair Pullen (Cofounder & CEO Genie) про файн-тюнинг линейки gpt4o моделей. Напомню, Genie — лаборатория, совсем недавно показавшая лучший результат на SWE-bench verified 43.8% (тщательно отфильтрованная версия оригинального бенчмарка), взяв за основу gpt4o и дотюнив ее с OpenAI под доп задачи. - Всего тюнили на 100М токенах, поддерживают 15 языков программирования. Правда не рассказали, какие конкретно данные были: только траектории из swe-bench или какая-то смесь из разных задач. - Эксперименты ставились в таком формате: сначала proof of concept на маленькой модели (gpt4o-mini), потом переходят на большую. Но для “picker model” (классификатор, который выбирает лучшее действия из набора кандидатов) оставили маленькую модель. - Alistair топит за eval-driven finetuning. Если опираться только на “vibe-checking” — все сойдется к субоптимальному результату. - Пример работы по улучшению тренировочных данных. Была проблема: модель фейлится в решении легкой задачи если у нее не получилось с первого раза. Причиной оказалось то, что в датасете такие задачи всегда решались с первого раза. Решение: наделать синтетики, где после неудачных попыток простая задача все-таки решается (не рассказали, как делали так, чтобы эта синтетика не провоцировала вторую попытку там, где раньше была одна, но обычно это делается простым маскированием части токенов) - Другой пример. Alistair говорит, что в обучении (видимо, имеется в виду весь цикл обучения моделей OpenAI) недостаточно представлена ошибка runtime errors, так что имеет смысл нагенерить для тюнинга сэмплов с ними и примерами, как действовать в таких случаях. - Интересная мысль, про которую упомянул вскользь: во время валидации находить части, где у генерации модели высокий uncertainty и дальше изучать трейн, чтобы понять, почему так произошло. Таким образом итеративно можно закрывать потенциальные пробелы, которые трудно уловить визуально. - Очевидно весь фокус направлен на данные, включая синтетику, а не настройку гиперпараметров. Учили все со значениями, предложенными OpenAI. - После тюнинга обобщающая способность модели падает: модель начинает хуже работать вне того распределения, на котором тюнили.

Хорошая статья с разбором RingAttention - метода вычисления Attention с задействованием всех доступных карт, при этом не сильно упираясь в communication costs. Мотивация растет как всегда из того, что у Attention квадратичная сложность, и на больших контекстах потребление памяти становится огромным. Но даже с учетом оптимизаций FlashAttention (где за счет ухода от хранения полной матрицы QxK сложность памяти снижается до линейной) мы упираемся в проблемы для контекстов в сотни тысяч/миллион и больше. Поэтому хочется использовать несколько карт, чтобы распределить нагрузку на память. Собственно метод предлагает вариант реализации такого алгоритма, а в статье довольно подробно описаны все технические моменты по типу разбиения матриц Q, K, V и агрегация результатов для подсчета коэффициентов.

Неделю назад закончилось соревнование LMSYS - Chatbot Arena Human Preference Predictions на kaggle, где нужно было предсказать модель-победитель из чата пользователя на lmsys chatbot арене (Место, где юзер общается сразу с двумя моделями и оценивает их ответ. Из множества оценок складывается общий рейтинг ELO всех моделей). Получилась эта история очень неоднозначной: два найденных лика, один из которых зарепортили в последние 2 дня (тк некоторые люди хотели использовать его, не сообщив организаторам, и взять первые места), пересбор всего private теста после окончания, необходимость в большом кол-во железа для обучения моделей и как следствие шейкап в конце. Все соревнование держался на 1-5 месте, но на прайвате на новых собранных датасетах откатились аж до 33 😢 У этой серии будет несколько постов (описание лика, технических подходов, выводов), но пока не знаю, сколько именно. Для начала пара слов о задаче: есть чаты пользователей с арены, то есть серия вопросов, на каждый из которых имеется два ответа: от модели A и от модели B. Пообщавшись какое-то время, пользователь может поставить winner model A/B, Tie, Tie (both bad). Этот ответ и нужно было научиться предсказывать, только Tie и Tie (both bad) склеены в одну метку. Упрощая, можно сказать, что у нас задача классификации на 3 класса с финальной метрикой logloss. Итак, часть 1. Лик за 2 дня до конца, который поставил все соревнование под угрозу Если мы зайдем на лидерборд арены, то увидим ссылку на колаб для воспроизведения всех результатов. В этом ноутбуке скачивается файл, содержащий оценки пользователей за всю историю. Но самих чатов (промптов + ответов моделей нет), то есть казалось бы, эти данные не несут опасности. Наверное, поэтому организаторы и не стали задумываться о возможности лика в этом месте, но если посмотреть внимательно, то в колонке conv_metadata лежит информация о кол-ве токенов из промпта и ответов моделей. А токенизатор, который используется для сбора этой статистики, тоже известен. Тогда получается следующая картина: во время теста на кегле можно токенизировать вопросы/ответы и для многих случаев однозначно смапить их со строчкой из файла арены, где финальный ответ уже есть. Из-за этого в один из последних дней, когда у первого места был скор 0.879, люди проснулись и увидели 0.656, потом 0.251, а потом и вовсе ~0.1. На форуме несколько дней стоял хаос, так как до конца оставалось совсем немного, а ответа от организаторов не было. Итоговое решение было такое: продлить соревнование на 2 недели, заморозить прием сабмитов и пересобрать тестовый набор. Интересный сайд эффект, с которым ничего не стали делать: после заморозки соревнования и сбора нового теста, ~10% участников не получили никакого места на лидерборде, тк все их решения упали по таймауту 🤯 В следующем посте расскажу про основные технические подходы.

Agent Q: Advanced Reasoning and Learning for Autonomous AI Agents. Интересная работа, объединяющая известные методы улучшения моделей для решения сложных агентских задач. Здесь и про SFT (STaR), и MCTS во время инференса, и DPO на траекториях, и AI feedback guidance: в общем хороший обзор популярного сейчас направления с агентами. Любопытным показался метод сбора данных для обучения DPO, но уже не на уровне траекторий, а на уровне каждого шага. Давайте про этот момент подробнее. В классическом DPO у нас есть промпт x и пара y_chosen, y_rejected, для которых мы и считаем функцию ошибки. Для агентских сред все сложнее — допустим, мы решаем задачу: забронировать столик в ресторане Bougainville на 30 августа на 2 человек на 19:00. Для этого агенту нужно сделать целый ряд действий, а в итоге мы знаем лишь одно: получилось выполнить задачу или нет. И если в каком-то состоянии у агента есть n действий на выбор, то возникает вопрос, а как их сравнивать между собой, чтобы сгенерировать датасет попарок? К этой проблеме можно подойти по-разному: тренировать RM, использовать другую LLM в качестве критика или вообще не сравнивать между собой действия, а сразу целые траектории (это и есть trajectory-level DPO эксперимент из статьи). Здесь же авторы предложили следующее: давайте применим MCTS (Monte Carlo Tree Seach) для того чтобы для состояния и выбранного действия уметь оценивать Q(s, a) функцию. Q функция говорит о том, какую среднюю награду мы можем получить из состояния s, сделав действие а. За счет большого кол-ва симуляций в алгоритме (то есть попыток решить задачу из различных состояний) мы как раз получаем такого рода оценки. Далее в любом состоянии мы выбираем такие a_chosen, a_rejected, что разность их Q функций отличается на определенное значение. Это означает, что действие лучше не потому, что так оценил человек/модель, а потому, что с этим состоянием чаще удалось решить задачу. Отдельно радует, что таким образом можно собирать примеры и для довольно сложных задач, достаточно, чтобы модель умела хоть изредка решать их, иначе все оценки Q функций будут нулевые. На одном из бенчмарков такой тюнинг давал существенный прирост поверх других методов. Далее еще на инференс добавить MCTS и получаем SOTA 🙂