Проблема «черного ящика» в Attention
Механизм Multi-Head Attention (MHA) в трансформерах остается одной из самых сложных для интерпретации частей архитектуры. Существующие инструменты, такие как Feature vis, позволяют исследовать остаточные потоки (residual stream) и слои MLP, но полностью игнорируют сам механизм внимания. Это создает пробел в понимании того, как именно модели формируют паттерны внимания и какие внутренние функции за этим стоят.
Решение: QK Transcoder Dashboard
Автор предлагает архитектуру, основанную на методе разреженного словарного обучения (Sparse Dictionary Learning, SDL). Вместо плотных векторов запросов (Query) и ключей (Key) используются разреженные функции. Архитектура включает два отдельных трансодера — для Query и для Key, которые обучаются воспроизводить паттерны внимания оригинальной модели GPT-2 small.
Ключевая особенность: активации функций разделяются между всеми головами (heads) в данном слое. Каждая функция имеет вектор декодера для каждой головы, что позволяет одной и той же семантической функции влиять на разные аспекты внимания модели.
Технические детали и метрики
Модель использует функцию потерь Ljoint (расхождение Кульбака-Лейблера между оригинальным и реконструированным вниманием) и вспомогательные потери для сохранения верности оригинальным Q/K векторам. Важнейшим нововведением стало использование «призрачных градиентов» (ghost gradients) для реанимации «мертвых» нейронов.
| Параметр / Метрика | Значение / Описание |
|---|---|
| Модель | GPT-2 small |
| Размер словаря (F) | Зависит от слоя, но активации разрежены |
| Активные функции на токен (L0) | ~10–20 (меньше в ранних слоях, больше в поздних) |
| Доля «мертвых» функций (без ghost gradients) | 80%–95% |
| Доля «мертвых» функций (с ghost gradients) | 4%–12% |
| Размерность декодера | F × 12 × 64 (на каждую голову) |
Возможности дашборда
Интерфейс позволяет исследователям:
- Выбирать пары токенов: Анализировать, как конкретный токен-источник влияет на токен-назначение (например, связь «Aunt» → «Sally» в индуктивных паттернах).
- Визуализировать вклад функций: Смотреть тепловые карты, показывающие, какие именно пары Query-Key функций вносят наибольший вклад в итоговый счет внимания до и после Softmax.
- Декомпозиция: Получать точное разложение счета внимания на вклад каждой функции и смещений (biases).
Значение для индустрии
Этот инструмент представляет собой важный шаг к механической интерпретируемости (mechanistic interpretability). Позволяя «заглянуть» внутрь механизма внимания и увидеть, какие моносемантические функции активируются, исследователи могут лучше понимать, как модели хранят и обрабатывают информацию, а также выявлять потенциальные уязвимости или нежелательные паттерны поведения на уровне отдельных нейронов.
Источник: LessWrong ↗