Многоголовное внимание (Multi-Head Attention) на простом примере

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

Исходные данные:

Предложение: "Кот ловит мышь" (3 слова)
Каждое слово представлено вектором размерности 4

import torch

# 3 слова, каждое - вектор размерности 4
x = torch.tensor([
    [0.1, 0.2, 0.3, 0.4],  # "Кот"
    [0.5, 0.6, 0.7, 0.8],  # "ловит"
    [0.9, 1.0, 1.1, 1.2]   # "мышь"
])  # shape: (3, 4) = (seq_len, d_model)

---

Шаг 1: Создаём 2 головы (Multi-Head)

Оригинальное внимание было бы одно на все 4 признака. Мы разобьём на 2 головы по 2 признака каждая:

Исходные векторы (4D):
"Кот":    [0.1, 0.2, 0.3, 0.4]
"ловит":  [0.5, 0.6, 0.7, 0.8]
"мышь":   [0.9, 1.0, 1.1, 1.2]

Голова 1 (первые 2 признака):
"Кот":    [0.1, 0.2]    ← возможно, отвечает за "одушевленность"
"ловит":  [0.5, 0.6]    ← возможно, отвечает за "глагол"
"мышь":   [0.9, 1.0]    ← возможно, отвечает за "одушевленность"

Голова 2 (последние 2 признака):
"Кот":    [0.3, 0.4]    ← возможно, отвечает за "размер"
"ловит":  [0.7, 0.8]    ← возможно, отвечает за "действие"
"мышь":   [1.1, 1.2]    ← возможно, отвечает за "размер"

На практике: Мы не знаем, что означают признаки, и не выбираем вручную. Модель сама учится, как разбивать признаки на осмысленные группы.

---

Шаг 2: Проецируем в пространства Q, K, V для каждой головы

Каждая голова имеет свои весовые матрицы:

# Веса для Головы 1
W_Q1 = torch.randn(4, 2)  # (d_model, d_k) - для Query
W_K1 = torch.randn(4, 2)  # (d_model, d_k) - для Key
W_V1 = torch.randn(4, 2)  # (d_model, d_v) - для Value

# Веса для Головы 2
W_Q2 = torch.randn(4, 2)  # (d_model, d_k) - для Query
W_K2 = torch.randn(4, 2)  # (d_model, d_k) - для Key
W_V2 = torch.randn(4, 2)  # (d_model, d_v) - для Value

# Получаем Q, K, V для каждой головы
Q1 = x @ W_Q1  # (3, 2) - Query для Головы 1
K1 = x @ W_K1  # (3, 2) - Key для Головы 1
V1 = x @ W_V1  # (3, 2) - Value для Головы 1

Q2 = x @ W_Q2  # (3, 2) - Query для Головы 2
K2 = x @ W_K2  # (3, 2) - Key для Головы 2
V2 = x @ W_V2  # (3, 2) - Value для Головы 2

Теперь каждая голова работает с векторами размерности 2 вместо 4.

---

Шаг 3: Вычисляем внимание в каждой голове

Голова 1:

# Attention для Головы 1
attention_scores1 = Q1 @ K1.T  # (3, 3)
# Сравниваем каждый Query со всеми Keys

# После softmax получаем веса внимания
attention_weights1 = torch.softmax(attention_scores1 / (2**0.5), dim=-1)
# (3, 3) - каждая строка: насколько слово смотрит на другие

# Взвешенная сумма Values
output1 = attention_weights1 @ V1  # (3, 2)

Что может "видеть" Голова 1:

Вес внимания для "Кот":
- на "Кот":    0.6    ← смотрит на себя
- на "ловит":  0.3    ← кто делает действие
- на "мышь":   0.1    ← на кого направлено

Голова 1, возможно, учится "одушевленности" или "субъект-объектным отношениям"

Голова 2:

# Attention для Головы 2
attention_scores2 = Q2 @ K2.T  # (3, 3)
attention_weights2 = torch.softmax(attention_scores2 / (2**0.5), dim=-1)
output2 = attention_weights2 @ V2  # (3, 2)

Что может "видеть" Голова 2:

Вес внимания для "ловит":
- на "Кот":    0.4    ← кто делает
- на "ловит":  0.5    ← само действие
- на "мышь":   0.1    ← направление действия

Голова 2, возможно, учится "действиям" или "семантическим ролям"

---

Шаг 4: Объединяем выходы голов

# Собираем выходы обеих голов
combined = torch.cat([output1, output2], dim=-1)  # (3, 4)
# Было: голова1 (3,2) + голова2 (3,2) = (3,4)

