Механизм внимания. Attention

Смысл зависит от соседей

Прочитайте два предложения:

• «Он открыл дверь ключом».
• «Из-под камня бил холодный ключ».

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

Это и есть внимание (attention). Впервые его применили в машинном переводе (Бахданау, Чо, Бенжио, 2014): чтобы перевести очередное слово, модель «смотрит» на подходящие слова исходной фразы. А в 2017 году статья «Attention Is All You Need» показала, что внимание может быть единственным способом связывать слова, — так появился трансформер.

Внимание как мягкий поиск

Представьте словарь (dict в Python): по ключу находим значение. Запрос d["кот"] даёт значение, ключ которого точно совпадает с запросом.

Внимание — «мягкий» словарь:

1. Есть запрос q (query) — «что я ищу».
2. Есть пары ключ k_j (key) — «что у меня есть» и значение v_j (value) — «что я отдам».
3. Сравниваем запрос с каждым ключом: оценка s_j = q \cdot k_j (чем больше, тем ключ «подходящее»).
4. Превращаем оценки в веса через softmax: \alpha = \text{softmax}(s), веса положительны и в сумме дают 1.
5. Ответ — взвешенное среднее значений: \text{out} = \sum_j \alpha_j v_j.

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

import numpy as np
import vis

# ключи — «о чём» каждая запись, значения — что запись отдаёт
keys = np.array([[1.0, 0.0],      # «фрукты»
                 [0.9, 0.3],      # «ягоды»
                 [0.0, 1.0],      # «машины»
                 [0.1, 0.9]])     # «велосипеды»
values = np.array([[1.0, 0.0, 0.0, 0.0],
                   [0.0, 1.0, 0.0, 0.0],
                   [0.0, 0.0, 1.0, 0.0],
                   [0.0, 0.0, 0.0, 1.0]])
names = ["фрукты", "ягоды", "машины", "велосипеды"]

def softmax(s):
    e = np.exp(s - s.max())
    return e / e.sum()

for q, desc in [([3.0, 0.0], "про еду"), ([0.0, 3.0], "про транспорт"), ([1.5, 1.5], "непонятно про что")]:
    scores = keys @ np.array(q)
    alpha = softmax(scores)
    print(f"запрос {desc:18}: оценки {scores.round(2)}, веса {alpha.round(2)}")
    print(f"{'':27}ответ = {(alpha @ values).round(2)}")

Запрос «про еду» берёт почти всё из записей «фрукты» и «ягоды», запрос «про транспорт» — из «машин» и «велосипедов», а неопределённый запрос — понемногу из всех.

Зачем делить на \sqrt{d}

Если компоненты q и k — случайные числа со средним 0 и дисперсией 1, то q \cdot k = \sum_{i=1}^d q_i k_i — сумма d слагаемых, и её дисперсия равна d. При d = 512 оценки будут порядка \pm 20, softmax от таких чисел почти one-hot — один вес около 1, остальные около 0. А там, где softmax насыщен, его градиенты почти нулевые, и обучение встаёт.

import numpy as np

rng = np.random.default_rng(0)
for d in (4, 64, 512):
    for scaled in (False, True):
        spreads, maxw = [], []
        for _ in range(200):                                   # усредним по 200 случайным запросам
            q, K = rng.normal(size=d), rng.normal(size=(10, d))
            s = K @ q / (np.sqrt(d) if scaled else 1)
            w = np.exp(s - s.max()); w /= w.sum()
            spreads.append(s.std()); maxw.append(w.max())
        print(f"d = {d:3}, {'с делением на √d' if scaled else 'без деления     '}: разброс оценок {np.mean(spreads):6.2f}, "
              f"максимальный вес в среднем {np.mean(maxw):.2f}")

Деление на \sqrt{d} возвращает разброс оценок к единице при любой размерности. Отсюда название — scaled dot-product attention.

Много запросов сразу: матричная форма

Обычно запросов много (например, по одному на каждое слово). Сложим запросы в матрицу Q формы (m, d), ключи — в K (n, d), значения — в V (n, d_v). Тогда все оценки — одно матричное умножение, а softmax применяется к каждой строке:

