Я со свеженьким, J++ Lens — апдейт к посту про J-Lens

Несколько месяцев назад писала про J-Lens — вышел он весьма громко. После по нему тоже выходили работы, в основном с применением и анализом границ метода.

Но! Теперь из вкусного на потрогать вышла интерактивная демка J++Lens от Kola Ayonrinde (Anthropic Fellows) и Jack Lindsey (Anthropic, один из основных авторов статьи про J-Lens).

Теперь содержательное.

Что было не так с J-Lens?

Напомню: J-Lens — это усреднённый якобиан. Для каждого слоя считаем, как малое изменение активации сдвигает финальный слой, усредняем по куче текстов и переводим результат в словарь.

Проблемы:

• осмысленные readouts (то, что мы можем читать и понимать) появлялись в основном во второй половине слоёв.
• многие промежуточные концепты, которые модель использует, линза не видела.

Это проверяли на задачах, где для ответа нужен промежуточный шаг, которого нет ни в промпте, ни в ответе. Например: «Число ног у животного, которое плетёт паутину, равно…» — чтобы ответить «8», модель должна сначала прийти к «паук». Что такой шаг модель действительно использует, проверяли вмешательствами (например подмена вектора «паук» на «муравей» меняла ответ с 8 на 6), а наличие нужного концепта смотрели в топе по readout по слоям.

• в топе было много пунктуации и мусора.

От чего проблемы (тут не от нервов 🙂 )

Авторы выдвинули гипотезу, что проблема в усреднении (все якобианы входят в среднее с одинаковым весом).
По теории выборочное среднее несмещённое и состоятельное (то есть с ростом выборки оно сходится к истинному среднему). Нюанс — к среднему чего.

В базовой постановке к среднему по всем позициям сразу, включая те, где якобиан плохо описывает модель. Случайный шум с нулевым средним при этом гасится. А систематический — когда целая группа активаций (часто — пунктуация) стабильно даёт «грязненькие» якобианы — остаётся.

Плюс к ранним слоям обычный градиент сильно раздувается. Якобиан от раннего слоя до выхода по цепному правилу — это произведение якобианов всех блоков между ними. Чем раньше слой, тем больше множителей и тем больше накопленной ошибки.

Отсюда идея: что делает J++

1. Jacobian Filtering. Активации слоя делят на 8 кластеров обычным k-means по геометрии (центрируют, нормируют и проецируют на 64 главные компоненты и кластеры подбирают на текстах из WikiText-103). На каждом кластере считают свой усреднённый якобиан — «эксперта». Потом на небольшом размеченном наборе задач учат веса, с которыми экспертов смешать.

И ещё прикольная для меня заметка, что когда миксовали экспертов, попробовали кусочно-линейную карту — свой якобиан на каждый кластер (то есть ИЛИ по экспертам). И такой подход работает хуже, поэтому именно нужно сделать линейную смесь экспертов.

2. LRP вместо обычного backward. LRP — классический метод атрибуции из XAI: вклад раздают от выхода назад по активациям по правилам, которые сохраняют суммарную «релевантность». Здесь от него берут правила обратного прохода. Через attention идёт обычный градиент (хотя, кстати, есть AttnLrp вариант).

3. Из readout убирают токены без букв и цифр.

Тут примечательно, что это было и так, на уровне фильтра для визуализаций J-Lens, но в J++ его сделали частью метода.

Вот таким методом везде получили приросты. Потрогать, потыкать, посмотреть тут:
Пост J++: LessWrong
Демо J-Lens vs J++: Neuronpedia
Код: GitHub
Линзы: HuggingFace

Пара мыслей вдогонку:

Мне очень приятно было увидеть тут классический метод интерпретабилити, если честно. Для меня это просто самое красивое, когда классика снова где-то возникает, а тут ещё и k-means, и ряд Тейлора, и эксперты! Ну восторг.

И опять чертова линейность! Она тоже шикарная (но на ранних слоях эксперты расходятся сильнее, и там связь заметно менее линейна, это тоже важно).

Ещё узнала для себя новую фразу, "эффект белого медведя". Это когда запрещаешь себе думать о чём-то, а мысль возвращается чаще. Формулировка от Достоевского, оказывается.

И всем отличного вторника.