Отключение 729 голов внимания в GPT-2-XL на 91,5% сократило дословное запоминание

Языковые модели умеют дословно воспроизводить куски обучающих данных, а это риск утечки персональной информации и претензий по авторскому праву. Исследователи показали, что часть этой проблемы можно снять без дообучения: достаточно перед инференсом обнулить «лишние» головы внимания. В GPT-2-XL (1,5 млрд параметров) отключение 729 таких голов сократило число дословно воспроизводимых последовательностей на 91,5%, а перплексия на WikiText-103 при этом улучшилась на 0,92.

Какие головы убирали и как их находили

Идея опирается на два наблюдения. Теоретические работы говорят, что многоголовое внимание работает как ассоциативная память, причём её ёмкость растёт с числом голов. Эмпирические работы давно показали, что значительную долю голов можно выключить почти без потери качества. Связать эти факты никто не пробовал.

Авторы замораживают все веса и обучают только бинарные маски на выходах голов. Маска ступенчатая, поэтому градиенты проходят через сигмоидную аппроксимацию (straight-through estimator). Штраф в функции потерь не даёт перплексии вырасти выше заданного порога и одновременно поощряет меньшее число активных голов. Оптимизатор AdamW, батчи по 8 последовательностей, 1000 шагов. Берётся не последняя конфигурация, а та, где зафиксировано минимальное отношение перплексии к исходной. Маски подбирают на отдельной части WikiText-103 или LAMBADA (датасет на угадывание последнего слова абзаца), не пересекающейся с оценочной.

Запоминание определяли через k-extractability: если подать модели начало последовательности, она при жадном декодировании выдаёт продолжение слово в слово. Тесты шли на GPT-2-XL с корпусом OpenWebText и на Pythia-2.8B и Pythia-6.9B (обе Deduped) с корпусом Pile-Deduped.

Что получилось в цифрах

На GPT-2-XL маски, подобранные по WikiText-103, убрали около 61% голов. Перплексия снизилась с 10,08 до 9,16, а число запомненных последовательностей в одной из конфигураций упало с 5590 до 315. Для обеих Pythia картина та же: перплексия 7,56 → 7,32 и 6,90 → 6,73, запоминание с 20 000 до 2875 и 3500 последовательностей. По всем моделям и двум бенчмаркам снижение составило от 44 до 94%.

Маски по LAMBADA отключают меньше голов, и запоминание падает слабее, зато точность предсказания последнего слова сохраняется или чуть растёт. Наборы удаляемых голов для двух бенчмарков и трёх моделей различаются, но эффект есть везде. Авторы делают вывод, что запоминание распределено по всему механизму внимания, а не сидит в узкой группе голов.

Контрольный эксперимент показывает, что важно не только число отключённых голов, но и какие именно. При случайном отключении того же количества запоминание иногда падало даже сильнее (330 против 473 последовательностей у GPT-2-XL в одной из конфигураций), но модель разрушалась. Перплексия GPT-2-XL на WikiText-103 вырастала с 10,08 до 146,82, у Pythia с 7,56 до 18,30 и с 6,90 до 17,20. Точность на LAMBADA у GPT-2-XL падала с 0,6128 до 0,2014. Сломанная модель просто ничего не воспроизводит, поэтому одного счётчика запоминания для оценки метода мало.

Почему это работает и где слабые места

Механистический разбор сделан на Pythia-2.8B и GPT-2-XL. На заученных последовательностях головы чаще «насыщают» softmax: почти всё внимание уходит на одну позицию. Кроме того, разные головы чаще указывают на одну и ту же позицию. После отключения этот разрыв между заученным и незаученным текстом сокращается. Отрыв логита лучшего токена от второго на заученных последовательностях был 8,79 и 7,68, примерно вчетверо выше, чем на обычном тексте. После отключения он упал до 3,60 и 3,86, а на незаученном тексте почти не изменился. Выходит, затронут именно механизм излишней уверенности, а не общая уверенность модели.

Ограничения авторы называют сами. Механистическая часть корреляционная: данные показывают, что насыщение, согласованность голов и уверенность меняются вместе, но причинность не доказана. Сравнение согласованности делается по пулу насыщенных голов, а он после отключения сжимается, так что часть эффекта может быть артефактом. Почему оставшиеся головы смотрят «шире», тоже не проверено. Анализ не охватывает Pythia-6.9B из-за нехватки вычислений.

Есть и методологическое ограничение. Качество измеряли на том же бенчмарке, по которому искали лишние головы, пусть и на отдельной выборке. Как обрезанная модель поведёт себя на других задачах, в работе не проверено. Эксперименты ограничены тремя небольшими моделями 2020–2023 годов, и на современных LLM с группированным вниманием результат может отличаться. Код обещают опубликовать после принятия статьи.

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