Визуализация механизма внимания

Зачем смотреть на внимание

Нейросеть — это миллионы чисел, и понять, что она делает, трудно. У трансформеров есть удобное окно внутрь: веса внимания. Для каждой головы каждого слоя это матрица T \times T, которая прямо говорит, откуда каждый токен берёт информацию. Если нарисовать её, часто видно, какой алгоритм нашла модель.

Два основных способа рисовать:

• тепловая карта: строки — запросы (кто смотрит), столбцы — ключи (на кого), яркость — вес; каждая строка в сумме даёт 1;
• диаграмма связей (как в инструменте BertViz): токены в два столбца, линии от запроса к ключам, толщина линии — вес.

В этом уроке мы обучим несколько маленьких трансформеров и посмотрим, что у них внутри.

Головы делят работу

Задача: для каждой позиции выдать сумму соседних токенов по модулю 10 (последовательность замкнута в кольцо): y_i = (x_{i-1} + x_{i+1}) \bmod 10. Одной позиции нужно посмотреть сразу в две стороны. Дадим модели слой внимания с двумя головами:

import numpy as np
import nanotorch as nt
import vis

T, V, D = 10, 10, 32
rng = np.random.default_rng(0)
def make_data(n):
    X = rng.integers(0, V, (n, T))
    return X, (np.roll(X, 1, axis=1) + np.roll(X, -1, axis=1)) % V     # x[i-1] + x[i+1] по кругу
X_train, Y_train = make_data(3000)
X_val, Y_val = make_data(500)

class Model(nt.Module):
    def __init__(self):
        self.tok = nt.Embedding(V, D)
        self.pos = nt.Embedding(T, D)
        self.attn = nt.MultiHeadAttention(D, num_heads=2)
        self.head = nt.Sequential(nt.Linear(D, 128), nt.ReLU(), nt.Linear(128, V))

    def forward(self, ids):
        h = self.tok(ids) + self.pos(np.arange(T))
        return self.head(h + self.attn(h))

nt.manual_seed(0)
model = Model()
opt = nt.Adam(model.parameters(), lr=0.01)
for epoch in range(8):
    perm = rng.permutation(len(X_train))
    for i in range(0, len(X_train), 64):
        b = perm[i:i + 64]
        loss = nt.cross_entropy(model(X_train[b]).reshape(-1, V), Y_train[b].reshape(-1))
        opt.zero_grad()
        loss.backward()
        opt.step()
pred = model(X_val).data.argmax(axis=-1)
print(f"точность: {(pred == Y_val).mean():.1%}")

x = X_val[0]
labels = [f"{t}·{p}" for p, t in enumerate(x)]            # «токен·позиция», чтобы различать одинаковые цифры
A = model.attn.attn[0]                                    # (головы, T, T)
for h in range(2):
    vis.heatmap(A[h], xlabels=labels, ylabels=labels, title=f"Голова {h}", fmt=".1f")
vis.attention(labels, A[0], title="Голова 0: связи")

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

Внимание по содержимому

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

import numpy as np
import nanotorch as nt
import vis

T, V, D = 10, 12, 32
rng = np.random.default_rng(0)
def make_data(n):
    X = rng.integers(0, V, (n, T))
    Y = ((X[:, :, None] == X[:, None, :]).sum(axis=-1) > 1).astype(int)   # 1, если есть такой же токен
    return X, Y
X_train, Y_train = make_data(3000)
X_val, Y_val = make_data(500)

class Model(nt.Module):
    def __init__(self):
        self.tok = nt.Embedding(V, D)
        self.attn = nt.MultiHeadAttention(D, num_heads=1)
        self.head = nt.Sequential(nt.Linear(D, 64), nt.ReLU(), nt.Linear(64, 2))

    def forward(self, ids):
        h = self.tok(ids)
        return self.head(h + self.attn(h))

nt.manual_seed(0)
model = Model()
opt = nt.Adam(model.parameters(), lr=0.003)
for epoch in range(8):
    perm = rng.permutation(len(X_train))
    for i in range(0, len(X_train), 64):
        b = perm[i:i + 64]
        loss = nt.cross_entropy(model(X_train[b]).reshape(-1, 2), Y_train[b].reshape(-1))
        opt.zero_grad()
        loss.backward()
        opt.step()
pred = model(X_val).data.argmax(axis=-1)
print(f"точность: {(pred == Y_val).mean():.1%}  (если всегда отвечать «есть повтор» — {Y_val.mean():.0%})")

