Всё смешалось — эксперты, модели. Как уменьшить LLM, но не проиграть в продуктивности
Накатившая волна LLM с правилом «чем больше, тем лучше» уже вызывала
естественный запрос — хочется повысить эффективность не за счет увеличения
размера модели. Похоже, что вторая волна будет (или уже есть) именно такая —
наращивание размеров уже не так впечатляет, появляется всё больше хитрых
подходов, которые делают ставку не на размер. Мы разберем некоторые из них.
Кто мы? Николай Михайловский, основатель и директор компании НТР, которая
разрабатывает информационные системы для крупнейших российских промышленных
компаний, и Наталья Дерюгина, научный коммуникатор и математик. Ведём канал об
искусственном интеллекте, машинном обучении и других задачах и проектах, которые
покажутся нам интересными, а теперь еще и блог на хабре.
MoE
Возможно, самый громкий из недавних подходов — смесь экспертов (mixture of experts,
MoE). Большую популярность приобрел Mixtral 8×7B, вышедший под конец 2023 года,
но истоки этой идеи уходят к работе Хинтона и коллег 1991 года. MoE архитектура
предполагает, что исходная задача разбивается на подзадачи и каждая передается
одному или нескольким «экспертам» (специализированным сетям). Там, где у
обычного трансформера Feed Forward слой, у MoE — несколько экспертов. Каждый из
них — это FFN (у Mixtral 8×7B) или что-то более сложное в общем случае (например,
сам эксперт тоже может быть MoE). Второй компонент любой MoE — это
маршрутизатор, который определяет какой токен отправится к каким экспертам.
Token Choice Routing. The routing algorithm picks the top-1 or top-2 experts with highest affinity scores for each token. The affinity scores can be trained together with model parameters.
картинка отсюда
Сам маршрутизатор состоит из обучаемых параметров. Вообще, в MoE он обучается одновременно с остальной сетью, но есть интересный вариант Beyonder, который в два раза меньше Mixtral, но дает сравнимый результат. Там авторы взяли четыре готовых эксперта (Openchat 3.5, CodeNinja, PiVoT-0.1-Starling-LM-RP и WizardMath) и дополнительно обучили только маршрутизатор. Получается, что можно не обучать одновременно всех экспертов и маршрутизаторов, а просто собрать несколько готовых моделей с готовыми специализациями, хотя изначально MoE были не об этом. Вопрос остается в том, насколько совместимы между собой должны быть эти заранее обученные эксперты.
Вернемся к классическим MoE. В самом базовом варианте каждый эксперт получает вход и дает выход, но понятно, что это не поможет ускорить работу. Поэтому нужно отсекать ненужных экспертов заранее (эта идея возникла не у Mistral, а раньше). Это conditional computation, когда отдельные части модели активируются в зависимости от входа. Таким образом плотную FFN превращают в разреженную сеть, скорость работы которой зависит от количества выбранных экспертов для каждого токена. Почему бы тогда не отправлять каждый токен единственному эксперту? В таком случае маршрутизатор не сможет научиться направлять к разным экспертам, так что выбранных экспертов должно быть минимум два.
У Mixtral 8×7B 8 экспертов по 7 миллиардов параметров, но авторы спешат оговориться, что это не означает 56B параметров вкупе. Отдельными экспертами считаются только FFN, а еще ведь есть общие слои — часть параметров относится к ним. для запуска Mixtral 8×7B понадобится тот же объем памяти, который нужен для плотной модели на около 45B параметров. В микстрале для каждого токена выбирается два лучших по весам эксперта. Поэтому по скорости работы микстраль аналогичен модели из 12В параметров (не 14, а 12, потому что 2В уходят на общие слои).
Работает модель не хуже GPT-3.5 и LLaMA 2 70B, а иногда даже и лучше. Mixtral доступна на huggingface, попробовать ее из браузера можно в perplexity labs (на русском тоже).
Mixtral сильно экономит вычисления, но всё ещё требует большого объема оперативной памяти. А тренд 2024 года требует, чтобы модели становились всё меньше и меньше. Авторы из Физтеха оптимизировали Mixtral 8×7B так, чтобы его можно было запустить на T4, RTX 3060 и RTX 3080 Mobile. Они обнаружили у микстраля две особенности, которые и позволили разгрузить модель. Во-первых, некоторые эксперты используются повторно для последовательных токенов. Во-вторых, скрытые состояния ранних уровней уже содержат в себе информацию о том, какие эксперты будут в итоге выбраны.
Исходя из этого авторы разгрузили Mixtral в два приема: использовали LRU кэш, чтобы сократить обмен данными между процессором и оперативкой (экспертов, которые используют несколько раз подряд, можно оставлять в памяти для следующего токена) и заранее прогнозировали экспертов, чтобы ускорить их загрузку.
SOLAR
Через две недели после Mixtral-8×7B компания Upstage предложили новый способ апскейлинга для LLM. Он называется глубинное масштабирование, depth up-scaling, DUS, а сама модель — SOLAR 10.7B. В отличие от MoE архитектуры, DUS (вдохновленный при этом Mistral 7B) не так сильно меняет структуру исходного трансформера. Подход значительно более прямой и простой.
Предлагается взять любой трансформер (авторы использовали 32-слойную Llama 2 с предобученными весами Mistral 7B) c n слоями. Когда-то моя школьная учительница по математике говорила: «n мало — бери m». Авторы буквально так и делают: копируют исходную модель с n слоями и накладывают дубли друг на друга, при этом убирая m последних слоев первой модели и m первых слоев второй. Получается двухслойный пирог из 2*(n-m) слоев. Чему будет равняться m, сколько слоев из серединки этого пирога выкинуть — зависит от ограничений на железо. Авторы выкидывали по 8 слоев и в итоге получили 2*(32–8) = 48 слоев.
В таком виде производительность модели конечно хуже, чем у изначальной. Поэтому затем модель дообучают и оказывается что, эффективность быстро восстанавливается. Успех авторы объясняют как раз этим трюком с выкидыванием середины. Можно было бы ведь и просто сложить два трансформера друг на друга, но в таком случае после, например, 32 слоя шел бы первый. А в случае с SOLAR после 24 слоя идет 9. Расхождение на «шве» становится меньше, поэтому дообучение быстрее восстанавливает результативность.
Пока SOLAR не получила такого же внимания как Mixtral, хотя по некоторым бенчмаркам SOLAR превзошла Mixtral 8×7B в донастроенной версии.
CALM
Все громкие новинки стремятся обойти по бенчмаркам предыдущие (а в идеале конечно GPT 3.5), но в определенном смысле это не всегда корректное сравнение, потому что им всегда требуется дообучение в том или ином виде.
Совсем свежая новинка от Google — CALM (Composition to Augment Language Models). Целью авторов этого подхода было научиться совмещать LLM так, чтобы улучшить производительность, не трогая исходные веса. Исходные данные такие: есть готовые опорная модель и дополнительная к ней; изменять веса ни той, ни другой нельзя; есть небольшой набор данных, отражающих комбинированный скилл двух моделей.
Почему не просто сложить две модели с каким-либо гейтинговым механизмом или маршрутизатором, как в МоЕ? Такой вариант не подойдет, утверждают авторы, если по отдельности модели не дают нужного результата. К тому же, непонятно как совмещать модели разного размера и обученных на разных объектах.
Фреймворк CALM дает возможность комбинировать совершенно разные модели, чтобы достигать результатов, недоступных для них по отдельности. Нам доступны веса (но менять их нельзя) обеих моделей и промежуточные представления, но нет доступа к обучающим данным и гиперпараметрам. Чтобы их совместить, авторы вводят два дополнительных набора данных, которые и составляют всю суть подхода. Первый — линейные преобразования, которые отображают представления из дополнительной модели на размерность представлений опорной модели. Это отображение позволяет реализовать динамическую связь между моделями, несмотря на то, что переучивать веса исходных моделей нельзя. Второй набор — слои кросс-внимания. Это и есть сама динамическая связь.
Чтобы обучить этот новый набор параметров используется тот самый комбинированный датасет. Например, дополнительная модель (ma на картинке) обучена на множестве отображений строка-целое число, а базовая модель (mb) — это LLM, которая может совершать арифметические действия. Снежинка у обеих означает, что модели «заморожены», их не изменяют. Задача для объединенной модели — решить уравнения в заданных обозначениях. Так что комбинированный датасет будет состоять из арифметических действий над набором обозначений. Подход классической тонкой настройки потребовал бы полного набора значений и обозначений. В данном же случае авторы показали, что обучать можно на небольшом количестве обозначений, этого достаточно, чтобы произошло обобщение на всё множество.
Больше обзоров AI-статей на канале Pro AI.
