Визуализация внимания GPT-2: как заставить модель «видеть» Bad Apple через оптимизацию эмбеддингов

Зачем заставлять GPT-2 смотреть Bad Apple?

Представьте, что вы взяли языковую модель, которая никогда в жизни не видела пикселей, и заставили ее "внимание" нарисовать силуэт танцующей девочки из культового видео Bad Apple. Звучит как абсурд. Но именно этот абсурд раскрывает кое-что важное: насколько гибкими могут быть эмбеддинги в замороженной модели.

Большинство статей про визуализацию внимания показывают, куда модель смотрит при генерации текста. Скучно. Мы пойдем другим путем - заставим модель смотреть туда, куда мы хотим, не меняя ни одного веса в трансформере. Только эмбеддинги.

Проблема: GPT-2 думает словами, а нам нужны картинки

GPT-2 обучена предсказывать следующее слово. Её механизм внимания выстроен вокруг лингвистических паттернов. Каждая голова внимания в 12 слоях модели (в GPT-2 Small) ловит свои зависимости: местоимения, синтаксические связи, тематические ассоциации.

А мы хотим, чтобы матрица внимания между 256 токенами вдруг стала похожа на кадр из Bad Apple размером 16x16 пикселей. Как будто просим пианиста, который играет Баха, вдруг отбить ритм драм-н-бейса - на том же самом инструменте.

Ключевое наблюдение: внимание в трансформере вычисляется как softmax(QK^T/√d). Q и K - проекции из эмбеддингов. Если эмбеддинги - это наш единственный рычаг, можем ли мы подобрать такие векторы, чтобы QK^T давала нужную нам картинку?

Решение: Заморозить всё, кроме эмбеддингов

Архитектура проста до безобразия:

1. Берём GPT-2 (или любую другую трансформерную модель) и замораживаем все параметры
2. Создаём learnable эмбеддинги для N токенов (например, 256)
3. Пропускаем их через модель, вытаскиваем матрицы внимания из выбранного слоя и головы
4. Сравниваем с целевым изображением (преобразованным в матрицу)
5. Оптимизируем эмбеддинги через градиентный спуск

Вся магия в функции потерь. Самый наивный подход - MSE между матрицей внимания и целевым изображением - работает плохо. Почему? Потому что attention scores после softmax - это распределение вероятностей. А нам нужно управлять сырыми logits перед softmax.

1 Готовим поле боя: установка и импорты

Нам понадобится PyTorch, transformers и много терпения. Если вы никогда не собирали трансформер с нуля, сначала посмотрите этот разбор - там объяснена каждая матрица.

import torch
import torch.nn as nn
import torch.optim as optim
from transformers import GPT2Model, GPT2Config
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image

2 Загружаем модель и отрубаем градиенты

Берём GPT-2 Small. На 2026 год эта модель - уже классика, как Fortran для научных вычислений. Но принцип работает для любой трансформерной архитектуры.

config = GPT2Config.from_pretrained('gpt2')
model = GPT2Model.from_pretrained('gpt2', config=config)

# Замораживаем ВСЁ
for param in model.parameters():
    param.requires_grad = False

model.eval()  # Переводим в режим оценки

Важный момент: мы не просто выключаем градиенты. Мы гарантируем, что веса не сдвинутся ни на йоту. Это критично - если модель начнёт адаптироваться под наш бред, эксперимент теряет смысл.

3 Создаём обучаемые эмбеддинги

Здесь начинается магия. Мы создаём тензор эмбеддингов для 256 токенов. Почему 256? Потому что 16x16 - это 256 пикселей. Каждый токен будет "отвечать" за один пиксель в итоговой карте внимания.

seq_length = 256
embed_dim = model.config.hidden_size  # 768 для GPT-2 Small

# Инициализируем эмбеддинги
learnable_embeds = nn.Parameter(
    torch.randn(seq_length, embed_dim) * 0.02  # Небольшая дисперсия
)

Инициализация случайная - но уже через несколько итераций оптимизации эти векторы станут чем-то осмысленным (с точки зрения нашей задачи).