Строка i матрицы весов — распределение внимания i-го запроса по всем ключам. Её удобно рисовать тепловой картой.

import numpy as np
import nanotorch as nt

rng = np.random.default_rng(0)
Q = nt.tensor(rng.normal(size=(3, 8)))      # 3 запроса
K = nt.tensor(rng.normal(size=(5, 8)))      # 5 ключей
V = nt.tensor(rng.normal(size=(5, 4)))      # 5 значений размерности 4
out, weights = nt.scaled_dot_product_attention(Q, K, V)
print("веса:", weights.shape, "\n", weights.data.round(2))
print("суммы строк:", weights.data.sum(axis=1).round(6))
print("выход:", out.shape)

Маски

Иногда на некоторые позиции смотреть нельзя:

• на паддинг — фиктивные токены, которыми дополнили короткие тексты;
• на будущие слова — когда модель генерирует текст слева направо, слово не должно подглядывать в те, что идут после него (это понадобится в GPT).

Решение простое: перед softmax прибавляем к запрещённым оценкам очень большое отрицательное число (-10^9, «минус бесконечность»). После экспоненты там получается 0, и весь вес распределяется между разрешёнными позициями.

import numpy as np
import nanotorch as nt

rng = np.random.default_rng(1)
Q = nt.tensor(rng.normal(size=(4, 8)))
K = nt.tensor(rng.normal(size=(4, 8)))
V = nt.tensor(np.eye(4))

pad = np.array([False, False, False, True])          # последний токен — паддинг
_, w = nt.scaled_dot_product_attention(Q, K, V, mask=np.tile(pad, (4, 1)))
print("маска паддинга:\n", w.data.round(2))

causal = np.triu(np.ones((4, 4), bool), k=1)          # True выше диагонали: будущее закрыто
_, w = nt.scaled_dot_product_attention(Q, K, V, mask=causal)
print("причинная маска (нельзя смотреть вперёд):\n", w.data.round(2))

С причинной маской первое слово видит только себя, второе — первые два и так далее: матрица весов нижнетреугольная.

Пример: выравнивание при переводе

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

import numpy as np
import vis

src = ["the", "black", "cat", "sat", "on", "the", "mat"]
dst = ["чёрная", "кошка", "сидела", "на", "коврике"]
W = np.array([
    [0.02, 0.90, 0.05, 0.01, 0.00, 0.01, 0.01],   # чёрная  ← black
    [0.10, 0.05, 0.82, 0.02, 0.00, 0.00, 0.01],   # кошка   ← cat (и немного the)
    [0.01, 0.01, 0.08, 0.88, 0.01, 0.00, 0.01],   # сидела  ← sat (род — от cat!)
    [0.00, 0.00, 0.00, 0.03, 0.94, 0.02, 0.01],   # на      ← on
    [0.01, 0.00, 0.00, 0.01, 0.04, 0.12, 0.82],   # коврике ← mat
])
vis.attention(src, W, query_tokens=dst, title="Куда смотрит переводчик")
vis.heatmap(W, xlabels=src, ylabels=dst, title="Те же веса тепловой картой")

Модель сама находит, какое слово чему соответствует, хотя её никогда этому не учили напрямую — она училась только переводить. Обратите внимание на «сидела»: чтобы выбрать женский род, полезно заглянуть и на «cat».

Итог

• Внимание позволяет каждому элементу «посмотреть» на другие и собрать из них нужную информацию.
• Запрос сравнивается с ключами, оценки через softmax превращаются в веса, ответ — взвешенное среднее значений.
• Оценки делят на \sqrt{d}, чтобы softmax не насыщался при большой размерности.
• Матричная форма: \text{softmax}(QK^\top/\sqrt{d})\,V; строка весов — распределение внимания одного запроса.
• Маски запрещают смотреть на паддинг или на будущее: к таким оценкам прибавляют -\infty.
• Всё дифференцируемо — запросы, ключи и значения можно обучать. Откуда они берутся — в следующем уроке.