Native Sparse Attention: новый вид разреженного внимания от DeepSeek. архитектура.. архитектура. внимание.. архитектура. внимание. искусственный интеллект.. архитектура. внимание. искусственный интеллект. трансформер.

Когда Маск выпустил Grok 3, а Сэм Альтман все еще колебался, стоит ли открывать исходный код, Лян Вэньфэн, как соавтор, работал с исследовательской группой DeepSeek над созданием шокирующей и сенсационной исследовательской статьи. DeepSeek официально представила свой последний научный прорыв — Native Sparse Attention (NSA)! Эта технология имеет большое значение. Она, скорее всего, значительно повысит способность следующего поколения больших языковых моделей обрабатывать длинные тексты, полностью учитывая при этом операционную эффективность. Нет сомнений, что это еще одна веха в области больших языковых моделей (LLM)!

1. Предпосылки исследования и общие выводы

Проблема, на решение которой направлена данная статья, заключается в высоких вычислительных затратах на моделирование длинных контекстов. Стандартный механизм внимания имеет высокую вычислительную сложность при обработке длинных последовательностей, что становится узким местом для производительности модели. Трудности исследования этой проблемы включают:

  • Как повысить эффективность, сохранив при этом возможности модели.

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

Сопутствующие исследовательские работы по этой проблеме включают:

  • Метод исключения кэша KV.

  • Метод выбора блочного кэша KV.

  • Метод выборки, кластеризации или выбора хэша и т. д.

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

Общий вывод

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

Ключевые инновации NSA

  1. Собственная конструкция разреженного внимания:

    • Позволяет проводить сквозную оптимизацию разреженного шаблона на этапе предварительного обучения.

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

  2. Иерархический механизм разреженного внимания:

    • Умело достигает баланса между локальной обработкой информации и глобальной обработкой информации.

    • Предоставляет более эффективное и комплексное решение для модели обработки длинных текстов.

Table_1

Table_1

Рисунок 1 | Сравнение производительности и эффективности между полным вниманием (Full Attention) и нашим NSA.
Слева: Несмотря на то, что NSA является разреженной (sparse) моделью, она в среднем превосходит базовый вариант с полным вниманием по общим бенчмаркам, задачам с длинным контекстом и оценке умозаключений.
Справа: Для обработки последовательностей длиной 64k NSA обеспечивает значительный выигрыш в скорости вычислений по сравнению с полным вниманием на всех этапах: декодировании, прямом проходе и обратном распространении ошибки.

2. Методы исследования

В данной статье предлагается NSA (Native Sparse Attention) — аппаратно-ориентированный и обучаемый механизм разреженного внимания, позволяющий решить проблему высоких вычислительных затрат при моделировании длительного контекста. Чтобы задействовать потенциал внимания в разреженном режиме, NSA заменяет исходные пары «ключ-значение» (k_t, v_t) на более компактный и информативный набор представлений (tilde{k}, tilde{v}), зависящих от каждого запроса (q_t). Так же инженеры стремились оптимизировать арифметическую интенсивность.

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

Table_2

Table_2

Рисунок 2 | Общий обзор архитектуры NSA.
Слева: Фреймворк обрабатывает входные последовательности через три параллельные ветви внимания. Для заданного запроса предыдущие ключи и значения преобразуются в сжатое внимание (compressed attention) для крупнозернистых паттернов, выборочное внимание (selected attention) для важных токенов и скользящее внимание (sliding attention) для локального контекста.
Справа: Визуализация различных паттернов внимания, которые формируются каждой ветвью. Зелёные области показывают зоны, где необходимо вычислять коэффициенты внимания, а белые области — это зоны, которые можно пропускать.

2.1 Общая схема алгоритма

Для эффективного использования механизма внимания в разреженном режиме в статье предлагается заменить исходные пары ключ-значение в классическом уравнении внимания трансформеров на более компактное и информационно насыщенное представление.

Это представление формируется для каждого запроса ( q_t ) следующим образом:

  • Оно основывается на текущем запросе ( q_t ) и контекстной памяти.

  • Динамически строится в зависимости от контекста.

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

В NSA используются три стратегии отображения ( C={ text{cmp}, text{slc}, text{win} } ), которые представляют собой сжатие, выбор и скользящее окно соответственно. Выход внимания в NSA формируется как комбинация результатов этих стратегий:

o'_t=sum_{c in C} g^c_t cdot Attn(q_t, tilde{k}^c, tilde{v}^c)

где каждая стратегия (c in C) генерирует свои (tilde{k}^c, tilde{v}^c). Величина (g^c_t in [0,1]) — это «gate»-оценка для каждой стратегии, вычисляемая на основе входных признаков с помощью MLP и сигмоидной активации. Она управляет вкладом каждой стратегии и общим количеством переназначенных ключей и значений:

Nt=sum{c in C} g^c_t cdot size[tilde{k}^c]

Таким образом, NSA поддерживает высокую степень разреженности, гарантируя, что N_t ll t, что позволяет снизить вычислительные затраты и повысить эффективность обработки длинных последовательностей. В частности, оптимизированный выход внимания формально определяется представленными выше уравнениями.

2.2 Динамическая многоуровневая разреженная стратегия:

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

  1. Сжатое внимание: имеет дело с крупнозернистыми шаблонами и собирает глобальную информацию путем сжатия блоков токенов.

  2. Выбранное внимание: обработка важных блоков токенов и выборочное сохранение подробной информации.

  3. Скользящее окно Внимание: Обработка локальной контекстной информации.

(1) Крупнозернистое сжатие (глобальное внимание):

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

Ниже приведена формула расчета для крупнозернистого сжатия:

K_{t}^{text{cmp}} ;=; f_{k}^{text{cmp}}bigl(k_{t}bigr) ;=; Bigl{,varphibigl(k_{id + 1:id+l}bigr);Bigm|;1 leqslant i leqslant Bigllfloorfrac{t - l}{d}Bigrrfloor Bigr}

где:

  • k_t – исходная последовательность векторов-ключей [k_1, k_2, dots, k_t] до момента t.

  • f_k^{text{cmp}} – функция сжатия, применяемая к исходным ключам.

  • varphi – обучаемый многослойный персептрон (MLP), который преобразует блок последовательных ключей в один сжатый вектор. При этом может использоваться позиционное кодирование, чтобы учитывать порядок токенов внутри блока.

  • k_{id+1:id+l} – блок из ll последовательных ключей, начинающийся с индекса id+1 и заканчивающийся индексом id+l; параметр l определяет длину блока.

  • d– шаг скольжения между блоками. Если d<l, блоки будут перекрываться, что помогает сохранить непрерывность информации.

  • lfloorfrac{t-l}{d}rfloor – число полных блоков, которые можно сформировать из последовательности длины t.

tilde{K}_{t}^{mathrm{cmp}} ;in; mathbb{R}^{,d_ktimesleftlfloor frac{t-l}{d} rightrfloor}

представляет собой тензор, состоящий из сжатых ключей. Обычно мы принимаем d < l, чтобы уменьшить фрагментацию информации. Для сжатых значений есть похожая формула. Сжатое представление фиксирует более грубую семантическую информацию более высокого уровня и снижает вычислительную нагрузку на внимание.

Как работает  обучаемый многослойный персептрон (MLP):

Функция varphi реализована как обучаемый многослойный перцептрон (MLP), который преобразует блок из ll последовательных ключей в один сжатый вектор размерности dmodeldmodel​. Вот пошаговое объяснение процесса:

1. Входные данные

  • Блок ключей:

    k{id+1:id+l} in mathbb{R}^{l times d_{text{model}}}

    где:

    • l — размер блока (например, 32 токена),

    • d_{text{model}} — размерность эмбеддинга каждого токена.

2. Подготовка входного вектора

  1. Конкатенация токенов:
    Все токены блока объединяются в один вектор:

    x_{text{concat}}=text{concat}(k_{id+1}, k_{id+2}, dots, k_{id+l}) in mathbb{R}^{l cdot d_{text{model}}}

    Например, для l=32 и dmodel​=512, размерность x_{text{concat}} будет 32 times 512=16384.

  2. Добавление позиционного кодирования:
    Каждому токену в блоке добавляется позиционное кодирование, чтобы сохранить информацию о порядке:

    x_{text{pos}}=x_{text{concat}} + text{PositionalEncoding}(id+1, id+2, dots, id+l)

3. Архитектура MLP

MLP состоит из следующих слоёв:

varphi(x)=text{LayerNorm}(W_2 cdot text{GELU}(W_1 cdot x_{text{pos}} + b_1) + b_2)

