Мы обнаружили в GPT-2 нейрон конкретного токена

Мы начали с вопроса: откуда GPT-2 знает, когда использовать слово an, а не a? Выбор зависит от того, начинается ли следующее за ним слово с гласной, однако GPT-2 может прогнозировать только одно слово за раз.

У нас по-прежнему нет полного ответа, однако мы нашли нейрон MLP в GPT-2 Large, который необходим для прогнозирования токена " an". Также мы выяснили, что веса этого нейрона соотносятся с эмбеддингом токена " an", что позволило нам найти другие нейроны, прогнозирующие конкретный токен.

Обнаружение нейрона

Выбор промпта

Удивительно сложно оказалось придумать промпт, при котором GPT-2 должен был вывести токен " an" (начальный пробел является частью токена) в качестве наиболее вероятного прогноза. На самом деле, мы отказались от GPT-2_small и перешли на GPT-2_large. Как будет видно ниже, даже GPT-2_large систематически занижает предсказание " an" в пользу " a". Возможно, это вызвано тем, что малые языковые модели склонны к повышенной частоте «a», чтобы дать наилучшую догадку. В конечном итоге, промпт, который давал высокую вероятность (64%) " an", оказался таким:

«I climbed up the pear tree and picked a pear. I climbed up the apple tree and picked»

Первое предложение необходимо, чтобы подтолкнуть модель к неопределённому артиклю, без него модель делала бы прогнозы, например, "[picked] up".

Прежде чем двигаться дальше, рекомендую прочитать краткое описание архитектуры трансформера. Каждый блок внимания и MLP получает вводы и добавляет выводы в остаточный поток.

Logit Lens

При помощи методики logit lens мы взяли логиты из остаточного потока между каждым из слоёв и составили график разности между logit(‘ an’) и logit(‘ a’). После MLP слоя 31 мы обнаружили большой всплеск.

Патчинг активации по слоям

Патчинг активации — это методика, предложенная Менгом и соавторами (2022 год) для анализа важности отдельного слоя в трансформере. Сначала мы сохранили активацию каждого слоя при пропускании исходного промпта через модель — «чистую активацию».

Затем мы пропустили через модель повреждённый промпт: !I climbed up the pear tree and picked a pear. I climbed up the lemon tree and picked". Заменив слово «apple» на «lemon», мы заставили модель спрогнозировать токен " a" вместо " an".

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



В остальной части поста мы по большей мере будем игнорировать внимание, однако стоит заметить, что эти результаты показывают, что именно в слое 26 " picked" начинает сильно думать об " apple", что очевидно необходимо для прогнозирования " an".


Здесь выделяются два слоя MLP: слой 0 и слой 31. Мы уже знаем, что MLP слоя 0 в общем случае важен для работы GPT-2 (хоть и не знаем, почему важно внимание в слое 0). [Замечание Нила Нанда о MLP 0:

«В GPT-2 Small часто наблюдается большая важность MLP0, и его удаление полностью уничтожает результативность модели. Сейчас моя основная гипотеза заключается в том, что первый слой MLP, по сути, действует как расширение эмбеддинга (по какой-то причине), и что когда позже слой хочет получить доступ а токенам ввода, он в основном считывает вывод первого слоя MLP, а не эмбеддинги токенов. В таких условиях первый слой внимания не так уж важен.

Если это так, то логично, что MLP0 важен для второго токена, потому что это позиция с другим входным токеном!

Я не совсем понимаю, почему это происходит, но предположу, что из-за того, что матрицы эмбеддингов и анэмбеддингов в GPT-2 Small одинаковы. Это довольно непринципиально, поскольку задачи токенов эмбеддингов и анэмбеддингов не взаимно обратны, но это распространённая практика, и, похоже, что модели для преодоления этого хотят выделить некоторые параметры.

У меня есть только косвенные свидетельства этого, и мне бы хотелось, чтобы кто-то изучил тему должным образом!»] Влияние слоя 31 более интересно. Наши результаты дают понять, что MLP слоя 31 играет важную роль в прогнозировании токена " an". (Если вам непонятно, как этот результат соответствует представленной выше logit lens, см. этот комментарий.)

Открытие 1: мы можем выявлять прогнозирующие нейроны патчингом активаций отдельных нейронов

Патчинг активаций использовался для послойного исследования трансформеров, но можем ли мы расширить возможности технологии и применять его к отдельным нейронам? Так как каждый MLP в трансформере имеет только один скрытый слой, активация каждого нейрона не влияет ни на какой другой нейрон в MLP. Поэтому мы сможем патчить отдельные нейроны, поскольку они независимы друг от друга в том же смысле, что и головы внимания в одном слое независимы друг от друга.

Мы выполнили патчинг активаций для каждого нейрона в MLP слоя 31 аналогично тому, как выполняли патчинг слоёв. Мы повторно внедряем в MLP чистую активацию каждого нейрона и пропускаем через модель испорченный промпт, изучая, как восстановление каждого нейрона влияет на разность логитов между " a" и " an".


