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

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

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

Четыре правки

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

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

Три чистые правки

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

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

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

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

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

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

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

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

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

Беспорядочная: Google

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

• Apple → Техас (P = 1,00)
• Microsoft → Техас (P = 1,00)
• Silicon Valley → Техас (P = 1,00)
• Stanford University → Техас (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) — частично увлечена
• Silicon Valley → Калифорния (0,81) ✓ сохранена
• Stanford → Калифорния (0,92) ✓ сохранена

Просто дав токену субъекта немного контекста перед вычислением h*, правка становится достаточно хирургической, чтобы Silicon Valley и Stanford выжили. 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 образцам. Правильная настройка ковариации стоит ожидания.