# Применяем финальную проекцию
W_O = torch.randn(4, 4)  # (num_heads * d_v, d_model)
final_output = combined @ W_O  # (3, 4)

Теперь у нас снова векторы размерности 4, но обогащенные информацией из обеих голов.

---

Шаг 5: Что мы получили в итоге?

Для каждого слова теперь есть новое представление, которое:
1. От Головы 1: Учитывает, кто на кого смотрит с точки зрения "одушевленности"
2. От Головы 2: Учитывает, кто что делает с точки зрения "действий"

"Кот" после Multi-Head Attention:
[0.1, 0.2, 0.3, 0.4] → [0.15, 0.25, 0.35, 0.45]
↑
Теперь "знает", что он субъект действия "ловит" 
и связан с объектом "мышь"

---

Визуализация процесса:

Исходные слова:
[Кот]   [ловит]   [мышь]
  ↓        ↓        ↓
  ┌─────────────────┐
  │   Multi-Head    │
  │    Attention    │
  └─────────────────┘
       ↓    ↓    ↓
[Кот'] [ловит'] [мышь']

Внутри Multi-Head Attention:
         ┌─────────────┬─────────────┐
         │   Голова 1  │   Голова 2  │
         │  (субъекты) │  (действия) │
         └─────────────┴─────────────┘
                ↓            ↓
          объединяем → проекция

---

Почему это лучше обычного внимания?

Обычное внимание (1 голова):

• Смотрит на ВСЕ признаки сразу
• Может "перепутать" разные аспекты
• Как один человек пытается понять всё сразу

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

• Голова 1: Смотрит на "кто есть кто" (субъект-объект)
• Голова 2: Смотрит на "что происходит" (действия)
• Голова 3: Может смотреть на что-то ещё
• Как команда экспертов, каждый со своей специализацией

---

Аналогия из жизни:

Задача: Проанализировать футбольный матч

• Одно внимание (1 эксперт):
  Один комментатор пытается заметить всё: тактику, технику, судейство, эмоции. Может что-то упустить.

• Многоголовое внимание (несколько экспертов):

• Эксперт 1: Анализирует тактику команд
• Эксперт 2: Смотрит на технику игроков
• Эксперт 3: Оценивает работу судьи
• Эксперт 4: Следит за эмоциями на поле

Потом все мнения объединяются в полный анализ.

---

Код полного Multi-Head Attention:

class SimpleMultiHeadAttention:
    def __init__(self, d_model=4, num_heads=2):
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads  # размерность для каждой головы

        # Весовые матрицы
        self.W_Q = torch.randn(d_model, d_model)  # для Query
        self.W_K = torch.randn(d_model, d_model)  # для Key
        self.W_V = torch.randn(d_model, d_model)  # для Value
        self.W_O = torch.randn(d_model, d_model)  # финальная проекция

    def forward(self, x):
        # 1. Линейные проекции
        Q = x @ self.W_Q  # (3, 4)
        K = x @ self.W_K  # (3, 4)
        V = x @ self.W_V  # (3, 4)

        # 2. Разбиваем на головы
        batch_size, seq_len, _ = x.shape
        Q = Q.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        K = K.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        V = V.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        # Теперь: (batch, num_heads, seq_len, d_k)

        # 3. Attention в каждой голове
        scores = Q @ K.transpose(-2, -1) / (self.d_k ** 0.5)
        attention = torch.softmax(scores, dim=-1)
        head_outputs = attention @ V  # (batch, num_heads, seq_len, d_k)

        # 4. Объединяем головы
        combined = head_outputs.transpose(1, 2).contiguous().view(
            batch_size, seq_len, self.d_model
        )

        # 5. Финальная проекция
        output = combined @ self.W_O
        return output

# Использование
mha = SimpleMultiHeadAttention(d_model=4, num_heads=2)
result = mha.forward(x.unsqueeze(0))  # добавляем batch dimension
print("Результат Multi-Head Attention:", result.shape)  # (1, 3, 4)

Давайте разберем еще раз как формируются головы под другим углом более детально.

Исходные данные:

У нас есть одно предложение из 3 слов, каждое слово — вектор размерности 4 (d_model=4).
Мы хотим 2 головы (num_heads=2), значит каждая голова будет работать с векторами размерности 2 (d_k = 4/2 = 2).

import torch

# Исходные данные: batch_size=1, seq_len=3, d_model=4
# Одно предложение из 3 слов
x = torch.tensor([
    [[0.1, 0.2, 0.3, 0.4],   # слово 1
     [0.5, 0.6, 0.7, 0.8],   # слово 2
     [0.9, 1.0, 1.1, 1.2]]   # слово 3
])  # shape: (1, 3, 4) = (batch, seq_len, d_model)

# После линейных проекций получаем Q, K, V той же размерности
Q = torch.tensor([
    [[1.1, 1.2, 1.3, 1.4],   # Query для слова 1
     [1.5, 1.6, 1.7, 1.8],   # Query для слова 2
     [1.9, 2.0, 2.1, 2.2]]   # Query для слова 3
])  # shape: (1, 3, 4)

---

Шаг 1: Разбиваем Q.view(batch_size, seq_len, num_heads, d_k)

batch_size = 1
seq_len = 3
num_heads = 2
d_k = 2  # d_model / num_heads = 4/2 = 2

# Разбиваем 4D вектор каждого слова на 2 головы по 2D каждая
Q_reshaped = Q.view(batch_size, seq_len, num_heads, d_k)
print("После view:")
print(Q_reshaped)
print("Shape:", Q_reshaped.shape)  # (1, 3, 2, 2)

Что произошло:

Было Q: (1, 3, 4) = [слово1[4D], слово2[4D], слово3[4D]]

Стало Q_reshaped: (1, 3, 2, 2)
[
  [  # batch 0
    [  # слово 0
      [1.1, 1.2],  # Голова 0: первые 2 признака слова 0
      [1.3, 1.4]   # Голова 1: последние 2 признака слова 0
    ],
    [  # слово 1
      [1.5, 1.6],  # Голова 0: первые 2 признака слова 1
      [1.7, 1.8]   # Голова 1: последние 2 признака слова 1
    ],
    [  # слово 2
      [1.9, 2.0],  # Голова 0: первые 2 признака слова 2
      [2.1, 2.2]   # Голова 1: последние 2 признака слова 2
    ]
  ]
]

Проблема: Ось головы (ось 2) находится между seq_len (ось 1) и d_k (ось 3). Для вычисления внимания нам нужно, чтобы ось головы была первой после batch.

---

Шаг 2: Меняем оси местами .transpose(1, 2)

Q_transposed = Q_reshaped.transpose(1, 2)
print("\nПосле transpose(1, 2):")
print(Q_transposed)
print("Shape:", Q_transposed.shape)  # (1, 2, 3, 2)

Что произошло после transpose:

Стало: (1, 2, 3, 2)
[
  [  # batch 0
    [  # Голова 0
      [1.1, 1.2],  # слово 0 для головы 0
      [1.5, 1.6],  # слово 1 для головы 0
      [1.9, 2.0]   # слово 2 для головы 0
    ],
    [  # Голова 1
      [1.3, 1.4],  # слово 0 для головы 1
      [1.7, 1.8],  # слово 1 для головы 1
      [2.1, 2.2]   # слово 2 для головы 1
    ]
  ]
]

Теперь структура идеальна для параллельной обработки:
- Ось 0: batch (1)
- Ось 1: головы (2) ← теперь каждая голова обрабатывается отдельно
- Ось 2: последовательность (3) ← слова для каждой головы
- Ось 3: признаки (2) ← вектор для каждого слова в голове

---

Шаг 3: Почему это нужно?

До transpose:

# Q_reshaped[0, 0] даёт ВСЕ головы для слова 0
print("До transpose: Q_reshaped[0, 0] =", Q_reshaped[0, 0])
# [[1.1, 1.2],  # голова 0 слова 0
#  [1.3, 1.4]]  # голова 1 слова 0
# Это неудобно для параллельных вычислений!

После transpose:

# Q_transposed[0, 0] даёт ВСЕ слова для головы 0
print("После transpose: Q_transposed[0, 0] =", Q_transposed[0, 0])
# [[1.1, 1.2],  # слово 0 головы 0
#  [1.5, 1.6],  # слово 1 головы 0
#  [1.9, 2.0]]  # слово 2 головы 0
# Идеально! Теперь можем обрабатывать каждую голову отдельно!

---

Шаг 4: Визуализация процесса

Было:

Q: (batch, seq, d_model) = (1, 3, 4)
Слова:     Слово1[1.1,1.2,1.3,1.4]  Слово2[1.5,1.6,1.7,1.8]  Слово3[1.9,2.0,2.1,2.2]

После view:

(1, 3, 2, 2)
Для каждого слова: [Голова0[2признака], Голова1[2признака]]

После transpose:

(1, 2, 3, 2)
Голова0: [Слово1[2призн], Слово2[2призн], Слово3[2призн]]
Голова1: [Слово1[2призн], Слово2[2призн], Слово3[2призн]]

Теперь можем вычислять attention параллельно для каждой головы!

---

Шаг 5: Что происходит дальше в attention?

После transpose у нас есть:
- Q: (1, 2, 3, 2) - 2 головы, в каждой 3 слова по 2 признака
- K: (1, 2, 3, 2) - то же самое для Key
- V: (1, 2, 3, 2) - то же самое для Value

Вычисляем attention:

# Для головы 0: Q[0,0] (3,2) @ K[0,0].T (2,3) = (3,3) - матрица внимания!
# Для головы 1: Q[0,1] (3,2) @ K[0,1].T (2,3) = (3,3) - матрица внимания!

# В коде это делается одним вызовом для всех голов:
scores = Q @ K.transpose(-2, -1)  # (1, 2, 3, 3)
# Каждая голова получила свою матрицу внимания 3x3!

---

Шаг 6: Пример с конкретными числами

# Упрощенный пример с меньшими числами
Q_simple = torch.tensor([[
    [[1, 2, 3, 4],    # слово 0
     [5, 6, 7, 8],    # слово 1
     [9, 10, 11, 12]] # слово 2
]])  # (1, 3, 4)

print("Исходный Q:")
print(Q_simple[0])
# tensor([[ 1,  2,  3,  4],
#         [ 5,  6,  7,  8],
#         [ 9, 10, 11, 12]])

# 1. view
Q_viewed = Q_simple.view(1, 3, 2, 2)
print("\nПосле view(1, 3, 2, 2):")
print(Q_viewed[0])
# tensor([[[ 1,  2],   # слово 0, голова 0
#          [ 3,  4]],  # слово 0, голова 1
#         
#         [[ 5,  6],   # слово 1, голова 0
#          [ 7,  8]],  # слово 1, голова 1
#         
#         [[ 9, 10],   # слово 2, голова 0
#          [11, 12]]]) # слово 2, голова 1

# 2. transpose
Q_final = Q_viewed.transpose(1, 2)
print("\nПосле transpose(1, 2):")
print(Q_final[0])
# tensor([[[ 1,  2],   # голова 0, слово 0
#          [ 5,  6],   # голова 0, слово 1
#          [ 9, 10]],  # голова 0, слово 2
#         
#         [[ 3,  4],   # голова 1, слово 0
#          [ 7,  8],   # голова 1, слово 1
#          [11, 12]]]) # голова 1, слово 2

---

Шаг 7: Зачем всё это нужно?

Без этого трюка пришлось бы:
1. Вручную разбивать тензор на головы
2. В цикле обрабатывать каждую голову отдельно
3. Потом собирать обратно

С этим трюком:
1. Одна операция разбивает на головы
2. Матричные операции работают для всех голов параллельно
3. Автоматическая оптимизация на GPU

---

Аналогия:

Представьте, что у вас есть 3 студента (слова), каждый с 4 характеристиками:
- [рост, вес, оценка_математика, оценка_физика]

Вы хотите, чтобы 2 учителя (головы) оценили студентов:
- Учитель 1 (физкультура): смотрит только на [рост, вес]
- Учитель 2 (наука): смотрит только на [оценка_математика, оценка_физика]

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

Решение:
1. view: Разделили характеристики каждого студента на 2 группы
2. transpose: Переставили, чтобы получить:
- Таблица для учителя 1: все студенты, только рост/вес
- Таблица для учителя 2: все студенты, только оценки

---

Q.view(batch, seq, heads, d_k).transpose(1, 2) делает:
1. view: Разбивает каждый d_model-вектор на num_heads частей по d_k элементов
2. transpose: Переставляет оси, чтобы головы обрабатывались параллельно

Результат: Из (batch, seq, d_model) получаем (batch, heads, seq, d_k) — идеальный формат для параллельного вычисления внимания в каждой голове!

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

---

Что там про головы и отдельные признаки

Как делится информация:

Исходные данные:

• d_model = 512 (размер эмбеддинга слова)
• num_heads = 8 (8 голов)
• d_k = d_model / num_heads = 512 / 8 = 64

Делим эмбеддинг на части:

Эмбеддинг слова (512-dim):
[████████████████████████████████████████████████████████████████]
  ↑       ↑       ↑       ↑       ↑       ↑       ↑       ↑
  64      64      64      64      64      64      64      64
  Голова1 Голова2 Голова3 Голова4 Голова5 Голова6 Голова7 Голова8

Каждая голова получает по 64 измерения из исходных 512.

Аналогия:

Представьте, что у вас есть картина (эмбеддинг слова), и 8 экспертов (голов):

Эксперт (голова)	Смотрит на
Эксперт 1	Цвета (красный, синий, зеленый...)
Эксперт 2	Формы (круги, квадраты, линии...)
Эксперт 3	Текстуры (гладкое, шершавое...)
Эксперт 4	Освещение (тени, блики...)
Эксперт 5	Композицию (расположение объектов)
Эксперт 6	Перспективу (глубину)
Эксперт 7	Эмоции (грустное, весёлое)
Эксперт 8	Стиль (реализм, импрессионизм)

Каждый эксперт специализируется на своём аспекте!

На практике в NLP:

Пример: слово "банк"

Эмбеддинг "банка": [финансы, учреждение, деньги, река, берег, ...]

Головы могут специализироваться:
- Голова 1: финансовый контекст
- Голова 2: географический контекст
- Голова 3: синтаксическая роль (подлежащее/сказуемое)
- Голова 4: часть речи (существительное)
- Голова 5: связи с другими словами
- Голова 6: временной контекст
- Голова 7: эмоциональная окраска
- Голова 8: формальная/неформальная стилистика

Математически:

# Эмбеддинг слова: 512-мерный вектор
word_embedding = [x1, x2, x3, ..., x512]

# Делим на 8 голов по 64 измерения:
head1 = [x1, x2, ..., x64]      # 1-64 измерения
head2 = [x65, x66, ..., x128]   # 65-128 измерения
...
head8 = [x449, x450, ..., x512] # 449-512 измерения

Почему это хорошо?

1. Параллелизм:

# Все 8 голов работают ОДНОВРЕМЕННО на GPU
# Вместо одного большого attention (512×512)
# 8 маленьких attention (64×64) параллельно

2. Разные представления:

Каждая голова учится разным аспектам:
- Одна может следить за синтаксисом
- Другая за семантикой
- Третья за прагматикой

3. Интерпретируемость:

Можно посмотреть, что "видит" каждая голова:

# Анализируем attention weights каждой головы
for head_idx in range(num_heads):
    attention_map = attention[0, head_idx]  # для первого батча
    # head 0: смотрит на следующее слово
    # head 1: смотрит на предыдущее слово  
    # head 2: смотрит на главное слово предложения
    # head 3: смотрит на знаки препинания

Пример из реальной модели:

В BERT (12 голов, 768 размерность):
- Голова 0: Смотрит на следующее слово (синтаксис)
- Голова 1: Смотрит на предыдущее слово
- Голова 2: Смотрит на все слова предложения
- Голова 3: Смотрит на знаки препинания
- Голова 4: Связывает подлежащее и сказуемое
- Голова 5-11: Другие лингвистические паттерны

Что будет, если голов мало/много?

Слишком мало голов (1-2):

• Модель "не видит" все аспекты
• Перегруженная голова пытается всё уловить
• Хуже качество

Слишком много голов (64+):

• Каждая голова получает слишком мало измерений (512/64=8)
• Не хватает информации для обучения
• Вычислительно дорого

Золотая середина:

• BERT: 12 голов при 768 измерениях
• GPT-3: 96 голов при 12288 измерениях
• Обычно: 8-16 голов

Важный нюанс:

Головы не жёстко закреплены за конкретными аспектами! Они учатся сами в процессе тренировки. Модель сама решает, какие аспекты выделять каждой голове.

Визуализация работы голов:

# Пример: "The animal didn't cross the street because it was too tired"
# Слово "it" - на что смотрят разные головы?

Голова 0: "it" → "animal"   (94%)  # понял, что "it" = "animal"
Голова 1: "it" → "street"   (2%)   # слабая связь
Голова 2: "it" → "tired"    (3%)   # связь через состояние
Голова 3: "it" → "because"  (1%)   # связь через причину
...

Итог:

1. Берём слова → делим их признаки между несколькими "экспертами" (головами)
2. Каждая голова учится своему аспекту (кто, что, как, где)
3. Каждая голова вычисляет внимание в своём подпространстве
4. Объединяем результаты всех голов
5. Получаем обогащённые представления слов
6. Получает часть эмбеддинга (не весь)
7. Специализируется на своём аспекте (учится во время тренировки)
8. Работает параллельно с другими головами
9. Вместе дают полную картину

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