Multi-head attention: зачем несколько голов (запускаемый пример)

Одного механизма внимания мало: трансформер запускает несколько «голов» параллельно. Разберём, зачем, и увидим, как две головы смотрят в разные стороны.

Multi-head attention — несколько параллельных механизмов внимания со своими матрицами Q/K/V; их результаты объединяются, чтобы модель улавливала разные типы связей одновременно.

Проблема одной головы

Одно внимание выдаёт для каждого токена одно распределение весов. Но между словами существует много разных связей сразу: синтаксическая (кто к чему относится), смысловая (что чему близко), кореферентная (какое местоимение к кому). Заставлять одну голову совмещать всё это — всё равно что одним взглядом одновременно следить за грамматикой, темой и стилем. Что-то обязательно потеряется.

Решение: много голов параллельно

Поэтому внимание запускают несколько раз параллельно, и у каждой «головы» свои матрицы W_Q, W_K, W_V. Значит, каждая голова строит свои query/key/value и формирует свою картину внимания. Одна может специализироваться на «к какому существительному относится прилагательное», другая — на «какое слово продолжает мысль». Покажем на двух головах, как они расходятся.

import math

tokens = ["банк", "у", "реки"]

def softmax(xs):
    m = max(xs); e = [math.exp(x - m) for x in xs]; s = sum(e)
    return [v / s for v in e]
def dot(a, b): return sum(x*y for x, y in zip(a, b))

# Две "головы" внимания используют РАЗНЫЕ проекции -> разные Q и K,
# поэтому фокусируются на разных токенах. Здесь зададим их вручную.
Q1 = [[3.0, 0.0]]                              # запрос токена "банк" в голове 1
K1 = [[3.0, 0.0], [0.0, 1.0], [0.0, 1.0]]      # голова 1: "банк" тянется к себе
Q2 = [[0.0, 3.0]]                              # запрос токена "банк" в голове 2
K2 = [[0.0, 0.0], [0.0, 1.0], [0.0, 3.0]]      # голова 2: "банк" тянется к "реки"

def head_weights(Q, K):
    scores = [dot(Q[0], K[j]) / math.sqrt(2) for j in range(len(tokens))]
    return softmax(scores)

w1 = head_weights(Q1, K1)
w2 = head_weights(Q2, K2)
print("Токен 'банк', куда смотрит каждая голова:")
print("  голова 1:", ", ".join(f"{tokens[j]}={w1[j]:.2f}" for j in range(len(tokens))))
print("  голова 2:", ", ".join(f"{tokens[j]}={w2[j]:.2f}" for j in range(len(tokens))))
print()
print("Разные головы фокусируются на разных токенах -> модель улавливает")
print("несколько типов связей сразу. Выходы голов затем склеиваются в один.")

Вывод:

Токен 'банк', куда смотрит каждая голова:
  голова 1: банк=1.00, у=0.00, реки=0.00
  голова 2: банк=0.00, у=0.01, реки=0.98

Разные головы фокусируются на разных токенах -> модель улавливает
несколько типов связей сразу. Выходы голов затем склеиваются в один.

Для одного и того же токена «банк» голова 1 смотрит на сам токен, а голова 2 — почти целиком на «реки». В реальной модели таких голов 8, 16 или больше, и каждая по-своему «нарезает» контекст. Это даёт модели богатый, многогранный взгляд на одну и ту же последовательность.

Как головы собираются обратно

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

head_1 = attention(Q1, K1, V1)   # своя картина внимания
head_2 = attention(Q2, K2, V2)
...
head_h = attention(Qh, Kh, Vh)

multi = concat(head_1, head_2, ..., head_h) @ W_O

Важно: общая «стоимость» вычислений почти не растёт — мы делим размерность между головами, а не дублируем её целиком. То есть многоголовость даётся почти бесплатно, а пользы много.

Что показывают исследования

Если посмотреть на реальные обученные модели, разные головы действительно специализируются: одни отслеживают согласование, другие — связь местоимения с существительным, третьи — позицию относительно знаков препинания. Это не запрограммировано — модель сама распределила «обязанности» между головами в ходе обучения.

Итог

• Одна голова внимания даёт лишь одну картину связей — этого мало для языка.
• Multi-head запускает несколько вниманий параллельно, у каждого свои Q/K/V.
• Разные головы специализируются на разных типах связей (синтаксис, смысл, кореференция).
• Выходы голов склеиваются и смешиваются матрицей W_O; стоимость почти не растёт.