Патчинг весов LLM вручную

Я написал реализацию ROME — редактирования модели ранга 1 — с нуля, используя только torch и transformers. Цель: переписать один факт внутри GPT-2 Medium с помощью одного сложения матриц и посмотреть, что при этом происходит со всем остальным, что знает модель.

Я выполнил четыре редактирования. Три были хирургически точными. Четвёртое показало, что «хирургичность» ROME чрезвычайно чувствительна к деталям реализации, о которых не обязательно думаешь заранее.

Четыре редактирования

Я выбрал факты, которые GPT-2 Medium действительно знает с уверенностью (P > 0.5 для правильного ответа), из четырёх разных областей:

Все четыре редактирования попали в цель. После обновления модель предсказывала новый ответ с вероятностью ≥ 0.98 для точного редакционного промпта. Так что ROME сам по себе работает. Меняется всё остальное.

Три чистых редактирования

Три из четырёх вели себя примерно так, как предсказывает статья ROME. Вот что произошло с несвязанными фактами после каждого редактирования:

Тако → Япония (норма обновления: 9% от нормы весов)

Только буррито — ближайший сосед в пространстве мексиканской кухни — был затронут.

Гарвард → Калифорния (норма обновления: 13%)

Йель — ближайший сосед Гарварда в пространстве Лиги плюща — последовал за ним. Всё остальное не сдвинулось.

Статуя Свободы → Лас-Вегас (норма обновления: 16%)

Два ориентира поколебались. Колокол Свободы особенно забавен — он находится не в Нью-Йорке, но содержит слово «Liberty» («Свобода»), и GPT-2 их перепутал.

Итог: чистые прямые редактирования, узкий побочный ущерб, разумные величины обновлений (9–16%).

Модель также охотно сочиняла альтернативные истории под стать изменениям. Мой любимый пример: после перемещения Гарварда в Калифорнию, на вопрос о том, когда был основан Гарвард, модель ответила: «В 1776 году французским иезуитским отцом Шарлем де Монтескьё.» Полноценное предложение, внутренне согласованное, полностью ложное. Это модельные априорные знания («Гарвард престижен, у знаменитых мест знаменитые основатели»), заполняющие дыру на месте исчезнувшего реального факта.

Проблемный случай: Google

Редактирование Google прошло плохо. Моя первая запись показала норму обновления 117% — изменение ранга 1 было больше нормы самой матрицы весов — и это обрушило весь калифорнийский технологический кластер:

• Apple → Техас (P = 1.00)
• Microsoft → Техас (P = 1.00)
• Кремниевая долина → Техас (P = 1.00)
• Стэнфордский университет → Техас (P = 0.75)
• И ещё: «Google был основан в 2000 году Стивом Джобсом».

Я написал пост в блог, сделав это главным выводом. Потом подумал ещё раз и понял, что 117% подозрительно. Редактирование ранга 1 не должно быть больше того, что оно редактирует.

Отладка Google

Обнаружились две проблемы.

Проблема 1: недостаточно приготовленная ковариация

Формула обновления ROME опирается на C — ковариацию промежуточных (post-GELU) векторов на целевом слое:

u = torch.linalg.solve(C + lambda * I, h_star)
delta_W = (u / (h_star @ u)).unsqueeze(1) @ (v_star - W @ h_star).unsqueeze(0)

Направление C⁻¹ @ h* и делает редактирование избирательным — оно выровнено с h*, но ортогонально типичным ключам. Если C плохо обусловлена, C⁻¹ @ h* взрывается в направлениях малых собственных значений, и обновление становится огромным.

Я оценивал C по 200 выборкам WikiText — примерно 10 600 токенов. Для ковариационной матрицы 4096×4096 это около 2.5 выборок на измерение. Матрица была крайне дефицитна по рангу. У меня была регуляризация 1e-4 * I, чего было совершенно недостаточно.

Исправление: 2000 выборок (~118 000 токенов, ~29 на измерение) и регуляризация, масштабированная по следу (1e-2 × mean(diag(C))).

Результат для того же редактирования: норма обновления упала со 117% до 50.5%. Лучшая обусловленность, вдвое меньшая величина обновления.

Но 50% всё ещё огромно, и контроли всё ещё ломались:

Так что часть «катастрофы» была ошибкой — но не вся. Утечка кластера была реальной.

