Гимба: Архитектура с гибридными головами для малых языковых моделей

В этой статье мы рассмотрим новую архитектуру малых языковых моделей (LLM), названную Гимба (Hymba), которая представляет собой семейство моделей с гибридными головами, объединяющими механизмы внимания трансформеров и модели состояния пространства (SSM). Эта архитектура разработана для улучшения эффективности обработки языка, сочетая в себе высокое разрешение воспроизведения внимания и эффективное обобщение контекста SSM.

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

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

Гибридная архитектура Гимба

Фьюжн гибридных модулей

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

Эффективность и оптимизация KV кэша

Головы внимания улучшают производительность задач, но увеличивают требования к KV кэшу и снижают пропускную способность. Для решения этой проблемы мы оптимизируем гибридный модуль, комбинируя локальное и глобальное внимание и используя обмен KV кэшем между слоями, что улучшает пропускную способность на 3× и уменьшает кэш почти в 4 раза.

Метатокены

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

Масштабирование

Абляционные исследования проводились на модели с 300 миллионами параметров с использованием 100 миллиардов обучающих токенов; окончательные модели были обучены на 1.5 триллионах токенов и масштабированы до моделей с 350 миллионами и 1.5 миллиарда параметров.

Архитектура Гимба

Слияние внимания и SSM

Гибридный модуль Гимба может быть представлен унифицированной и симметричной формулировкой. Входная последовательность ( \tilde{X} ) проецируется на запросы, ключи и значения для голов внимания, а также на входные признаки и ворота для голов SSM. Выходы обоих типов голов нормализуются и усредняются, а затем проецируются на выходной слой.

Интерпретация с точки зрения памяти

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

Анализ важности голов

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

Оптимизация KV кэша

Комбинация глобального и локального внимания

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

Обмен KV кэшем между слоями

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

Метатокены

Мы наблюдали, что начальные токены, хотя и не семантически важны, часто получают значительные внимательные оценки от последующих токенов. Чтобы направить внимание на токены, которые значительно вносят вклад в производительность задачи, мы вводим набор обучаемых метатокенов ( R = [r_1, r_2, ..., r_m] ), которые добавляются в начало входной последовательности.

Роль метатокенов

Метатокены выполняют несколько функций:

• Предотвращение перезаписи токенов: Метатокены действуют как независимые от входных данных токены, которые помогают модели избежать перезаписи и перефокусировки внимания на несущественные токены.
• Инициализация KV кэша и состояния SSM: Обучение начальных токенов можно рассматривать как форму обученного настройки подсказок или обученной инициализации. Для вывода метатокены фиксированы, и ключи и значения могут быть предварительно вычислены оффлайн и сохранены.

Заключение

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