Много-запросное внимание объяснение
Изысканное и полное объяснение для всех интересующихся
Multi-Query Attention (MQA) является типом механизма внимания, который может ускорить скорость генерации токенов в декодере, при этом гарантируя производительность модели.
Он широко используется в эпоху больших языковых моделей, многие такие модели используют MQA, такие как Falcon, PaLM, StarCoder и другие.
Многоголовое Внимание(MHA)
Прежде чем ввести MQA, давайте сначала рассмотрим механизм внимания по умолчанию в трансформере.
Многоголовое внимание является механизмом внимания по умолчанию в модели трансформера, как показано на Рисунке 1:
- Познакомьтесь с mPLUG-Owl2 Многофункциональной моделью основного фонда, которая преобразует мультимодальные модели языка большого объема (MLLM) с участием модальностей
- Тенденции карьеры в области AI Что актуально в мире искусственного интеллекта?
- Резюме Выявление шаблонов на изображениях.
Однако у авторегрессивных языковых моделей, основанных на декодерах трансформера, есть проблема при генерации текста.
Во время обучения у нас есть доступ к истинной последовательности целевых значений и мы можем эффективно использовать параллелизм.
Однако во время вывода каждый запрос на позиции обращается ко всем парам ключ-значение, сгенерированных на данной позиции или раньше. Другими словами, вывод слоя самовнимания на конкретной позиции влияет на генерацию следующего токена. Из-за невозможности выполнения параллельных вычислений, декодирование замедляется.
Ниже приведен процесс декодирования слоя самовнимания в авторегрессивной языковой модели, основанной на декодерах трансформера:
def MHAForDecoder(x, prev_K, prev_V, P_q, P_k, P_v, P_o): q = tf.einsum("bd, hdk−>bhk", x, P_q) new_K = tf.concat([prev_K, tf.expand_dims(tf.einsum ("bd, hdk−>bhk", x, P_k), axis = 2)], axis = 2) new_V = tf.concat([prev_V, tf.expand_dims(tf.einsum("bd, hdv−>bhv", x, P_v), axis = 2)], axis = 2) logits = tf.einsum("bhk, bhmk−>bhm", q, new_K) weights = tf.softmax(logits) O = tf.einsum("bhm, bhmv−>bhv", weights, new_V) Y = tf.einsum("bhv, hdv−>bd", O, P_o) return Y, new_K, new_VПеременные:
- x: входной тензор на текущем шаге, который является шагом m+1, с формой [b, d]
- P_q, P_k: тензоры проекции запроса и ключа, с формой [h, d, k]
- P_v: тензор проекции значения, с формой [h, d, v]
- P_o: изученные линейные проекции, с формой [h, d, v]
- Prev_K: тензор Key с предыдущего шага, с формой [b, h, m, k]
- Prev_V: тензор Value с предыдущего шага, с формой [b, h, m, v]
- new_K: тензор Key с добавлением текущего шага, с формой [b, h, m+1, k]
- new_V: тензор Value с добавлением текущего шага, с формой [b, h, m+1, v]
Размерности:
- m: количество выполненных предыдущих шагов
- b: размер пакета
- d: размерность входного и выходного
- h: количество голов
- k: другая размерность тензоров Q, K
- v: другая размерность тензора V
Многозапросное Внимание(MQA)
Multi-Query Attention является вариацией много-головного внимания.
Подход MQA заключается в том, чтобы сохранить первоначальное количество голов для Q, но иметь только одну голову для K и V. Это означает, что все головы Q используют один набор голов K и V, отсюда и название Multi-Query, как показано на рис. 2:
Код процесса декодирования для MQA в основном такой же, как код для MHA, за исключением того, что буква “h”, представляющая размер голов, удаляется из уравнения tf.einsum для K, V, P_k и P_v:
def MQAForDecoder(x, prev_K, prev_V, P_q, P_k, P_v, P_o): q = tf.einsum("bd, hdk−>bhk", x, P_q) new_K = tf.concat([prev_K, tf.expand_dims(tf.einsum ("bd, dk−>bk", x, P_k), axis = 2)], axis = 2) new_V = tf.concat([prev_V, tf.expand_dims(tf.einsum("bd, dv−>bv", x, P_v), axis = 2)], axis = 2) logits = tf.einsum("bhk, bmk−>bhm", q, new_K) weights = tf.softmax(logits) O = tf.einsum("bhm, bmv−>bhv", weights, new_V) Y = tf.einsum("bhv, hdv−>bd", O, P_o) return Y, new_K, new_VПроизводительность
Насколько MQA действительно улучшает скорость? Давайте взглянем на график результатов, представленный в оригинальной статье:
Из таблицы выше видно, что улучшение скорости MQA на кодировщике не очень значительное, но в декодере оно весьма значительное.
В статье также проведены эксперименты по качеству, которые показывают, что MQA имеет лишь незначительное снижение производительности по сравнению с базовым значением. Для получения более подробной информации обратитесь к статье, ссылка на которую указана в конце данной статьи.
Почему MQA может достичь ускорения вывода?
Более эффективное использование памяти
В MQA размер тензоров ключа и значения составляет b * k и b * v, в то время как в MHA размер ключа и значения составляет b * h * k и b * h * v, где h представляет количество голов.
Меньшая вычислительная сложность
При использовании кэша KV, вычислительные затраты на вычисление тензора Key и Value на каждом шаге MQA составляют 1/h от MHA, где h представляет количество голов.
Резюме
В целом MQA достигает ускорения вывода за счет следующих методов:
- Размер кэша KV сокращается в h раз (количество голов), что означает, что тензоры, которые должны храниться в памяти GPU, также уменьшаются. Сэкономленное пространство можно использовать для увеличения размера пакета, тем самым повышая эффективность.
- Количество данных, считываемых из памяти, уменьшается, что сокращает время ожидания вычислительных устройств и повышает использование вычислительных ресурсов.
- MQA имеет относительно небольшой кэш KV, который может поместиться в кэш (SRAM). MHA, с другой стороны, имеет больший кэш KV, который не может быть полностью сохранен в кэше и должен считываться из памяти GPU (DRAM), что требует времени.
Заключение
Следует отметить, что MQA была предложена в 2019 году, и ее применение в то время было не столь широким. Это связано с тем, что предыдущие модели не требовали учета этих аспектов, например, LSTM требовала только поддержки одного состояния, без необходимости сохранения какого-либо кэша.
Когда впервые был предложен трансформер, он в основном использовался в задачах Seq2Seq, конкретно в моделях Encoder-Decoder. Однако модели были не очень масштабными, и не было особого практического спроса на них, поэтому MQA не привлекло много внимания.
Позже представительная модель BERT, которая также основана на структуре трансформера-кодировщика, сделала прямой пропуск.
Это произошло только тогда, когда недавние крупные модели языка на основе трансформера-декодера, такие как GPT, получили широкое применение, что был обнаружену узкое место вывода. В результате люди пересмотрели фокусы несколько лет назад и обнаружили их очень полезными. Другими словами, это главным образом вызвано практическим спросом на масштабные генеративные модели GPT-стиля.
Наконец, если здесь есть ошибки или упущения, пожалуйста, не стесняйтесь указывать на них.
Ссылки
Статья MQA: Быстрое декодирование трансформера: вам нужна только одна записывающая головка
