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

tldr; в статье рассмотрены приёмы, позволяющие ускорить обучение больших языковых моделей (LLM) и нарастить в них вывод (inference). Для этого нужно использовать большое контекстное окно, в котором умещается до 100K входных токенов. Вот эти приёмы: ALiBi с подмешиванием в вектор позиции слова в последовательности (positional embedding), разреженное внимание (Sparse Attention), мгновенное внимание (Flash Attention),  многозапросное внимание, условные вычисления и GPU A100 на 80 ГБ.  

Недавно прошло несколько анонсов по поводу новых больших языковых моделей (LLM), настроенных на приём исключительно крупного контекстного окна, целых 65K токенов (MPT-7B-StoryWriter-65k+ от MosaicML) или даже 100K токенов (описано в статье Introducing 100K Context Windows от Antropic). Google в техническом отчёте по Palm-2 не упоминает конкретный размер контекстного окна, но указывает, что удалось «значительно увеличить длину контекстного окна в модели».

Для сравнения: современная модель GPT-4 может работать, получая контекст длиной 32K входных токенов. Большинство же опенсорсных моделей LLM работают с контекстным окном длиной 2K токенов.

Это впечатляет, так как при работе с таким огромным контекстным окном затравка (prompt) может буквально сравниться по размеру с целой книгой. Так, в романе «Великий Гэтсби» насчитывается 72K токенов, 210 страниц, и этот роман прочитывается за 6 часов при скорости чтения 1,7 минут/страница. Так что вот какой объём «специфичной» информации может просканировать модель и далее хранить для обработки запросов!

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

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

  • Каковы главные ограничения оригинальной архитектуры трансформеров, если применять её для обработки длинных контекстов

  • Вычислительная сложность трансформерной архитектуры

  • Какие в настоящее время существуют приёмы оптимизации, позволяющие ускорить трансформер и увеличить длину контекста до 100K

"Краткая" аннотация

Здесь и далее я буду в качестве синонимов использовать термины «длина контекста», «окно контекста» и «количество входных токенов», обозначая их как n.