где:

  • W_1 in mathbb{R}^{(l cdot d_{text{model}}) times h} — матрица весов первого слоя,

  • W_2 in mathbb{R}^{h times d_{text{model}}} — матрица весов второго слоя,

  • h — размер скрытого слоя (обычно h=d_{text{model}}),

  • GELU — активационная функция,

  • LayerNorm — слой нормализации.

4. Как это работает

  1. Сохранение семантики:
    MLP обучается выделять наиболее важные признаки из блока токенов. Например, для текста он может акцентировать ключевые слова или фразы.

  2. Динамическая адаптация:
    Веса W_1 и W_2 обучаются в процессе тренировки модели, что позволяет φ гибко адаптироваться к разным типам данных.

  3. Позиционная информация:
    Добавление позиционного кодирования помогает модели учитывать порядок токенов внутри блока.

5. Пример сжатия

  • Исходный блок: 32 токена → 32 × 512 = 16384 параметров.

  • После сжатия: 1 × 512 параметров.Сокращение в 32 раза по числу токенов.

6. Преимущества подхода

  • Вычислительная эффективность: Уменьшает размерность ключей/значений для механизма внимания.

  • Адаптивность: MLP учится агрегировать информацию лучше, чем фиксированные методы (усреднение, максимум).

  • Сохранение контекста: Позиционное кодирование и нелинейности сохраняют семантику блока.

Итог

Функция φ динамически “резюмирует” блок токенов в компактный вектор, сохраняя ключевую информацию. Это позволяет NSA обрабатывать длинные последовательности с низкими вычислительными затратами.

Как работает формула:

  1. Разбиение последовательности: Исходная последовательность ключей k_t​ разбивается на блоки длины l с шагом d. Например, если l=32и d=16, каждый новый блок начинается на 16 токенов дальше предыдущего, а сами блоки могут частично перекрываться.

  2. Применение функции сжатия: Для каждого блока k_id+1:id+l​ функция φ (MLP) преобразует этот набор токенов в один сжатый вектор. Этот процесс позволяет “сжать” информацию из группы токенов в компактное представление, сохраняя при этом важные семантические связи.

  3. Формирование множества сжатых ключей: Результатом является множество сжатых ключей K_t^{text{cmp}}​, состоящее из lfloorfrac{t-l}{d}rfloor элементов, что существенно меньше исходного числа токенов tt. Эти сжатые ключи далее используются в механизме внимания вместо исходных, что значительно снижает вычислительную сложность.

Интуитивное понимание:

Представьте, что вам нужно быстро ознакомиться с содержанием длинного текста, например, романа или новостной сводки. Вместо того чтобы читать каждую страницу, вы разбиваете текст на главы и просматриваете краткие резюме каждой из них. Здесь:

  • Исходный текст – это последовательность ключей k_t​,

  • Глава – это блок k_{id+1:id+l},

  • Резюме главы – это сжатый вектор, полученный посредством varphi.

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

Преимущества данного метода:

  • Вычислительная эффективность: Сокращение числа токенов снижает сложность механизма внимания, что особенно важно для длинных последовательностей.

  • Масштабируемость: Модель способна обрабатывать значительно большие контексты.

  • Адаптивность: Обучаемая функция φφ позволяет динамически выделять наиболее важную информацию в каждом блоке, что делает метод более гибким по сравнению с простыми операциями усреднения или пуллинга.

Примечание: В юпитер тетрадке ‘CompressedAttention.ipynb’ представлен пример использования данного метода.

На основе проведенных экспериментов можно сделать следующие выводы:

Ключевые выводы при тестировании

  1. Эффективность зависит от длины последовательности:

    • Для последовательностей <1000 токенов стандартное внимание эффективнее

    • Для последовательностей >10K токенов сжатое внимание дает существенное преимущество

    • Потенциальное ускорение растет с увеличением длины последовательности

  2. Оптимальные параметры сжатия:

    • Для длинных последовательностей эффективнее использовать большие размеры блоков (256 вместо 32)

    • Оптимальное соотношение размера блока и шага (stride/block_size) близко к 0.5, что обеспечивает хороший баланс между сжатием и сохранением информации

  3. Компромисс между эффективностью и точностью:

    • Сжатое внимание жертвует гранулярностью внимания на уровне отдельных токенов ради обработки длинных контекстов

    • Для полной архитектуры NSA этот компромисс компенсируется комбинированием с другими механизмами внимания

  4. Аппаратные соображения:

    • Сжатое внимание не только ускоряет вычисления, но и значительно сокращает требования к памяти

    • Потенциально позволяет обрабатывать контексты длиной 100K+ токенов на стандартном оборудовании

