КПД

SwitchHead: Accelerating Transformers with Mixture-of-Experts Attention
[Статья] [Код]

Смеси экспертов для MLP блоков были неоднократно с успехом применены в больших языковых моделях (в частности в недавно вышедшей Mixtral, и GPT-4 по слухам тоже ей является). Однако, кроме MLP в трансформере можно добиться ускорения за счет использования за счет части голов в MultiHeadAttention. На больших последовательностях Attention слои требуют много вычислений и памяти, и использование части голов даст некоторую экономию.

Даже если бы статья была анонимной, то по цитированиям Шмидхубера при каждом упоминании Attention, можно было бы угадать одного из авторов)

Метод

Наивный подход был бы следующий - иметь предиктор, который предсказывает скор для каждой головы и берет только с наибольшим скором. Однако, проблема в том, что при авторегрессивной генерации могут активироваться разные эксперты и в итоге придется все равно хранить кэш keys и values на все головы (проблема, правда, нынче лечится MultiQueryAttention).

Авторы предлагают следующее - иметь экспертов внутри каждой головы и активировать только часть экспертов на прямом проходе. В этом по сути и вся суть метода.

Эксперименты

Подход валидируют на TransformerXL, который обучают на С4/peS2o 100к шагов. Не самый типичный setting для 2023. Качество замеряют по перплексии на WikiText103, C4, peS2o.

В бейзлайновом трансформере 2 или 10 голов (с одинаковым числом параметров), в предложенном SwitchHead 2 головы с 5 экспертами. Во всех случаях смесь экспертов имеет столько же параметров, сколько и исходная модель. В разных конфигурациях Q, K, V и O проекции могут быть как разбиты на экспертов, так и нет. Перплексия везде примерно +/- одинаковая и будто бы метод не проигрывает сильно в качестве исходному трансформеру с 10 головами (и уверенно опережает двуглавый). Вопрос в том, правда, насколько хорош бейзлайн.

Сравниваются с работой MoA, при том же качестве их подход гораздо экономичнее по памяти, с экономией в ~4-5 раз, в то время, как у MoA выходит ~1.5-3 раза.

Далее, можно сделать MLP тоже экспертами и получить Switch-All модель.
И Switch-All при примерно той же перплексии расходует снова заметно меньше памяти.

Выводы

Специализация attention в трансформере - полезное направление, однако эксперименты выглядят не слишком убедительно. Нет и замеров ускорения, а лишь сравнение MACs между SwitchHead(All) и бейзлайном. Да и в популярном нынче MultiQueryAttention можно было бы активировать часть query проекций и иметь профит в плане памяти и вычислений.

👍7

www.tgoop.com/quant_prune_distill/154

2.58K viewsDec 23, 2023 at 18:42

SwitchHead: Accelerating Transformers with Mixture-of-Experts Attention
[Статья] [Код]

Смеси экспертов для MLP блоков были неоднократно с успехом применены в больших языковых моделях (в частности в недавно вышедшей Mixtral, и GPT-4 по слухам тоже ей является). Однако, кроме MLP в трансформере можно добиться ускорения за счет использования за счет части голов в MultiHeadAttention. На больших последовательностях Attention слои требуют много вычислений и памяти, и использование части голов даст некоторую экономию.

Даже если бы статья была анонимной, то по цитированиям Шмидхубера при каждом упоминании Attention, можно было бы угадать одного из авторов)

Метод

Наивный подход был бы следующий - иметь предиктор, который предсказывает скор для каждой головы и берет только с наибольшим скором. Однако, проблема в том, что при авторегрессивной генерации могут активироваться разные эксперты и в итоге придется все равно хранить кэш keys и values на все головы (проблема, правда, нынче лечится MultiQueryAttention).

Авторы предлагают следующее - иметь экспертов внутри каждой головы и активировать только часть экспертов на прямом проходе. В этом по сути и вся суть метода.

Эксперименты

Подход валидируют на TransformerXL, который обучают на С4/peS2o 100к шагов. Не самый типичный setting для 2023. Качество замеряют по перплексии на WikiText103, C4, peS2o.

В бейзлайновом трансформере 2 или 10 голов (с одинаковым числом параметров), в предложенном SwitchHead 2 головы с 5 экспертами. Во всех случаях смесь экспертов имеет столько же параметров, сколько и исходная модель. В разных конфигурациях Q, K, V и O проекции могут быть как разбиты на экспертов, так и нет. Перплексия везде примерно +/- одинаковая и будто бы метод не проигрывает сильно в качестве исходному трансформеру с 10 головами (и уверенно опережает двуглавый). Вопрос в том, правда, насколько хорош бейзлайн.

Сравниваются с работой MoA, при том же качестве их подход гораздо экономичнее по памяти, с экономией в ~4-5 раз, в то время, как у MoA выходит ~1.5-3 раза.

Далее, можно сделать MLP тоже экспертами и получить Switch-All модель.
И Switch-All при примерно той же перплексии расходует снова заметно меньше памяти.

Выводы

Специализация attention в трансформере - полезное направление, однако эксперименты выглядят не слишком убедительно. Нет и замеров ускорения, а лишь сравнение MACs между SwitchHead(All) и бейзлайном. Да и в популярном нынче MultiQueryAttention можно было бы активировать часть query проекций и иметь профит в плане памяти и вычислений.

BY КПД