Смесь рекурсий обеспечивает в 2 раза более быстрый вывод — вот как это реализовать

Исследователи из KAIST AI и Mila представили новую архитектуру Transformer, которая повышает эффективность больших языковых моделей (LLM) с точки зрения памяти и вычислительных ресурсов. Эта архитектура, получившая название «смесь рекурсий» (MoR), значительно повышает точность моделей и обеспечивает более высокую производительность по сравнению с обычными трансформерами, даже при тех же ограничениях на количество параметров и вычислительный бюджет.

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

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

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

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

MoR расширяет этот рекурсивный подход двумя ключевыми компонентами. Первый — это легковесный маршрутизатор, который интеллектуально назначает каждому токену определённую глубину рекурсии. Эта концепция аналогична механизму маршрутизации в моделях «Смесь экспертов» (MoE), где маршрутизатор направляет токены в специализированные экспертные сети. Однако в MoR «эксперты» — это различные глубины рекурсии, что позволяет модели динамически выбирать объём вычислений, применяемых к каждому токену. Она решает, сколько раз следует применять общий блок слоёв, основываясь на сложности токена или его требуемой «глубине мышления». Это направляет вычисления только туда, где они наиболее необходимы, избегая ненужных циклов на легко обрабатываемых частях входных данных.

Вторым компонентом является более эффективная стратегия кэширования пар «ключ-значение» (KV). Кэширование KV — это стандартный метод, позволяющий сохранять информацию из предыдущих токенов для ускорения генерации, но в рекурсивных моделях оно становится узким местом по памяти. MoR представляет механизм кэширования KV «по рекурсии», который выборочно сохраняет и извлекает пары «ключ-значение» только для токенов, оставшихся активными на данном шаге рекурсии. Такое целевое кэширование сокращает трафик памяти и повышает пропускную способность без необходимости сложных модификаций после обучения.

Как утверждают исследователи в своей статье, «По сути, MoR позволяет моделям эффективно корректировать глубину своего мышления на потокенной основе, объединяя эффективность параметров с адаптивными вычислениями».

Чтобы протестировать свою структуру, исследователи обучили модели MoR с использованием от 135 миллионов до 1,7 миллиарда параметров и сравнили их с ванильными и стандартными рекурсивными базовыми моделями по потерям при проверке и показателям точности при небольшом количестве попыток.

Результаты демонстрируют значительный прирост производительности. При равном бюджете вычислений для обучения модель MoR достигла более высокой средней точности при небольшом количестве выстрелов (43,1% против 42,3%), чем базовая модель, несмотря на использование почти на 50% меньше параметров. При обучении на том же объёме данных модель MoR сократила время обучения на 19% и пиковое потребление памяти на 25% по сравнению с базовой моделью.

Архитектура MoR также демонстрирует масштабируемость. Хотя она немного уступала по производительности базовой модели при минимальном масштабе в 135 млн параметров, разрыв быстро сокращался по мере увеличения размера модели. Для моделей с более чем 360 млн параметров MoR достигала производительности стандартных Transformers или превосходила их, особенно при ограниченном вычислительном бюджете. Более того, архитектура MoR значительно повышает производительность вывода. Одна конфигурация MoR достигла ускорения в 2,06 раза по сравнению с базовой моделью. Для компании, работающей в условиях масштабирования, это может привести к значительной экономии операционных расходов.

Сангмин Бэ, соавтор статьи и аспирант KAIST, подробно описал практические результаты в электронном письме VentureBeat. «Хотя точные цифры предоставить сложно, в целом, сокращение размера параметров модели и объёма кэша KV означает, что мы можем выполнять вывод на гораздо большем количестве выборок одновременно», — сказал он. «Это приводит к увеличению количества одновременно обрабатываемых токенов и позволяет обрабатывать более длинные контекстные окна».

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

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

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

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

«Мы очень рады возможности его распространения на мультимодальные сценарии, где повышение эффективности имеет решающее значение», — сказал Бэ.

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