Вывод

Эксперимент убедительно демонстрирует, что механизм сжатого внимания является эффективным решением проблемы квадратичной сложности стандартного внимания для длинных последовательностей. Фактическое ускорение составило (4.69×), оно будет увеличиваться с ростом длины последовательности.

(2) Выборочное внимание (локальное внимание):

Выборочное внимание (Selected Attention) — вторая ветвь в архитектуре NSA, которая фокусируется на выборе наиболее релевантных блоков токенов для текущего запроса. В отличие от крупнозернистого сжатия, которое обрабатывает глобальный контекст путем сжатия информации, выборочное внимание сохраняет детальную информацию на уровне отдельных токенов, но только для наиболее важных частей последовательности. Этот механизм выбора гарантирует, что модель сможет сохранить ключевую локальную информацию при обработке длинных последовательностей и избежать потери важных деталей из-за сжатия. Формула выглядит следующим образом:

1. Вычисление оценок важности для сжатых блоков:

p_{t}^{slc}=text{Softmax} left( q_{t}^{T} tilde{K}_{t}^{cmp} right)

где:

  • q_t​ — текущий запрос (query vector) для позиции tt

  • tilde{K}_{t}^{cmp}— матрица сжатых ключей, полученных на этапе крупнозернистого сжатия

  • p_{t}^{slc}— вектор вероятностей (оценок важности), где каждый элемент p_{t}^{slc}[i] представляет относительную важность i-го сжатого блока для текущего запроса

Здесь мы вычисляем скалярное произведение запроса q_t​ с каждым сжатым ключом, а затем применяем функцию Softmax для нормализации этих оценок в распределение вероятностей. Это позволяет определить, насколько каждый сжатый блок релевантен для текущего запроса.

2. Выбор наиболее важных блоков:

I_{t}=left{ i mid text{rank} left( p_{t}^{slc}[i] right) leq n right}

где:

  • I_t​ — множество индексов выбранных блоков

  • text{rank}(p_{t}^{slc}[i]) — ранг i-го элемента в векторе вероятностей, где ранг 1 соответствует наивысшей вероятности

  • n — гиперпараметр, определяющий число выбираемых блоков

Эта формула означает, что мы выбираем n блоков с наивысшими оценками важности. Параметр n является ключевым для определения компромисса между вычислительной эффективностью и качеством моделирования.

(3) Скользящее окно:

Механизм скользящего окна используется для обработки локальной контекстной информации. Он сохраняет самые последние токены (например, самые последние 512 токенов), гарантируя, что модель может быстро адаптироваться к изменениям в локальных шаблонах при обработке текущего токена. Этот механизм предотвращает «короткое замыкание» модели локальными шаблонами во время сжатия и выбора, тем самым гарантируя, что модель сможет обрабатывать как локальные, так и долгосрочные зависимости. Формула выглядит следующим образом:

tilde{K}_{t}^{win}=k_{t-w:t}, tilde{V}_{t}^{win}=v_{t-w:t}

Давайте разберем каждый элемент формулы:

  • tilde{K}_{t}^{win} — модифицированная матрица ключей для механизма скользящего окна в позиции t

  • tilde{V}_{t}^{win} — модифицированная матрица значений для механизма скользящего окна в позиции t

  • k_{t-w:t} — подпоследовательность векторов-ключей от позиции t−w до текущей позиции t

  • v_{t-w:t}​ — соответствующая подпоследовательность векторов-значений

  • w — размер окна (гиперпараметр, например, равный 512 токенам)

Принцип работы скользящего окна

Механизм скользящего окна предельно прост, но при этом чрезвычайно эффективен. Вместо использования всей истории векторов-ключей и значений с начала последовательности (что привело бы к квадратичной сложности), модель учитывает только фиксированное количество последних токенов:

  1. Для каждого текущего запроса qtqt​ формируется локальное окно внимания шириной ww.

  2. В это окно попадают только ww последних токенов.

  3. Токены за пределами окна игнорируются при вычислении внимания.

Вычислительные преимущества

