Griffin
Griffin : сочетание закрытых линейных повторений с локальным вниманием для эффективных языковых моделей
arXiv
Модельная архитектура
Все наши модели содержат следующие компоненты: (i) остаточный блок, (ii) блок MLP и (iii) блок временного смешивания. Хотя (i) и (ii) одинаковы для всех моделей, мы рассматриваем три блока временного смешивания: глобальное внимание к множественным запросам (MQA), локальное (скользящее окно) MQA и предложенный нами рекуррентный блок. В составе рекуррентного блока мы используем Real-Gated Linear Recurrent Unit (RG-LRU) – новый рекуррентный слой, вдохновленный Linear Recurrent Unit Orvieto et al., 2023b.
Остаточный блок, как показано на рисунке 2 (а), определяет глобальную структуру наших моделей и вдохновлен Трансформаторами до нормы (Xiong et al., 2020). После внедрения входной последовательности мы передаем ее через $Н$ такие блоки ( $Н$ обозначающий глубину модели), а затем мы применяем RMSNorm Zhang and Sennrich, 2019 для получения окончательных активаций. Для вычисления вероятностей токенов мы применяем последний линейный слой, за которым следует softmax. Веса этого слоя используются совместно с входным слоем внедрения.
Остаточный блок
Рисунок 2: а) Основной основой нашей архитектуры режимов является остаточный блок, который складывается $Н$ раз. б) Закрытый блок MLP, который мы используем. в) Рекуррентный блок, который мы предлагаем в качестве альтернативы Multi Query Attention (MQA). Он использует предложенный нами уровень RG-LRU, определенный в разделе 2.4.
Остаточный блок содержит два компонента, применяемых по порядку. Первый компонент принимает скрытое состояние $чи$ и применяет RMSNorm Zhang and Sennrich, 2019, за которым следует блок временного смешивания. Затем мы объединяем вывод с пропущенным соединением из $чи$ через дополнение. Аналогично, второй компонент применяет RMSNorm, за которым следует блок MLP, а затем объединяет его выходные данные с пропущенным соединением со входа RMSNorm. Этот блок проиллюстрирован на рисунке 2 (а).
блок МЛП
Мы используем вентильный блок MLP Dauphin et al., 2017 (показан на рисунке 2(b)), который создает две ветви на основе входных данных измерения. $Д$ . Применяем линейный слой с выходным размером $MD$ на каждой ветке, где $М$ обозначает коэффициент расширения. Для простоты мы используем $М=3$ на протяжении всей этой работы. Мы применяем нелинейность GeLU Hendrycks and Gimpel, 2016 к одной из ветвей перед их объединением путем поэлементного умножения, аналогично GeGeLU Shazeer, 2020. Однако в нашем блоке MLP мы применяем последний линейный слой с выходным измерением. $Д$ на выходах слоя GeGeLU.
Временно-смешивающие блоки
Блок временного смешивания — это компонент нашей модели, который объединяет активации скрытого слоя в разных временных точках последовательности. Мы рассматриваем три блока временного смешивания: глобальный MQA Shazeer, 2019, локальный MQA Beltagy et al., 2020 и предлагаемый нами Recurrent блок.
Глобальное внимание к нескольким запросам
Если не указано иное, мы используем MQA, а не MHA, чтобы улучшить скорость вывода наших базовых линий Transformer Shazeer, 2019. Мы используем фиксированный размер головы. $D_{head}=128$ , и фиксируем количество голов внимания $H$ такой, что $HD_{head}=D$ . Для этого требуется размер модели $Д$ быть кратным 128. Мы не используем какие-либо абсолютные позиционные встраивания, но мы используем вращающееся позиционное встраивание (RoPE) Su et al., 2021 в качестве относительного позиционного встраивания.
Внимание к локальному скользящему окну
Одним из ключевых недостатков использования глобального внимания является то, что его вычислительная сложность растет квадратично с длиной последовательности. Чтобы решить эту проблему, в нескольких работах начали использовать местное внимание Beltagy et al., 2020, также известное как внимание скользящего окна. Это позволяет каждой позиции обслуживать только фиксированное количество токенов в прошлом. Это не только уменьшает количество вычислительных операций FLOP, но также ограничивает размер KV-кеша размером окна, что делает его более не квадратичным по длине последовательности. Все остальные детали такие же, как в глобальном MQA.
Рекуррентный блок
Наш рекуррентный блок (рис. 2(c)) аналогичен блоку GSS Mehta et al., 2022 и блоку, используемому Мамбой Гу и Дао, 2023. Мы принимаем входные данные размерности $Д$ и примените два линейных слоя с выходным размером $D_{RNN}$ параллельно создавая две ветки. К первой ветке мы применяем небольшой отделимый слой Conv1D, вдохновленный Shift-SSM в H3 Dao et al., 2022b, с размерностью временного фильтра, равной 4. Обратите внимание, что этот слой Conv1D очень мал, всего лишь $4D$ параметры. Мы следуем за слоем Conv1D с предложенным нами слоем RG-LRU (определенным ниже). На второй ветви мы применяем нелинейность GeLU, а затем объединяем ветви путем поэлементного умножения. Затем мы применяем последний линейный слой с выходным размером $Д$ .
Линейный рекуррентный блок с реальным вентилем (RG-LRU)
Предлагаемый нами уровень RG-LRU имеет простую рекуррентность, вдохновленную Linear Recurrent Unit (LRU) Orvieto et al., 2023b, но включает в себя механизм пропуска, мотивированный литературой по нелинейным RNN, в частности LSTM Hochreiter и Schmidhuber, 1997 и GRUs Chung et al., 2014. Уравнения, описывающие слой, следующие:
$$begin{align} r_t &= sigma(W_{a} x_t + b_a), & text{ворота повторения} \ i_t &= sigma(W_{x} x_t + b_x), & text{ входной вентиль} \ a_t &= a^{cr_t}, & text{} \ h_t &= a_t odot h_{t-1} + sqrt{1 - a_t^2} odot (i_t odot x_t). & text{} end{align}$$
Выход слоя $y_t=h_t$ , а нелинейность $сигма$ в уравнениях – сигмоидальная функция. Рекуррентный вес $а$ в уравнении (4) является диагональной. Следовательно, все операции выполняются поэлементно. Мы параметризуем $а$ в уравнении (3) как $a=sigma(Lambda)$ , где $Лямбда$ это обучаемый параметр. Это гарантирует, что $0 <= а <= 1$ , гарантируя, что рецидив является стабильным. Переменная $c$ — скалярная константа, равная 8. Для обеспечения численной стабильности на практике мы вычисляем $a^{cr_t}$ в лог-пространстве (см. Приложение А). Слой имеет вентили как на входе, так и на входе. $х$ и текущий вес $а$ . Однако ни один из ворот не зависит от рекуррентного состояния. $h_{t-1}$ , что гарантирует эффективное выполнение вычислений на устройстве. Мы инициализируем оба $W_{a}$ и $W_{b}$ используя LeCun init LeCun et al., 2002. Мы инициализируем $Лямбда$ такой, что $а^с$ равномерно распределяется между $0,9$ и $0,999$ в начале обучения, аналогично (Орвието и др., 2023б.). В отличие от многих недавних работ в литературе по SSM, RG-LRU не использует инициализацию, вдохновленную теорией ортогональных полиномов Gu et al., 2020, а также не определяется как дискретизация базовой непрерывной системы Gu et al., 2021а. В отличие от исходного слоя LRU, мы не используем комплексную алгебру в рекурсии. Хотя использование сложных повторений привело бы к более выразительному слою, Orvieto et al., 2023a, мы обнаружили, что сложные повторения не приносят пользы для практического моделирования языка, как также заметили Гу и Дао, 2023 (см. Приложение B).
Поведение ворот
Входные ворота $i_t$ похож на тот, что в LSTM, который может фильтровать (или уменьшать) входные данные $x_t$ . Однако, насколько нам известно, наши ворота возврата $r_t$ отличается от других воротных механизмов, описанных в литературе. Например, механизм выбора, предложенный в «Мамба Гу и Дао, 2023», можно сравнить с механизмом обновления ГРУ, который интерполирует $x_t$ . Его влияние на скрытое состояние позволяет ему сбросить свое состояние и забыть любую информацию, которую он хранит из прошлого, аналогично воротам забывания в LSTM. Напротив, наш вентиль повторения может приблизительно интерполировать между стандартным обновлением LRU из Orvieto et al., 2023a и предыдущим скрытым состоянием, что позволяет ему эффективно отбрасывать входные данные и сохранять всю информацию из предыдущей истории (дополнительную информацию см. в Приложении A). ). Мы считаем, что ключевая роль этих ворот заключается в том, чтобы позволить модели достичь суперэкспоненциальной памяти за счет уменьшения влияния неинформативных входных данных.