Мы видим, что патчинг нейрона 892 восстанавливает 50% от разности логитов чистого промпта, а патчинг всего слоя на самом деле восстанавливает всего на 49%.

Открытие 2: активация «an-нейрона» коррелирует с прогнозированием токена " an".

Примеры максимальной активации нейрона 892 слоя 31 в Neuroscope

Neuroscope — это онлайн-инструмент, отображающий примеры максимальной активации в большом датасете для каждого нейрона GPT-2. Взглянув на нейрон 892 слоя 31, мы увидим, что нейрон максимально активируется на токенах, для которых последующим токеном будет " an".

Однако Neuroscope показывает лишь 20 примеров максимальной активации. Будет ли это трендом для более широкого диапазона активаций?

Тестирование нейрона на более крупном датасете

Чтобы проверить наличие тренда, мы пропустили через эту модель датасет pile-10k. Это разнообразный набор примерно из десяти миллионов токенов, взятых из The Pile, разделённый на промпты из 1024 токенов. Мы составили график пропорции прогнозов " an" для диапазона активаций нейронов:


Видно, что пропорция прогнозов " an" увеличивается с повышением активации нейрона до того момента, когда " an" становится постоянным верхним прогнозом. Тренд довольно шумный, и это намекает, что в модели могут быть другие механизмы, влияющие на прогнозирование " an". А может быть, во время увеличения логита " an" увеличиваются и другие логиты.

Обратите внимание, что модель спрогнозировала " an" только примерно 1500 раз, хотя на самом деле этот токен в датасете встречается 12 тысяч раз. Неудивительно, что найти хороший промпт было так сложно!

Выходные веса нейрона имеют большое скалярное произведение с токеном " an"

Как нейрон влияет на вывод модели? Выходные веса нейрона имеют большое скалярное произведение с эмбеддингом для токена " an". Мы называем это конгруэнтностью нейрона с токеном. По сравнению с другими случайными токенами наподобие " any" и " had", конгруэнтность нейрона с " an" очень высока:


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


Похоже, нейрон, по сути, добавляет эмбеддинг " an" в остаточный поток, что увеличивает выходную вероятность для " an", поскольку этап анэмбеддинга состоит из получения скалярного произведения окончательного остатка с каждым токеном. [Что ещё он мог сделать? Он мог подавить логит для " a", что могло оказать то же влияние на разность логитов. Или он мог добавить совершенно другое направление для остаточного потока, что заставило бы нейрон в последующем слое увеличить логит " an".]

Существуют ли другие нейроны, также конгруэнтные с " an"? Чтобы выяснить это, можно вычислить конгруэнтность всех нейронов с токенов " an":


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

Если это так, то можем ли мы использовать конгруэнтность для поиска нейрона, идеально коррелирующего с прогнозированием одного токена?

Открытие 3: можно использовать выходную конгруэнтность нейронов для поиска конкретных нейронов, прогнозирующих токен

Поиск нейрона, связанного с токеном

Можно попробовать найти связанный с конкретным токеном нейрон, выполнив следующий поиск:

1. Для каждого токена находить нейрон с наибольшей выходной конгруэнтностью.
2. Для каждого из таких конгруэнтных нейронов находить, насколько более они конгруэнтны по сравнению с следующим наиболее конгруэнтным нейроном для того же токена.
3. Взять нейрон или нейроны, которые наиболее сильно выделяются конгруэнтностью.

При таком поиске мы хотим найти нейроны, уникальным образом отвечающие за токены. Наша гипотеза заключается в том, что у нейрона, наиболее отвечающего за токен, активация будет сильнее коррелировать с прогнозированием токена, так как любой прогноз этого токена будет «полагаться» на этот нейрон.

Давайте выполним поиск и составим график наиболее конгруэнтных нейронов для каждого токена:


При таком поиске мы видим, что для токенов наподобие " off" и «though» есть нейроны, конгруэнтность которых выделяется. Давайте пропустим нейрон " though" (нейрон 1921 слоя 28) через датасет и посмотрим, стал ли график более чистым!


Ого, он гораздо более грязный, чем граф для an-нейрона. Что происходит?

Изучение данных нейрона в Neuroscope показывает нам, что нейрон с максимальной активацией прогнозирует токены и " though", и " however". Это усложняет анализ: похоже, этот нейрон коррелирует с группой семантически схожих токенов (соединительных наречий). [Стоит заметить, что в то время, как though-нейрон конгруэнтен группе семантически схожих токенов, an-нейрон коррелирует с группой синтаксически схожих токенов (например, " an" и " Ancients").]


При вычислении конгруэнтности нейрона для всех токенов мы обнаружим, что выбросами становятся одни и те же токены:


В показанном выше графике корреляции с большим датасетом случаи, когда нейрон активируется и прогнозируется " however" вместо " though" будут считаться отрицательными примерами, поскольку " though" не был верхним прогнозом. Это также может объяснить часть шума в корреляции " an", где нейрон также конгруэнтен с "An", " An" и "an".

