Как исследователи нарушают привычные подходы в ИИ, исключая матричное умножение

Привет, на связи Елена Кузнецова, специалист по автоматизации Шерпа Роботикс. Сегодня я перевела для вас статью, посвященную использованию моделей ИИ без лишней математики. Все мы знаем, что нейросети — дело энергозатратное. А описанные в статье исследования могут помочь снизить энергопотребление в работе нейросетей.

Исследователи из Университета Калифорнии в Санта-Круз, UC Дэвис, LuxiTech и Университета Сучжоу объявили о разработке нового подхода к оптимизации работы языковых моделей ИИ, который исключает матричное умножение из вычислительных процессов. Это принципиально меняет операции нейронных сетей, которые в настоящее время ускоряются с помощью графических процессоров (GPU). Результаты работы, изложенные в недавней препринте, могут существенно повлиять на экологическую устойчивость и операционные расходы систем ИИ.

Матричное умножение, или MatMul, является ключевым элементом большинства вычислительных задач нейронных сетей. GPU особенно эффективны в выполнении этих операций благодаря своей способности выполнять множество умножений параллельно. Эта способность временно сделала Nvidia самой ценной компанией в мире на прошлой неделе: сейчас она занимает около 98% рынка GPU для дата-центров, которые широко используются для запуска таких ИИ-систем, как ChatGPT и Google Gemini.

В новой работе, озаглавленной «Масштабируемое моделирование языка без MatMul», исследователи описывают создание специализированной модели с 2.7 миллиарда параметров, которая демонстрирует результаты, сопоставимые с традиционными крупными языковыми моделями, но без использования MatMul. Они также показывают возможность обработки 1.3 миллиарда параметров на скорости 23.8 токенов в секунду на GPU, ускоренном с помощью программируемого FPGA-чипа, который потребляет всего 13 ватт (без учета энергозатрат самого GPU). Это открывает путь к разработке более эффективных и адаптивных архитектур, утверждают авторы исследования.

Хотя метод еще не прошел рецензирование, исследователи — Руй-Цзе Чжу, Юй Чжан, Итэн Сифферман, Тайлер Шивз, Ицзяо Ван, Дастин Ричмонд, Пэн Чжоу и Джейсон Эшрагян — утверждают, что их работа бросает вызов сложившемуся мнению о том, что матричные операции умножения являются незаменимыми для создания высокоэффективных языковых моделей. Они настаивают на том, что их подход может сделать крупные языковые модели более доступными, эффективными и устойчивыми, особенно для использования на аппаратуре с ограниченными ресурсами, такой как смартфоны.

Избавление от матричной математики

В недавней статье исследователи ссылаются на BitNet — технику «1-битного» трансформера, ставшую популярной в октябре как препринт, как на важный предшественник своего труда. По словам авторов, BitNet продемонстрировал возможность использования бинарных и тернарных весов в языковых моделях, успешно масштабируясь до 3 миллиардов параметров и при этом сохраняя конкурентоспособные показатели.

Тем не менее, авторы отмечают, что BitNet все же полагался на матричные умножения в своем механизме самовнимания. Ограничения BitNet стали мотивом для настоящего исследования, побудив авторов разработать совершенно «свободную от MatMul» архитектуру, которая могла бы поддерживать производительность, исключая матричные операции даже в механизме внимания.

Исследователи предложили два основных новшества: во-первых, они создали кастомную языковую модель, ограничив ее использование только тернарными значениями (-1, 0, 1) вместо традиционных чисел с плавающей запятой, что позволяет упростить вычисления. Во-вторых, они перепроектировали ресурсозатратный механизм самовнимания традиционных языковых моделей, заменив его на более простой и эффективный элемент, названный ими «MatMul-free Linear Gated Recurrent Unit» (MLGRU), который обрабатывает слова последовательно, используя базовые арифметические операции вместо матричных умножений.

Третье нововведение — адаптация управляемого линейного модуля  (GLU) — механизма управления потоком информации в нейронных сетях, для использования тернарных весов при смешивании каналов. Под смешиванием каналов понимается процесс комбинирования и преобразования различных аспектов или характеристик данных, с которыми работает ИИ, аналогично тому, как диджей смешивает разные аудиоканалы для создания целостной композиции.

Эти изменения в сочетании с кастомной аппаратной реализацией для ускорения тернарных операций с помощью FPGA-чипа позволили исследователям достичь, по их утверждению, производительности, сопоставимой с современными моделями, при этом значительно снизив потребление энергии. Хотя они проводили сравнения на GPU для оценки своих моделей по сравнению с традиционными, модели, свободные от MatMul, оптимизированы для работы на оборудовании, подходящем для более простых арифметических операций, таких как FPGA. Это предполагает, что данные модели могут быть эффективно запущены на различных типах аппаратуры, включая те, что имеют более ограниченные вычислительные ресурсы по сравнению с GPU.

Для оценки своего подхода исследователи сравнили свою MatMul-free модель с воспроизведенной моделью Llama-2 (которая была названа «Transformer++») по трем размерам моделей: 370M, 1.3B и 2.7B параметров. Все модели были предобучены на наборе данных SlimPajama, причем более крупные модели обучались на 100 миллиардах токенов каждая. Исследователи утверждают, что MatMul-free модель показала конкурентоспособные результаты по сравнению с базовой моделью Llama 2 на нескольких бенчмарках, включая ответы на вопросы, общее рассуждение и физическое понимание.

Помимо снижения энергопотребления, MatMul-free модель значительно уменьшила использование памяти. Их оптимизированная реализация на GPU снизила потребление памяти до 61% во время обучения по сравнению с неоптимизированной базовой моделью.

Стоит отметить, что модель Llama-2 с 2.7 миллиарда параметров находится далеко от текущих лучших LLM на рынке, таких как GPT-4, которая, по оценкам, содержит более 1 триллиона параметров. GPT-3 вышла в 2020 году с 175 миллиардами параметров. Число параметров, как правило, означает большую сложность (и, грубо говоря, способность) модели, но в то же время исследователи находят способы достигать более высокого уровня производительности LLM с меньшим количеством параметров.

Таким образом, мы еще не достигли уровня обработки, сравнимого с ChatGPT, но техника из UC Santa Cruz не исключает возможность достижения такого уровня производительности при наличии дополнительных ресурсов.

Экстраполяция в будущее

Исследователи утверждают, что законы масштабирования, наблюдаемые в их экспериментах, предполагают, что языковая модель, свободная от MatMul, может превзойти традиционные LLM на очень больших масштабах. Авторы прогнозируют, что их подход теоретически может пересечься с производительностью стандартных LLM и даже превзойти её при достижении масштабов около 1023 FLOPS, что приблизительно соответствует вычислительным ресурсам, необходимым для обучения таких моделей, как Llama-3 от Meta с 8 миллиардами параметров или Llama-2 с 70 миллиардами.

Тем не менее, авторы отмечают, что их работа имеет ограничения. Модель MatMul-free еще не была протестирована на экстрабольших моделях (например, с более чем 100 миллиардами параметров) из-за вычислительных ограничений. Они призывают организации с большими ресурсами инвестировать в масштабирование и дальнейшую разработку этого легковесного подхода к языковому моделированию.