4 Готовим цель: Bad Apple в матричном виде

Берём кадр из Bad Apple, преобразуем в чёрно-белый, уменьшаем до 16x16. Нам нужна матрица 16x16, где значения от 0 до 1 представляют яркость пикселя.

def load_target_image(image_path, target_size=16):
    img = Image.open(image_path).convert('L')  # В градации серого
    img = img.resize((target_size, target_size))
    img_array = np.array(img) / 255.0  # Нормализуем к [0, 1]
    return torch.tensor(img_array, dtype=torch.float32)

target_image = load_target_image('bad_apple_frame.png')
target_matrix = target_image.flatten().reshape(1, -1)  # Преобразуем в вектор

5 Пишем функцию потерь: хитрость с QK проекциями

Вот где собака зарыта. Если брать внимание после softmax, оптимизация становится нестабильной - градиенты почти нулевые там, где attention scores близки к 0 или 1. Нужно работать с сырыми logits (QK^T/√d).

def compute_attention_loss(model, embeddings, target, layer_idx=5, head_idx=2):
    """
    Вычисляем loss между QK проекциями и целевым изображением.
    
    Args:
        layer_idx: какой слой трансформера использовать (0-11)
        head_idx: какая голова внимания (0-11)
    """
    # Пропускаем эмбеддинги через модель
    outputs = model(inputs_embeds=embeddings.unsqueeze(0), output_attentions=True)
    
    # Достаём QK проекции ДО softmax
    # В transformers внимание сохраняется как кортеж (layer_output, attentions)
    # Но чтобы получить QK, нужно лезть в сам механизм внимания
    # Для простоты используем готовые attention scores, но работаем с ними до softmax
    
    # Альтернатива: вытащить Q и K напрямую из модели
    layer = model.h[layer_idx]
    attn = layer.attn
    
    # Проецируем эмбеддинги в Q, K, V
    q = attn.c_attn(embeddings)
    q = q[..., :attn.embed_dim]  # Делим на три части
    
    # Здесь нужно разделить q на головы и т.д.
    # Полный код смотрите в репозитории на GitHub
    
    # Для демонстрации используем упрощённый вариант
    attention_scores = outputs.attentions[layer_idx][0, head_idx]
    
    # Берём logits до softmax (они уже поделены на √d в реализации)
    # Сравниваем с целевой матрицей
    loss = torch.nn.functional.mse_loss(attention_scores, target)
    return loss

На самом деле, в реальном коде нужно аккуратно вытащить Q и K проекции до деления на √d и до softmax. Это требует копания во внутренностях Hugging Face реализации. Но идея ясна: оптимизируем сырые скалярные произведения, а не вероятности.

6 Оптимизация: multi-start и хитрые learning rate

Случайная инициализация эмбеддингов может завести в плохой локальный минимум. Делаем 10 разных случайных инициализаций, выбираем лучшую.

def optimize_embeddings_multi_start(model, target, n_starts=10, n_iters=500):
    best_embeds = None
    best_loss = float('inf')
    
    for start in range(n_starts):
        print(f"Запуск {start+1}/{n_starts}")
        
        # Новая случайная инициализация
        embeds = nn.Parameter(torch.randn(seq_length, embed_dim) * 0.02)
        optimizer = optim.Adam([embeds], lr=0.01)
        scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, n_iters)
        
        for i in range(n_iters):
            optimizer.zero_grad()
            loss = compute_attention_loss(model, embeds, target)
            loss.backward()
            optimizer.step()
            scheduler.step()
            
            if i % 50 == 0:
                print(f"  Итерация {i}, loss: {loss.item():.4f}")
        
        if loss.item() < best_loss:
            best_loss = loss.item()
            best_embeds = embeds.data.clone()
    
    return best_embeds, best_loss

CosineAnnealingLR здесь не просто для красоты. Learning rate падает от 0.01 до почти нуля - это помогает "притировать" эмбеддинги в конце оптимизации, не прыгая вокруг минимума.