[Почему " an" имеет более чёткую корреляцию с другими конгруэнтными токенами? Мы не знаем этого точно. Одно из возможных объяснений: "An" и " An" просто являются менее частыми токенами, поэтому мало влияют на корреляцию, а "an" имеет гораздо меньшую конгруэнтность с нейроном, чем верхние три.

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

Можем ли мы найти более подходящий нейрон, предпочтительно, такой, который прогнозирует только один токен?

Поиск чётко ассоциированного нейрона

Чтобы нейрон был «чётко ассоциирован» с токеном, их конгруэнтность друг с другом должна быть взаимно исключительной, то есть:

1. Нейрон гораздо более конгруэнтен с токеном, чем любой другой нейрон.
2. Нейрон гораздо более конгруэнтен с токеном, чем с любым другим токеном.

(Помните, что конгруэнтность — это просто скалярное произведение.)

Оба критерия помогают упростить взаимосвязь между нейроном и его токеном. Если конгруэнтность нейрона с токеном — это показатель того, насколько он вносит вклад в прогнозирование этого нейрона, то первый критерий можно рассматривать как гарантию того, что только этот нейрон отвечает за прогнозирование этого токена, а второй критерий можно рассматривать как гарантию того, что этот нейрон отвечает за прогнозирование только этого токена.

Тогда наш поиск будет выглядеть так:

1. Для каждого токена находим наиболее конгруэнтный нейрон.
2. Для каждого нейрона находим наиболее конгруэнтный токен.
   
   [При изучении наиболее конгруэнтных токенов для каждого нейрона мы видим знакомых возмутителей спокойствия, встречающихся с очень высокой конгруэнтностью:
   
   
   
   
   Поначалу кажется, что все эти «запретные токены» связаны с «запретным нейроном» (нейроном 3354 слоя 35), с которым все они очень конгруэнтны. Но на самом деле, если мы создадим график наиболее конгруэнтных токенов многих других нейронов, мы также увидим, что на нём встречается часть их этих странных токенов. Наша экспериментальная гипотеза заключается в том, что это как-то связано с hubness-эффектом.]
3. Находим пары «токен-нейрон», находящиеся в обоих списках, то есть пары, в которых наиболее конгруэнтный нейрону токен — это токен, наиболее конгруэнтный с этим нейроном!
4. Вычисляем, насколько они уникальны, умножая их разность конгруэнтности двух верхних нейронов с их разностью конгруэнтности верхних двух нейронов.
5. Находим пары с максимальной взаимно исключительной конгруэнтностью.

В случае GPT-2_large высочайшим по этой метрике является пара нейрона 4142 слоя 33 с токеном "i". Изучение в Neuroscope [для этого нейрона данные Neuroscope были недоступны, поэтому мы взяли примеры максимальной активации датасета из датасета pile-10k. Тексты 1, 2, 3 — это примеры 1755, 8528 и 6375.] подтверждает эту связь:


А затем мы составим график пропорции верхних прогнозов по активации для пар с пятью наибольшими значениями. [Стоит заметить, что один из пяти верхних токенов — это 'an', но он довольно сильно отличается от ' an', о котором мы говорили выше, поскольку он редко будет использоваться как начало слова или отдельное слово. И аналогично, нейрон, парой которому он является (нейрон 4549 слоя 34), не является an-нейроном, о котором мы говорили выше.]:


Мы видим, что для каждой пары и в самом деле есть плавные корреляции!

Что всё это значит?

Действительно ли конгруэнтность нейрона с токеном измеряет степень, в которой нейрон прогнозирует токен? Мы не знаем. Может быть множество причин, по которым пары «токен-нейрон» с высокой взаимно исключительной конгруэнтностью не всегда коррелируют:

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

Однако мы выяснили, что пара «токен-нейрон» с наибольшей взаимно исключительной конгруэнтностью («i» и «i-нейрон») на самом деле имеет сильную корреляцию. Мы не тестировали другие пары, но ожидаем, что многие другие пары, имеющие высокое значение этой метрики, также коррелируют.

Краткая сводка

1. Мы использовали патчинг активации на уровне нейронов, чтобы найти нейрон, важный для прогнозирования токена " an" в конкретном промпте.
2. Активация «an-нейрона» действительно коррелирует с прогнозированием " an" в общем случае.
3. Это может быть вызвано тем, что выходные веса нейрона имеют большое скалярное произведение с токеном " an" (нейрон имеет высокую конгруэнтность с токеном). Более того, этот нейрон имеет скалярное произведение с этим токеном больше, чем с любым другим токеном. И этот нейрон имеет большее скалярное произведение с этим токеном, чем с любым другим нейроном (они имеют высокую взаимно исключительную конгруэнтность).
4. Конгруэнтность между нейроном и токеном — это отлично. Мы нашли пару нейрон-токен «i», имеющую наибольшую взаимно исключительную конгруэнтность по сравнению с любой другой парой «нейрон-токен». Активация этого нейрона сильно коррелирует с прогнозированием токена "i".

Воспроизводящий наши результаты код можно найти здесь.