Этот пост получился несколько длинным, так что сделаю здесь подборку основных моментов и приёмов, рассмотренных ниже:

  • 1-я задача – это квадратичная сложность (пространственная и временная) у некоторого слоя внимания относительно количества n токенов, подаваемых на вход.

  • Когда размер вектора d > n2-я задача сводится к квадратичной временной сложности линейных слоёв относительно вектора, имеющего размер d.

  • 3-я задача – это позиционный синусоидальный эмбеддинг, применяемый в оригинальной архитектуре.

  • В трансформерной архитектуре очертания доступных для изучения весов матрицы не зависят от количества входных токенов n.

  • Таким образом, трансформер, обученный на 2K контекстных примеров может потреблять токены любой длины, даже 100K. Но на стадии вывода модель не произведёт осмысленного результата на материале 100K токенов, если её не обучали на 100K.

  • Обучение бесхитростного трансформера на гигантском корпусе текстов и только с использованием очень длинных контекстов оказывается невыполнимо затратным из-за квадратичной сложности относительно n и d. По имеющейся оценке, обучение LLaMA на контексте длиной 2K обошлось бы в ~$3M. Следовательно, обучение LLaMA на 100K стоило бы ~$150M.

  • Есть вариант: обучить модель на контексте в 2K токенов, а затем тонко настроить её на более длинных контекстах (например, 65K). Но такой подход не сработает с оригинальной архитектурой трансформеров, всё дело – в позиционной синусоидальной кодировке.

  • [Фокус #1] Можно избавиться от позиционной синусоидальной кодировки и воспользоваться ALiBi; это простой и изящный позиционный эмбеддинг, который не портит точности. С его помощью можно обучить модель на 2K, а далее тонко настраивать на 100K.

  • [Фокус #2] Не требуется вычислять показатели внимания между всеми токенами. Одни токены важнее других, поэтому можно воспользоваться Разреженным Вниманием. Так можно одновременно ускорить и обучение, и вывод.

  • [Фокус #3] Flash Attention эффективно реализует слой внимания для GPU. Этот инструмент использует замощение и обходится без материализации больших промежуточных матриц (n, n), не умещающихся в SRAM GPU. Так можно одновременно ускорить и обучение, и вывод.

  • [Фокус #4] Использовать многозапросное внимание вместо многоголового. В таком случае все головы смогут совместно использовать веса, при этом линейно проецируя K и V. Так удаётся радикально ускорить поступательный вывод.

  • [Фокус #5] При условных вычислениях можно не применять всех параметров модели ко всем токенам, содержащимся во входной последовательности. CoLT5 применяет тяжеловесные вычисления только с самыми важными токенами и процессами, а для обработки остальных токенов применяются облегчённые варианты слоёв. Так можно одновременно ускорить и обучение, и вывод.

  • [Фокус #6] Чтобы вместить большой контекст, вам потребуется много RAM в GPU, поэтому в так их случаях принято использовать GPU на 80 ГБ A100.

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

Теперь давайте обсудим все эти моменты более подробно.

Почему длина контекста так важна

Длина контекста – один из важнейших лимитирующих факторов при работе с большими языковыми моделями. Увеличить контекст до 100K – это уже невероятное достижение, ставшее реальностью (занятно, как этот тезис будет восприниматься через год).

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

В настоящее время обходить это ограничение пробуют по-разному, а именно:

  • При помощи приёмов резюмирования и изощрённых сцепленных затравок

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

  • Когда это возможно – тонко настраивать LLM на данных, предоставляемых пользователем (такая возможность предоставляется не во всех коммерческих LLM, а для опенсорсных LLM это не самая тривиальная задача)

  • Разработка специализированных сравнительно небольших LLM для конкретных данных, которые нас интересуют (опять же, не самая тривиальная задача)

При наличии длинного контекстного окна уже имеющаяся в вашем распоряжении большая языковая модель (видевшая целый Интернет) может изучить имеющийся у вас контекст и данные, а затем взаимодействовать с вами на совершенно ином уровне, предполагающем более высокую персонализацию. Всё это – без изменения весов модели, когда всё «обучение» производится на лету, «в памяти». В целом, чем больше контекстное окно, тем более высокая точность, беглость и изобретательность приобретается моделью.

В качестве аналогии здесь можно рассмотреть ОЗУ компьютера, где операционная система хранит в режиме реального времени актуальный контекст для всех ваших приложений. LLM, располагая достаточно длинным контекстом, сравнима с «рассуждающим компьютером», учитывающим широкий контекст, предоставляемый пользователем.

Оригинальный трансформер и длина контекста

Важно отметить, что в архитектуре транмформеров формы всех весов матриц, доступных для обучения, не зависят от количества подаваемых на вход токенов n. Все параметры, поддающиеся обучению (поиск по векторам, слои проекций, слой softmax и слои внимания) не зависят от длины входного фрагмента и должны быть в состоянии обрабатывать такие фрагменты варьирующейся длины. Просто отлично, когда такое свойство в архитектуре предоставляется прямо «из коробки».

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

Если требуется обучить трансформер на таком большом контексте, то можно, например, обучать его в два этапа: сначала базовую модель на окне контекста длиной 2K токенов, а потом продолжить обучение (в качестве тонкой настройки) на более длинных контекстах (например, 65K или 100K). Именно это и было сделано с моделью MosaicML. Но вот загвоздка: такой подход не сработает с оригинальной архитектурой Трансформеров, поэтому придётся прибегать к определённым ухищрениям (см. Фокус #1 ниже в этом посте).

Напомню, что такое "многоголовый подход к вниманию"

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

Q — запросы, K — ключи и V — значения. Такие обозначения типичны для статьи, касающейся извлечения информации, где в систему подаётся «запрос», а далее подыскивается ближайший к нему «ключ».   
n —  количество входных токенов
d — размерность текстового векторного представления
h — количество голов в слое внимания
k— размер линейной проекции для Q и K
v — размер линейной проекции для V

Многоголовое внимание

  1. Можно выполнить в слое вектора такой поиск, по результатам которого для заданного токена возвращается вектор размером (1, d). Соответственно, для последовательности из n токенов получаем матрицу текстовых векторных представлений X размера (n, d). Затем суммируем её с позиционным синусоидальным эмбеддингом.

  2. Слой многоголового внимания нужен для того, чтобы рассчитать новый вектор для данной последовательности токенов. Эта последовательность такова, что может считаться оригинальным текстом, кодирующим X, но при этом является взвешенной (1) по относительной важности среди всех токенов относительно контекста и (2) по относительным позициям токенов.

  3. Эту матрицу вектора X (n, d) мы обрабатываем параллельно на h слоях внимания (головах). Чтобы получить QK и для всех голов в слое внимания, мы линейно  проецируем X на размерности kk и v, соответственно. Это делается путём умножения X на h матриц формы (d, k)(d, k) и (d, v). Можете трактовать эту операцию как перемножение (n, d) на (h, d, k)(h, d, k) и (h, d, v).

  4. Головы в слое внимания возвращают матриц с баллами внимания размера (n, v). Затем сцепляем фрагменты от всех голов (n, h*v) и линейно проецируем их на следующие этапы.

Общий вид архитектуры внимания; схема взята из статьи Attention is All You Need
Общий вид архитектуры внимания; схема взята из статьи Attention is All You Need

Взвешенное скалярное внимание

Теперь давайте как под лупой рассмотрим одну голову в слое внимания.

  1. Q, K, V – это 3 линейные проекции размера (n, k)(n, k) и (n, v), получаемые умножением доступных для обучения весов отдельно для каждой головы.

  2. Получаем баллы внимания, отдельно вычисляя расстояние (скалярное произведение) между Q и (в транспонированном виде). Умножаем матрицу (n, k) на (k, n) и получаем матрицу (n, n). Затем умножаем её на маскирующую матрицу, чтобы обнулить некоторые токены (требуемые в декодере). Затем масштабируем её и применяем softmax в диапазоне от 0 до 1. Таким образом, получаем матрицу формы (n, n) с n_ij – это относительный балл внимания от 0 до 1 между i-м и j-м токеном, демонстрирующий степень «близости» этих токенов в данном конкретном контексте длины n.

  3. Затем умножаем эту матрицу баллов внимания (n, n) на “значения” V размером (n, d), чтобы получить текстовое векторное представление, взвешенное по этим относительным баллам внимания.

В оригинальной статье матрица баллов внимания для одной головы вычисляется по этой формуле
В оригинальной статье матрица баллов внимания для одной головы вычисляется по этой формуле

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

Очень красивый код, демонстрирующий формы матриц на каждом шаге работы со слоем внимания. Код взят из статьи Multi-Query .
Очень красивый код, демонстрирующий формы матриц на каждом шаге работы со слоем внимания. Код взят из статьи Multi-Query .

Сложность трансформера и длина контекста

Сложность перемножения 2 матриц размером (a,b)*(b,c) равна O(a*b*c).
Для простоты предположим, что k*h = O(d) и будем  от этого отталкиваться, чтобы вывести сложность внимания.

Сложность слоя внимания складывается из двух частей:

  1. Линейные проекции для получения Q, K, V: умножение матрицы эмбеддинга размером (n, d) на h матриц, доступных для обучения (d, k)(d, k) и (d, v). Следовательно, сложность ~ O(nd²)

  2. Операции умножения Q на K в трансформированном виде и далее умноженных на V: (n,k) * (k,n) = (n,n) и (n,n)*(n,v) = (n,v). Сложность ~ O(n²d)

Таким образом, сложность уровня внимания равна O(n²d + nd²), где n — это длина контекста (количество входных токенов), а d — размер эмбеддинга. Так что здесь и далее видим, что сложность вычислений на уровне внимания квадратично возрастает относительно количества входных токенов n и квадратично возрастает относительно размера эмбеддинга d.

Член O(nd²) важен при d > n (например, в LLaMa, n=2K и d=4K).
Член O(n²d) важен при n > d (например, при обучении MosaicML, когда n=65K и d=4K).

Просто напомню, как лихо растёт квадратичная функция:
2 000² = 4 000 000, 100 000² = 10 000 000 000.

Приведу пример, как такое квадратичное усложнение влияет на стоимость обучения модели. Оценочный бюджет обучения LLaMa составлял ~$3M, при 65 млрд. параметров, длине контекста 2K и размере эмбеддинга 4K. Требуемое примерное время тратится в основном на обучение с GPU. Если увеличить длину контекста с 2K до 100K (в 50 раз), то и время обучения вырастет ~50 раз (поскольку окно контекста увеличится, нам потребуется меньше итераций, но на каждую итерацию будет уходить больше времени). Таким образом, обучение LLaMA на окне контекста в 100K обойдётся примерно в $150M.

Немного подробнее расскажу об этих вычислениях:

Поскольку количество токенов равно n, сложность внимания составляет O(n²d + nd²), и на обучение требуется M итераций. Если увеличить длину контекста по формуле n → p*n, то потребуется M/p итераций – опять же, поскольку увеличится длина контекста (давайте для простоты предположим, что рост здесь линейный, но эта оценка может оказаться завышенной или заниженной в зависимости от задачи). Теперь имеем 2 уравнения:
(1) Сложность для 
n ~M * (n²d + nd²)
(2) Сложность для p*n ~ M/p * ((p*n)²d + (p*n)d²)
После ряда упрощений и операций деления, получаем соотношение (2)/(1) ~(d + p*n)/(d + n)

Если d << n, то увеличение n в p раз даст ~ в p раз больше итераций.
Если 
d ~ n, то увеличение n в p раз даст ~ в p/2 раз больше итераций.

Разница между стадиями обучения и вывода при работе с трансформером

Прежде, чем углубиться в изучение приёмов оптимизации, осталось обсудить ещё одну вещь: чем отличаются вычисления на этапе обучения и на этапе вывода.

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

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

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

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

[Фокус #1] Улучшенная позиционная кодировка — ALiBi

При обучении трансформера с большим контекстным окном можно организовать процесс в два этапа: сначала обучить базовую модель на контекстном окне в 2K токенов, а потом дообучить (тонко настроить) на более длинных контекстах (например,  65K). Но выше я уже сказала, что с оригинальной архитектурой трансформеров такой подход неприменим. Почему?

Всё дело в позиционной синусоидальной кодировке, которая не экстраполируется. Авторы статьи ALiBI[4] продемонстрировали, что позиционная синусоидальная кодировка – не самый надёжный инструмент для расширения контекстного окна в период вывода. Стоит добавить несколько токенов, как производительность начинает портиться. Поэтому непригодность к такой «экстраполяции»  в принципе означает, что вы не можете рассчитывать на увеличение длины контекста при выводе/тонкой настройке, хотя при обучении такое увеличение возможно. Термин «экстраполяция» и сравнение различных вариантов позиционной кодировки даны в [4].

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

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

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

Этот приём позволяет ускорить обучение

При вычислении баллов внимания для каждой головы методом ALiBi к каждому баллу внимания (qi · kj , влево) добавляется постоянный сдвиг (вправо). Затем к этим баллам применяется функция softmax, как это делается в неизменённом слое внимания. Остальные этапы вычислений проходят без изменений. m – это скаляр, специфичный для каждой головы; он устанавливается на этапе обучения, но сам при обучении не изменяется. Из статьи об ALiBI.
При вычислении баллов внимания для каждой головы методом ALiBi к каждому баллу внимания (qi · kj , влево) добавляется постоянный сдвиг (вправо). Затем к этим баллам применяется функция softmax, как это делается в неизменённом слое внимания. Остальные этапы вычислений проходят без изменений. m – это скаляр, специфичный для каждой головы; он устанавливается на этапе обучения, но сам при обучении не изменяется. Из статьи об ALiBI.

[Фокус #2] Разреженное внимание

В контекстном окне на 100K не все токены одинаково релевантны. Чтобы сократить объём вычислений, можно  при вычислении баллов внимания учитывать лишь некоторые токены. Цель – проредить таким образом выборку, чтобы вычисления усложнялись относительно n линейно, а не квадратично. Есть несколько подходов, позволяющих выбрать принцип связи между токенами, а сам этот подход отлично проиллюстрирован в следующем посте от Google:

Совокупное внимание можно трактовать как полный граф. Из статьи Sparse Attention Methods
Совокупное внимание можно трактовать как полный граф. Из статьи Sparse Attention Methods
Методы разреженного внимания

Например, внимание со скользящим окном (также именуемое локальным) – это подход, при котором применяется окно внимания с фиксированной длиной, объемлющее каждый токен. При таком паттерне внимания, если мы имеем окно с фиксированным размером w, каждый токен следит за w/2 токенами по обе стороны от себя. Вычислительная сложность такого паттерна O(n*w), она масштабируется линейно при работе с входной последовательностью длины n. Для обеспечения эффективности этого метода w должно быть невелико в сравнении с n. Фокус в том, что информация о внимании «проплывает» всё контекстное окно с захватом близлежащих токенов и аппроксимирует полный граф.   

Метод оценки внимания BigBird сочетает глобальные, локальные и случайные механизмы. Авторы статьи об этом методе сделали важнейшее наблюдение о неотъемлемом конфликте между тем, насколько мало метрик схожести приходится вычислять, и тем, насколько активен поток информации между разными узлами. (Речь о том, насколько один токен может влиять на другие и все токены друг на друга).

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

[Фокус #3] FlashAttention — эффективная реализация слоя внимания для работы на GPU

В слое внимания есть несколько операций, которые приходится повторять снова и снова:

1.       S = Q*K

2.      = softmax(S)

3.      O = P*V

Запомните, что здесь упоминаются буквы PS и O; с ними будем работать позже. Авторы метода FlashAttention (мгновенное внимание) «сплавили» эти операции: реализовали такой алгоритм слоя  внимания, при котором активно задействуется память GPU и вычисляется точное значение внимание.

Чтобы GPU мог выполнить операцию над входными данными, эти данные должны быть представлены в «легкодоступной» памяти, именуемой «статическая память с произвольным доступом» (SRAM). Данные копируются из «медленной» HBM (памяти с высокой пропускной способностью) в SRAM, а затем возвращаются в HBM, как только вычисления будут закончены. SRAM-память работает гораздо быстрее, чем HBM-память, но при этом намного меньше последней (20 МБ против 40ГБ на A100 40 ГБ GPU).

Иерархия памяти в A100 GPU из статьи о FlashAttention
Иерархия памяти в A100 GPU из статьи о FlashAttention

Итак, обращение к HBM – это дорогостоящая операция

Основная проблема, возникающая в слое внимания и касающаяся задействования памяти   GPU – это работа с «промежуточными» результатами умножения P, S и O, которые обычно велики (n, n). Приходится сохранять их в HBM и снова считывать между актами применения внимания.  Перенос P, S и O из HBM в SRAM и обратно – это узкое место, устранить которое и пытались авторы упоминаемой статьи.

Главная идея, на которой основан алгоритм мгновенного внимания – разделить входные значения Q, матрицы K и V поблочно, загружая эти блоки из HBM в SRAM, а затем вычисляя результирующее внимание относительно этих блоков. Эта процедура называется замощением.

Слева: В алгоритме FlashAttention используется замощение, чтобы не возникла слишком большая матрица внимания n × n (на рисунке обозначена прерывистой линией) в HBM. Во внешнем цикле (красные стрелки) FlashAttention перебирает блоки матриц K и V и загружает их в SRAM. В каждом блоке FlashAttention перебирает блоки матрицы Q (голубые стрелки), загружает их в SRAM и пишет результат вычисления внимания обратно в HBM. Справа: ускорение в 7,6 раз. Из статьи FlashAttention
Слева: В алгоритме FlashAttention используется замощение, чтобы не возникла слишком большая матрица внимания n × n (на рисунке обозначена прерывистой линией) в HBM. Во внешнем цикле (красные стрелки) FlashAttention перебирает блоки матриц K и V и загружает их в SRAM. В каждом блоке FlashAttention перебирает блоки матрицы Q (голубые стрелки), загружает их в SRAM и пишет результат вычисления внимания обратно в HBM. Справа: ускорение в 7,6 раз. Из статьи FlashAttention

Операция “перемножение матриц” уже оптимизирована для выполнения на GPU. Можно считать, что этот алгоритм мгновенного внимания реализует «слой внимания» в оптимизированном для GPU виде. Авторы «объединили» несколько этапов умножения, а также softmax и замощение, попутно оптимизировав обращение к HBM.

Хороший обзор работы FlashAttention дан в статье по этому алгоритму.

Совсем недавно мгновенное внимание встроили в PyTorch 2.0. Авторы FlashAttention также предлагают реализацию этого алгоритма на языке Triton.

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

[Фокус #4] многозапросное внимание (MQA)

При оригинальном подходе с многоголовым вниманием (MHA) в каждой голове предусматривался отдельный линейный слой для матриц K и V.

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

Многозапросное внимание (MQA) – это оптимизация, при которой предлагается совместно использовать веса при работе всех голов, в то же время, линейно проецируя K и V. Поэтому нам требуется хранить всего две матрицы размером (n, k) и (n, v). В большой модели (например, GPT-3) может быть до 96 голов, а это значит, что при применении MQA можно 96-кратно сократить потребление памяти в кэше декаодера, где хранятся ключи и значения.

Такая оптимизация особенно полезна при генерации длинных текстов. Например, если мы работаем с широким контекстным окном и запрашиваем длинный осмысленный анализ или резюмирование.

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

[Фокус #5] Условные вычисления

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

В разделе о разреженном внимании мы уже обсуждали, что некоторые токены важнее других. Следуя той же логике, авторы статьи CoLT5 разделили на две ветки все вычисления прямого распространения и внимания: получилась тяжёлая ветка и лёгкая ветка. Лёгкие слои применяются ко всем токенам, а тяжёлые – только к важным.

«Лёгкие и тяжёлые ветки прямого распространения отличаются только в том, какие в них есть скрытые размерности. В лёгких ветках таких размерностей меньше, чем в стандартном слое прямого распространения T5, а в тяжёлых ветках – больше».

Как показывает практика, этот подход оказывается быстрее и точнее, чем уже существующая модель LongT5, если речь идёт об обработке исключительно длинных последовательностей до 64K входных токенов.

Обзор трансформерного слоя COLT5 при применении условных вычислений. Все токены обрабатываются с применениём лёгких слоёв внимания и MLP, тогда как q-направляемые токены запросов обрабатываются с применением более тяжеловесного внимания, чем v-направляемые токены ключей-значений. При этом m-направляемые токены обрабатываются более тяжёлыми MLP. Из статьи CoLT5
Обзор трансформерного слоя COLT5 при применении условных вычислений. Все токены обрабатываются с применениём лёгких слоёв внимания и MLP, тогда как q-направляемые токены запросов обрабатываются с применением более тяжеловесного внимания, чем v-направляемые токены ключей-значений. При этом m-направляемые токены обрабатываются более тяжёлыми MLP. Из статьи CoLT5

[Фокус #6] Увеличение оперативной памяти в GPU

Это не фокус, а необходимость. Чтобы уместить большой контекст, вам нужно много оперативной памяти в GPU, поэтому и используются GPU A100 на 80 ГБ.

Заключение

Уф, длинная статья получилась. Честно, я это не нарочно.

Надеюсь, она получилась полезной. Я сама многое узнала, работая над ней (и вы, надеюсь, тоже), и теперь могу хотя бы предположить, как эти большие языковые модели с миллиардами параметров обучались на беспрецедентно широких контекстных окнах величиной  65-100K токенов.

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

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

Ссылки на основные источники

[1] Introducing 100K Context Windows от Antropic
[2] MPT-7B от MosaicML
[3] Palm-2 Technical report от Google
[4] ALiBI: Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation
[5] FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
[6] Multi-Query attention: Fast Transformer Decoding: One Write-Head is All You Need
[8] Attention is All You Need
[9] Video on Positional Sinusoidal Embedding
[10] Overview of the FlashAttention paper
[11] Sliding Window Attention
[12] Constructing Transformers For Longer Sequences with Sparse Attention Methods
[13] FlashAttention: реализация на языке Triton 
[14] How to Accelerate HuggingFace Throughput by 193% with Triton and ClearML
[15] ClearML Serving
[16] Analyzing the Pros and Cons of NVIDIA Triton Inference Server vs. Other Inference Engines
[17] COLT5: Faster Long-Range Transformers with Conditional Computation
[18] LongT5: Efficient Text-To-Text Transformer for Long Sequences
[19] PaLM
[20] Механизм внимания BigBird 

Комментарии (0)