SAE: введение, пояснение и код

Привет, друзья! В прошлой статье мы разобрали идею применения автоэнкодеров к трансоформерам. Там весь наш pipeline проходил на идее сжатия признакового пространства так, чтобы поделить кошек и собак. Но что делать, если у нас не задача классификации, а задача next token prediction? Да и признаки не соответствуют "собакам" и "кошкам", а охватывают все богатство естественного языка...

Ответ сообщества сейчас такой — давайте использовать SAE.

SAE: Sparse Autoencoders

Разреженным автокодировщиком назовем модель , (все эти обозначения нам понадобятся), такую, что: где — функция, возвращающая размерность.

Говоря словами — это автокодировщик, у которого скрытая рамерность "раздута" и сильно больше размерности эмбеддинга.

Идею SAE популяризовали Antropic, в октябре 2023 показав эффективность SAE на игрушечной модели, а потом, в мае 2024, расширив результаты для Claude 3.

Что делает AE SparseAE?

Обучение SAE происходит по стандартной минимизации функции потерь, которая делает результат выхода декодера близким к начальному . Базовый пример такой функции — MSE. Сама по себе раздутая размерность не гарантирует эффективную разреженость, потому что задача имеет бесконечно много решений и совершенно не факт, что лучшие по MSE — разрежены. Обычно это реализуется через дополнительный член в функции потерь:

или через аналогичную регуляризацию, поощряющую нули, почти нули или редкие всплески активаций.

Почему? (Можете пропустить этот блок)

Давайте сделаем вид, что энкодер у нас уже фиксирован и что штрафа (члена с ) нет. Тогда

Оптимизируем по (z):

Берём производную:

На константу нам всё равно. Она минимума не изменит. Поэтому условие оптимума (стационарная точка):

Почему решений много (недоопределённость при (D>d))

Матрица имеет размер , но её ранг ,
а значит вырождена (ее определитель — 0, обратной матрицы нет, всё грустно), стало быть решений либо нет, либо их бесконечно много.

И размерность множества решений может быть как минимум , что при D=16 000 (стандартный SAE размер), а d = 2304 — размер эмбеддингов — есть 13 696. Не самое разреженное в мире.

Именно поэтому без L1/top-k и т.п. MSE не “выбирает” разреженный я: оптимум не единственный, и плотненькие решения столь же хороши по MSE. Поэтому мы штрафуем. И на практике — активны 100-200 нейронов на input +/-.

Зачем такой “толстый” скрытый слой?

Расширение пространства — это способ бороться с полисемантичностью — явлением, когда 1 нейрон кодирует много признаков. Эмпирически показано, что если латентное пространство в 10–100 раз больше embedding space, то модель может позволить себе паттерн:

• один нейрон — один или одно осмысленное подмножество признаков (нейрон golden bridge, нейрон немецкого языка и пр). Он включается редко, но всегда по одному и тому же смыслу.

В какой точке обучают SAE?

Чтобы ответить на этот вопрос, давайте рассмотрим процесс трансформации эмбеддинга при проходе через конкретный блок трансформера.

​Если посмотреть на SAE, обученные на конкретных моделях, то можно встретить следующие точки:

### Скрипт загрузки таблицы со списком SAE из sae-lens ### Полный ноутбук-туториал — в конце статьи from sae_lens.loading.pretrained_saes_directory import get_pretrained_saes_directory df = pd.DataFrame.from_records( {k: v.__dict__ for k, v in get_pretrained_saes_directory().items()} ).T.reset_index(drop=True) print("SAEs in the GTP2 Small Resid Pre release") for k, v in df.loc[df.release == "gpt2-small-res-jb", "saes_map"].values[0].items(): print(f"SAE id: {k} for hook point: {v}") print("-" * 50) print("SAEs in the feature splitting release") for k, v in ( df.loc[df.release == "gpt2-small-res-jb-feature-splitting", "saes_map"] .values[0] .items() ): print(f"SAE id: {k} for hook point: {v}") print("-" * 50) print("SAEs in the Gemma base model release") for k, v in df.loc[df.release == "gemma-2b-res-jb", "saes_map"].values[0].items(): print(f"SAE id: {k} for hook point: {v}") ### OUTPUT: #.... #SAE id: blocks.11.hook_resid_pre for hook point: blocks.11.hook_resid_pre #SAE id: blocks.11.hook_resid_post for hook point: blocks.11.hook_resid_post #....

Нас интересуют hook_resid_pre и hook_resid_post. Соотнося с картинкой выше — это вход в attention блок и выход после MLP. Иногда также можно встретить SAE, обученные на hook_resid_mid — представление после attention до MLP.

Выбор зависит от цели анализа.

• hook_resid_pre — для анализа глобальных признаков, выделенных до слоя;

• hook_resid_mid — для выделения вклада attention;

• hook_resid_post — для анализа эффекта блока целиком.

На сегодняшний день единого «рецепта» или стандарта, какой именно слой/точка лучший всего для обучения SAE — нет. И наиболее частое место — hook_resid_pre.

Работа с обученным SAE

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

И, как я говорила выше, для любого будет активно небольшое подмножество латентных нейронов. К примеру, при типичном размере SAE D=16 000 ненулевыми оказываются порядка 100–200 компонент. Это и позволяет рассматривать латентное пространство не как плотную смесь признаков, а как набор активируемых направлений.

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

Размечать признаки руками — долго, поэтому для себя можно и нужно работать с библиотеками. Давайте рассмотрим анализ нейронов в том же SAE Lens.

from IPython.display import IFrame # Посмотрим на признак со случайным индексом feature_idx = 3777 html_template = "https://neuronpedia.org/{}/{}/{}?embed=true&embedexplanation=true&embedplots=true&embedtest=true&height=300" def get_dashboard_html(sae_release="gpt2-small", sae_id="7-res-jb", feature_idx=0): return html_template.format(sae_release, sae_id, feature_idx) html = get_dashboard_html( sae_release="gpt2-small", sae_id="7-res-jb", feature_idx=feature_idx ) IFrame(html, width=1200, height=600)

• Важное тут — SAE не дают «истинного» или «единственно правильного» разложения модели. Они зависят от слоя, точки в блоке, гиперпараметров и даже конкретной инициализации. Тем не менее, SAE-шки на сегодня — один из самых популярных инструментов. С ними можно делать многое, в том числе искать circuts — цепочки преобразований внутри модели, управлять поведением — например заставлять говорить только о собаках (или умнее — заменять/удалять факты из модели) , и много другое.

Базовая теория на этом заканчивается. И чтобы было лучше — предлагаю закрепить ей на практике с этим ноутбуком.

В нём описано:

• извлечение активаций через hook_resid_pre;

• процесс извлечения признаков на аннотированном и нет SAE;

• анализ разреженных латентных признаков + пара кивков для самопроверки;

Пусть он будет вам полезен!
Спасибо за чтение! Другие туториалы публикую в блоге (https://t.me/jdata_blog) и на [гитхаб](https://github.com/SadSabrina/XAI-open_materials/tree/main)!

Удачи и успешных решений,
Ваш Дата-автор!