x = X_val[0]
print("токены:  ", x)
print("повторы: ", Y_val[0])
labels = [str(t) for t in x]
vis.heatmap(model.attn.attn[0, 0], xlabels=labels, ylabels=labels, title="Токены смотрят на такие же токены", fmt=".2f")

Яркие клетки стоят там, где токены одинаковые: каждый токен ищет свои копии. Если копия нашлась, вес делится между ней и самим токеном; если нет — токен смотрит в основном на себя, и по этому классификатор отличает одно от другого. Именно так — через совпадение запросов и ключей — языковые модели находят в тексте слова, связанные с текущим.

Как узор появляется при обучении

Вернёмся к развороту последовательности из прошлого урока и посмотрим на карту внимания в разные моменты обучения:

import numpy as np
import nanotorch as nt
import vis

T, V, D = 8, 10, 32
rng = np.random.default_rng(0)
X_train = rng.integers(0, V, (2000, T))
Y_train = X_train[:, ::-1].copy()
x_show = rng.integers(0, V, (1, T))

class Reverser(nt.Module):
    def __init__(self):
        self.tok = nt.Embedding(V, D)
        self.pos = nt.Embedding(T, D)
        self.attn = nt.MultiHeadAttention(D, num_heads=1)
        self.head = nt.Sequential(nt.LayerNorm(D), nt.Linear(D, V))

    def forward(self, ids):
        h = self.tok(ids) + self.pos(np.arange(T))
        return self.head(h + self.attn(h))

nt.manual_seed(0)
model = Reverser()
opt = nt.Adam(model.parameters(), lr=0.003)
snapshots, titles, losses = [], [], []
step = 0
for epoch in range(3):
    perm = rng.permutation(len(X_train))
    for i in range(0, len(X_train), 32):
        if step in (0, 10, 20, 30, 60):
            model(x_show)
            snapshots.append(model.attn.attn[0, 0].copy())
            titles.append(f"шаг {step}")
        b = perm[i:i + 32]
        loss = nt.cross_entropy(model(X_train[b]).reshape(-1, V), Y_train[b].reshape(-1))
        opt.zero_grad()
        loss.backward()
        opt.step()
        losses.append(loss.data.item())
        step += 1

vis.images(snapshots, titles=titles, cols=5)
vis.line(losses, title="Потери при обучении", xlabel="шаг")

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

Несколько слоёв: attention rollout

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

import numpy as np
import vis

T = 6
tokens = ["я", "видел", "кота", "он", "спал", "."]
# слой 1: каждый токен смотрит на предыдущий
A1 = np.zeros((T, T)); A1[0, 0] = 1
for i in range(1, T):
    A1[i, i - 1] = 1
# слой 2: «он» смотрит на «кота» через позицию после него
A2 = np.eye(T); A2[3] = 0; A2[3, 3] = 0.2; A2[3, 2] = 0.8

def rollout(attns):
    R = np.eye(T)
    for A in attns:
        A = 0.5 * A + 0.5 * np.eye(T)
        A = A / A.sum(axis=1, keepdims=True)
        R = A @ R
    return R

vis.heatmap(A2, xlabels=tokens, ylabels=tokens, title="Слой 2 сам по себе", fmt=".2f")
vis.heatmap(rollout([A1, A2]), xlabels=tokens, ylabels=tokens, title="Rollout двух слоёв", fmt=".2f")

Слой 2 «смотрит» с «он» на позицию «кота», но на этой позиции после слоя 1 лежит смесь «кота» и «видел». Rollout показывает это: у строки «он» вес распределён между «он», «кота» и «видел».

Осторожно с выводами

Карты внимания — полезная подсказка, но не доказательство того, почему модель так ответила:

• Важно не только куда смотрят, но и что передаётся. Большой вес на токене, у которого значение v почти нулевое, ничего не значит.
• После внимания идут FFN и остаточные связи, а слоёв и голов десятки — информация многократно перемешивается.
• Исследования показали, что веса внимания часто можно заметно изменить, почти не изменив ответ модели.

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

Итог

• Веса внимания — матрица T \times T на каждую голову; строка — куда смотрит токен, сумма строки 1.
• Рисуют тепловыми картами и диаграммами связей; полезно подписывать токены вместе с позициями.
• Головы сами делят работу; внимание бывает позиционным (соседи, разворот) и по содержимому (одинаковые токены).
• Узор внимания часто возникает скачком вместе с резким падением потерь.
• Для нескольких слоёв используют attention rollout; веса внимания — подсказка, а не полное объяснение.