Расчёт весов внимания: softmax(QKᵀ)·V вручную (запускаемый пример)

Сердце трансформера — одна формула. В этом уроке мы посчитаем self-attention руками на трёх токенах и увидим каждый шаг вживую.

Self-attention вычисляется как softmax(QKᵀ / √d) · V: похожести запросов и ключей превращаются в веса, которыми усредняются значения.

Формула по шагам

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

Похожести. Скалярное произведение его query с key каждого токена: QKᵀ. Чем сильнее совпадение, тем больше число.
Масштабирование. Делим на √d (корень из размерности ключа), чтобы числа не «разъезжались» при большой размерности.
Softmax. Превращаем похожести в веса внимания — положительные, в сумме равные 1.
Взвешенная сумма. Складываем value всех токенов с этими весами — это и есть выход внимания для нашего токена.

Шаг softmax отдельно

Поскольку softmax — ключевая операция, посмотрим на неё изолированно. Она превращает произвольные числа (логиты) в вероятности.

import math

def softmax(scores):
    # Вычитаем максимум для численной устойчивости (результат не меняется).
    m = max(scores)
    exps = [math.exp(s - m) for s in scores]
    total = sum(exps)
    return [e / total for e in exps]

raw = [2.0, 1.0, 0.1]
probs = softmax(raw)

print("Сырые числа (логиты):", raw)
print("После softmax:")
for s, p in zip(raw, probs):
    print(f"  логит {s:>4} -> вероятность {p:.3f}")
print("Сумма вероятностей:", round(sum(probs), 5))

Вывод:

Сырые числа (логиты): [2.0, 1.0, 0.1]
После softmax:
  логит  2.0 -> вероятность 0.659
  логит  1.0 -> вероятность 0.242
  логит  0.1 -> вероятность 0.099
Сумма вероятностей: 1.0

Большее число получает большую долю, но и маленькие не обнуляются полностью, а сумма всегда равна 1. Именно так похожести «запрос-ключ» становятся весами внимания.

Считаем внимание целиком

Теперь полный расчёт. Три токена «кошка», «пьёт», «молоко», у каждого вручную заданы Q, K, V размерности 2. Запустите и проследите веса и выход для каждого токена.

import math

# Три токена, каждый описан вектором Q, K, V размерности 2.
# В настоящей модели Q,K,V получаются умножением эмбеддинга на обучаемые матрицы.
tokens = ["кошка", "пьёт", "молоко"]
Q = [[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]]
K = [[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]]
V = [[10.0, 0.0], [0.0, 10.0], [5.0, 5.0]]

def dot(a, b):
    return sum(x * y for x, y in zip(a, b))

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]

d_k = len(Q[0])  # размерность ключа, для масштабирования

# Считаем выход внимания для КАЖДОГО токена-запроса.
for qi, token in enumerate(tokens):
    # 1. Похожесть запроса этого токена со всеми ключами: QK^T / sqrt(d_k)
    scores = [dot(Q[qi], K[kj]) / math.sqrt(d_k) for kj in range(len(tokens))]
    # 2. Превращаем в веса внимания (сумма = 1)
    weights = softmax(scores)
    # 3. Берём взвешенную сумму значений V
    out = [0.0, 0.0]
    for kj in range(len(tokens)):
        out[0] += weights[kj] * V[kj][0]
        out[1] += weights[kj] * V[kj][1]
    pretty = ", ".join(f"{tokens[j]}={weights[j]:.2f}" for j in range(len(tokens)))
    print(f"Токен '{token}':")
    print(f"  веса внимания -> {pretty}")
    print(f"  выход         -> [{out[0]:.2f}, {out[1]:.2f}]")

Вывод:

Токен 'кошка':
  веса внимания -> кошка=0.40, пьёт=0.20, молоко=0.40
  выход         -> [6.02, 3.98]
Токен 'пьёт':
  веса внимания -> кошка=0.20, пьёт=0.40, молоко=0.40
  выход         -> [3.98, 6.02]
Токен 'молоко':
  веса внимания -> кошка=0.25, пьёт=0.25, молоко=0.50
  выход         -> [5.00, 5.00]

Разберём вывод. Для токена «кошка» query совпал с ключами «кошка» и «молоко» сильнее, чем с «пьёт», поэтому веса 0.40 / 0.20 / 0.40 — и выход стал смесью их значений: [6.02, 3.98]. У каждого токена своя картина внимания. Заметьте: это self-attention — query, key и value берутся из одной и той же последовательности, токены смотрят сами на себя.

Почему делим на √d

При большой размерности скалярные произведения становятся большими по модулю, softmax «насыщается» (почти весь вес уходит одному токену), и обучение портится. Деление на √d держит числа в разумном диапазоне. Это маленькая, но важная деталь устойчивости.

Матрично и параллельно

В коде мы шли по токенам в цикле — ради наглядности. На деле всё считается одной матричной операцией softmax(QKᵀ/√d)·V сразу для всех токенов. Именно это и делает трансформер быстрым: вся последовательность обрабатывается параллельно, без пошагового прохода как в RNN.

Итог

Веса внимания = softmax от масштабированных скалярных произведений query и key.
Выход = взвешенная сумма value с этими весами.
Деление на √d удерживает softmax от насыщения и стабилизирует обучение.
В self-attention Q, K, V берутся из одной последовательности; всё считается одной матричной операцией.