Проблема 2: позиция имеет огромное значение

«Google» — это один BPE-токен. В моём промпте — «Google is a company headquartered in the state of» — он стоит на позиции 0. Это означает, что h* (промежуточный вектор, используемый для редактирования) вычисляется из токена, который не видел никакого предшествующего контекста. Это голое представление.

Что если разместить Google не на позиции 0? Я изменил промпт на «The technology company known as Google is headquartered in the state of» — теперь Google стоит на позиции 5, имея «The technology company known as» в качестве предшествующего контекста.

Та же цель редактирования. Та же новая ковариация. Результат:

• Норма обновления: 9.1% (снизилась с 50.5%, снизилась со 117%)
• Apple → Техас (0.97) — всё ещё затронут
• Microsoft → Техас (0.57) — частично затронут
• Кремниевая долина → Калифорния (0.81) ✓ сохранена
• Стэнфорд → Калифорния (0.92) ✓ сохранён

Просто добавив субъектному токену немного контекста перед вычислением h*, редактирование становится достаточно хирургическим, чтобы Кремниевая долина и Стэнфорд выжили. Apple и Microsoft всё ещё получили толчок, так что некоторая реальная утечка по соседству с Google существует — но ничего похожего на исходный апокалипсис.

Чему это на самом деле учит

Я хотел, чтобы этот пост стал «смотрите, ROME не может редактировать хаб-концепции — вот как Google всё разнёс». Этот фрейминг был ошибочным. Правда интереснее и менее драматична:

1. ROME чувствителен к деталям реализации, которые не видны в статье. Количество выборок для ковариации. Сила регуляризации. Где стоит субъектный токен в промпте. Ошибитесь в любом из этих пунктов, и ваш «катастрофический побочный ущерб» может оказаться вашим собственным кодом.

2. Одиночные токены-субъекты на позиции 0 — наихудший случай. Их h* наименее дискриминативен, и любая численная погрешность в C при обращении превращается в чрезмерно большое обновление. Если хотите чистого редактирования — дополните субъект предшествующим контекстом.

3. Утечка хаб-концепций реальна, но невелика. Даже при правильной ковариации и предшествующем контексте редактирование Google слегка сдвигает Apple и Microsoft. «Google» находится в плотном семантическом окружении, и редактирование ранга 1 касается этого окружения. Можно сократить это ещё в 2–4 раза с помощью многослойного распределения в стиле MEMIT, но полностью устранить нельзя.

4. Норма обновления — надёжный диагностический показатель. Ниже 15% нормы весов: скорее всего, всё в порядке. Выше 50%: скорее всего, что-то сломано — либо из-за ошибки, либо потому что вы редактируете хаб. Проверяйте перед тем, как доверять результату редактирования.

Конфабуляции реальны

При каждом успешном редактировании модель придумывала согласованные альтернативные факты в соответствие с изменениями:

• Гарвард (теперь в Калифорнии) был основан французским иезуитом в 1776 году.
• Тако (теперь из Японии) подаются с рисом, а язык, наиболее ассоциируемый с тако, — испанский.
• Паромная переправа к Статуе Свободы теперь отправляется из «Международного аэропорта Лас-Вегаса».
• Google (теперь в Техасе) был «основан в 2000 году Стивом Джобсом».

Это не шум — это модель, применяющая свои априорные знания к изменённому факту. Как только она убеждается, что Google — техасская компания, «основана Стивом Джобсом» — это не случайная галлюцинация; это лучшее предположение модели о том, какой должна быть история основания знаменитой техасской технологической компании.

Знания внутри языковой модели — это не список независимых фактов. Это граф фактов, взаимно подкрепляющих друг друга. Отредактируйте один узел — и граф породит вокруг него согласованный (и полностью ложный) новый регион.

Реализация

Вся реализация занимает около 500 строк в нескольких файлах: причинная трассировка, оценка ковариации, градиентный спуск для v*, обновление весов ранга 1 и сквозной скрипт для четырёх редактирований.

Зависимости: torch, transformers, datasets. Ничего специфичного для ROME.

Работает на CPU. Каждое редактирование занимает около 3 минут при ковариации из 200 выборок и около 15 минут при 2000 выборках. Настройка с правильной ковариацией стоит ожидания.