Скользящее окно обеспечивает постоянную вычислительную сложность O(w) для каждого токена, в отличие от стандартного внимания, где сложность растет линейно с длиной последовательности O(t). Это дает существенные преимущества:

  1. Постоянное потребление памяти: Для каждой позиции хранится фиксированное количество пар ключ-значение.

  2. Стабильное время вычислений: Время обработки токена не зависит от длины предыдущего контекста.

  3. Оптимизация для GPU: Фиксированный размер окна позволяет эффективно использовать GPU-память и вычислительные блоки.

Роль скользящего окна в архитектуре NSA

В общей архитектуре NSA скользящее окно выполняет три критические функции:

  1. Локальная контекстуализация: Обеспечивает доступ к непосредственному контексту, что важно для семантического понимания текущего токена.

  2. Предотвращение “короткого замыкания”: Даже если две другие компоненты NSA (сжатое и выборочное внимание) пропустят важную локальную информацию, скользящее окно гарантирует, что недавние токены будут учтены.

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

Практический пример

Представим задачу генерации текста с длинным контекстом. При генерации слова в позиции 10000:

  • Скользящее окно (w=512): Учитывает токены с позиций 9488-10000

  • Сжатое внимание: Предоставляет общую информацию обо всем предыдущем контексте

  • Выборочное внимание: Выделяет ключевые токены из всего контекста, важные для текущего запроса

При этом скользящее окно гарантирует, что непосредственный локальный контекст не будет потерян из-за сжатия или селективного выбора, которые могут упустить важные локальные связи.

Все три механизма дополняют друг друга, обеспечивая как глобальное понимание контекста, так и детальное восприятие локальной информации.

Итог

Шаг 1: Сжатие (обобщение фраз):

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

Шаг 2: Выборка (выбор важных слов):

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

Шаг 3: Скользящее окно (сохранение локального контекста):

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

2.3 Дизайн обучения восприятия:

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

2.4 Система аппаратного выравнивания:

Для достижения ускорения на уровне FlashAttention во время обучения и предварительного заполнения, в статье реализовано аппаратно-оптимизированное ядро разреженного внимания на Triton. Оптимизировано блочное разреженное внимание для использования Tensor Core и доступа к памяти, обеспечивая сбалансированную арифметическую интенсивность. Конкретно, были применены следующие оптимизации:

  • Блочный режим доступа к памяти: Увеличивает использование Tensor Core за счет объединенной загрузки, уменьшая избыточную передачу ключей и значений (KV).

  • Циклическое планирование: Умное расположение циклов в ядре для устранения избыточной передачи KV.

  • Групповая загрузка данных: Для каждого внутреннего цикла загружаются все запросы внутри группы и их общие индексы разреженных блоков ключей/значений.

  • Совместное получение KV: Внутри внутреннего цикла последовательно загружаются непрерывные блоки ключей/значений в SRAM для минимизации загрузки памяти.

  • Внешний цикл с планированием сетки: Поскольку длина внутреннего цикла почти одинакова для разных блоков запросов, циклы запросов/выводов размещаются в планировщике сетки Triton для упрощения и оптимизации ядра.

Table_3

Table_3

Рисунок 3 | Конструкция ядра NSA. Ядро загружает запрос через группу GQA (сетчатый цикл), извлекает соответствующие разреженные блоки «ключ-значение» (внутренний цикл) и выполняет вычисления внимания в SRAM. Зеленые блоки представляют данные на SRAM, а синие блоки представляют данные на HBM.

3. Настройки обучения

3.1 Настройки предварительного обучения

Следуя передовым практикам в современных крупных языковых моделях (LLM), в их экспериментах используется архитектура, сочетающая групповое внимание запросов (GQA) и смесь экспертов (MoE). Модель содержит 27 миллиардов параметров, из которых 3 миллиарда являются активными. Модель состоит из 30 слоев с размерностью скрытого слоя 2560. Для GQA количество групп установлено на 4, с общим количеством голов внимания 64. Для каждой головы размерности скрытых слоев для запросов, ключей и значений настроены как (d_q=d_k=192) и (d_v=128). Для MoE используется структура DeepSeekMoE с 72 маршрутизируемыми экспертами и 2 общими экспертами, при этом количество выбираемых экспертов (K) установлено на 6. Для обеспечения стабильности обучения, первый слой MoE заменен на многослойный перцептрон (MLP) в форме SwiGLU.

3.2 Параметры архитектуры NSA

