RSAE: новый разреженный автоэнкодер улучшил интерпретацию языковых моделей

В новой научной работе, опубликованной на arXiv, представлен Rational Sparse Autoencoder (RSAE) — усовершенствованная версия разреженного автоэнкодера для механистической интерпретируемости нейросетей.

Существующие разреженные автоэнкодеры (SAE) используют фиксированные нелинейности, такие как ReLU, JumpReLU и TopK, что ограничивает их гибкость. Авторы утверждают, что такой жёсткий подход искажает баланс между качеством реконструкции и разреженностью.

RSAE решает эту проблему заменой фиксированной активации на обучаемую рациональную функцию. Она достаточно гибка, чтобы аппроксимировать активации из существующих семейств SAE, и при этом адаптируется к реальной геометрии пред-активаций.

Метод реализован в два этапа: сначала веса предварительно обученного базового SAE копируются, а коэффициенты рациональной функции инициализируются с помощью ослабленного алгоритма Ремеза на синтетических данных. Затем проводится тонкая настройка под стандартным регуляризованным реконструктивным лоссом.

Эксперименты на активациях residual-stream трёх открытых языковых моделей показали, что RSAE строго улучшает метрики реконструкции и качества downstream-поведения по сравнению с исходными SAE. При этом интерпретируемость на уровне признаков, оцениваемая через sparse probing, не пострадала.

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

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