7 Визуализация: превращаем внимание обратно в картинку

После оптимизации достаём матрицу внимания из выбранного слоя и головы, преобразуем обратно в изображение 16x16.

def visualize_attention(model, embeddings, layer_idx=5, head_idx=2):
    with torch.no_grad():
        outputs = model(inputs_embeds=embeddings.unsqueeze(0), output_attentions=True)
        attention = outputs.attentions[layer_idx][0, head_idx].cpu().numpy()
    
    # Масштабируем к [0, 1] для отображения
    attention = (attention - attention.min()) / (attention.max() - attention.min())
    
    plt.figure(figsize=(8, 8))
    plt.imshow(attention.reshape(16, 16), cmap='viridis')
    plt.colorbar()
    plt.title(f"Attention map: layer {layer_idx}, head {head_idx}")
    plt.show()

Что пошло не так: нюансы, о которых молчат в учебниках

Первый же запуск покажет, что всё не так просто. Вот какие подводные камни ждут:

Самая неприятная проблема - это интерференция голов внимания. Вы оптимизируете одну голову, а остальные 11 в том же слое продолжают делать свою работу. Их выходы складываются, и ваша красивая картинка тонет в шуме.

Решение? Либо усреднять loss по всем головам (тогда каждая будет пытаться приблизиться к целевой картинке), либо точечно выключать другие головы через mask в forward pass. Второе - читерство, но работает.

Что это вообще значит для интерпретируемости?

Когда вы видите, что замороженная GPT-2 может "смотреть" на Bad Apple, возникает два противоречивых чувства:

1. Ура, мы можем управлять вниманием модели!
2. Чёрт, значит, карты внимания ничего не значат?

И то, и другое одновременно. Этот эксперимент показывает, что механизм внимания в трансформерах - это гибкий инструмент, который можно настроить на паттерны, совершенно не связанные с исходной задачей модели. Но это же ставит под вопрос интерпретацию attention maps в explainable AI.

Если одна голова внимания в GPT-2 может смотреть на Bad Apple, значит ли это, что когда она смотрит на местоимение в тексте, это тоже просто артефакт оптимизации? Не совсем. В обученной модели эмбеддинги фиксированы (слова), а веса QKV выучены под языковые паттерны. В нашем эксперименте мы крутим эмбеддинги, оставляя веса неизменными.

А что с более новыми моделями?

На 2026 год GPT-2 - уже музейный экспонат. GPT-4, Claude 3, Gemini 2.0 - у них другие архитектуры, другие механизмы внимания (может быть, даже не dot-product). Но принцип остаётся тем же: если есть обучаемые входные представления и фиксированные веса, можно оптимизировать вход под нужный выход.

Современные LLMs часто используют multi-query attention или grouped-query attention. Это усложняет задачу, потому что у вас меньше "ручек" для управления. Но суть эксперимента от этого не меняется.

Попробуйте повторить этот трюк с современной архитектурой - и увидите, насколько устойчивы принципы, заложенные в 2017 году в "Attention is All You Need".

Финальный совет: не верьте картинкам слепо

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

Если вы можете заставить GPT-2 смотреть на Bad Apple, значит, где-то в пространстве эмбеддингов существует конфигурация, которая даёт такую матрицу внимания. Но это не значит, что модель "видит" Bad Apple в том смысле, в каком это делает человек или даже CNN.

Это значит, что механизм внимания - универсальный сопоставитель паттернов. И иногда эти паттерны могут быть удивительно далеки от того, для чего модель обучалась.

Попробуйте повторить эксперимент с другими изображениями. Попробуйте оптимизировать не одну голову, а сразу несколько. Попробуйте использовать другие функции потерь (SSIM вместо MSE). Каждый раз вы будете открывать что-то новое о том, как работает эта, казалось бы, простая формула softmax(QK^T/√d).

И помните: если где-то в интернете вы увидите красивую визуализацию внимания LLM, спросите себя - а что, если автор просто подобрал эмбеддинги?