Архитектура обеспечивает эффективный баланс между вычислительными затратами и производительностью модели. Для NSA установлены следующие параметры:

  • Размер сжатого блока (l=32),

  • Шаг скользящего окна (d=16),

  • Размер выбранного блока (l′=64),

  • Количество выбранных блоков (n=16) (включая 1 фиксированный начальный блок и 2 локальных блока),

  • Размер скользящего окна (w=512).

Модели с полным вниманием и разреженным вниманием предварительно обучаются на 270 миллиардах токенов с длиной текста 8k, а затем дообучаются на текстах длиной 32k с использованием YaRN для адаптации к длинному контексту. Обе модели обучаются до полной сходимости для обеспечения справедливого сравнения.

4. Результаты и анализ

4.1 Общая оценка:

В таких бенчмарках, как MMLU, MMLU-PRO, CMMLU, BBH, GSM8K, MATH, DROP, MBPP и HumanEval , NSA превосходит базовый уровень полного внимания по большинству показателей, несмотря на его более высокую разреженность.

Table_4

Table_4

4.2 Длинная оценка контекста :

Авторы оценили методику собственного разреженного внимания (NSA) на тестовой платформе LongBench и сравнили ее эффективность с современным методом разреженного внимания (SOTA) и базовым уровнем полного внимания. Чтобы обеспечить единообразие разреженности между различными методами, исследователи установили количество токенов, активируемых для каждого запроса во всех базовых уровнях разреженного внимания, равным 2560, что в точности соответствует среднему количеству токенов, активируемых NSA при обработке последовательностей длиной 32 тыс.

Согласно правилам StreamLLM, бюджет в 2560 токенов включает в себя первые 128 токенов и 512 локальных токенов. Учитывая, что некоторые подмножества в LongBench имеют в целом низкие оценки по всем моделям, участвующим в оценке, и не могут служить ценным источником для сравнения производительности различных методов, авторы исключили эти подмножества из области оценки. Результаты оценки показаны в Таблице ниже. NSA показывает исключительно хорошие результаты, лидируя со средним баллом 0,469, что значительно превосходит все базовые показатели. В частности, его средний балл на 0,032 выше, чем у метода полного внимания, и на 0,046 выше, чем у метода Exact-Top.

Table_5

Table_5
Table_6

Table_6

4.3 Оценка цепного мышления:

При оценке рассуждений по инструкциям AIME NSA-R после контролируемой тонкой настройки превосходит базовый уровень полного внимания-R при длине контекста как 8 тыс., так и 16 тыс.

Table_7

Table_7

Таблица 3 | Оценка на основе инструкций AIME после обучения с учителем. Наш NSA-R демонстрирует лучшую производительность по сравнению с Full Attention-R как при длине последовательности 8k, так и при 16k.

Рисунок 6 | Сравнение ядра NSA на базе Triton с ядром FlashAttention-2 на базе Triton. Наша реализация значительно снижает задержку при всех длинах контекста, и улучшение становится более заметным по мере увеличения длины входных данных.

5. Подведем итог

5.1. Инновации которые были достигнуты в данной работе

Значительное ускорение: Благодаря оптимизации разработки алгоритма и современной аппаратной реализации, NSA достигает значительного ускорения по сравнению с моделью полного внимания на этапах декодирования, прямого распространения и обратного распространения, особенно при обработке последовательностей длиной 64 тысячи токенов.

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

Стратегия динамической иерархической разреженности: NSA сочетает грубое сжатие токенов и точный выбор токенов для сохранения как глобальной осведомленности о контексте, так и локальной точности.

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

Разработка алгоритма с учетом обучения: Стабильное сквозное обучение достигается за счет эффективных алгоритмов и обратных операторов, поддерживающих эффективное развертывание и сквозное обучение.

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


🔥Не пропустите важные обновления и углубленные материалы!🔥  

Хотите быть в курсе самых свежих обзоров и исследований в мире ML и AI? Переходите по ссылкам ниже, чтобы получить доступ к эксклюзивному контенту:  

📌 Все обзоры также доступны в нашем Telegram канале TheWeeklyBrief📢

📌 Более подробный обзор с математической формализацией и программным кодом ждет вас в нашем репозитории Weekly-arXiv-ML-AI-Research-Review 👩‍💻📂✨  

Не упустите шанс глубже погрузиться в мир технологий! 🚀

Автор: Verbasik

Источник

Rambler's Top100