Предложен слой MoWE: нейросеть с читаемыми координатами без SAE

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

Ответ, слой Mixture of Wide Experts (MoWE). Это разреженный feed-forward слой (полносвязный блок трансформера), в котором выбранные эксперты читают всё представление целиком, а пишут только в фиксированные, не пересекающиеся между собой выходные координаты. На контролируемой задаче с известными скрытыми причинами такая конструкция, по словам авторов, сохраняет качество предсказаний и при этом локализует причины яснее, чем плотная модель и стандартный Mixture-of-Experts (MoE).

Эксперименты: при примерно совпадающих вычислениях авторы обучили по два 19-слойных прогона каждой архитектуры на 10 млрд токенов. Оценочная кросс-энтропия (чем ниже, тем лучше) у MoWE составила 2,417, у плотной модели, 2,452, а у стандартного MoE, 2,22. То есть MoWE чуть лучше плотной модели, но заметно уступает обычному MoE по этому показателю.

Читаемость «родных» координат измеряли так: языковая модель-судья описывает каждую координату по отрывкам, где та сильно активна, а затем это описание проверяют на новых отрывках. С поправкой на описания, подходящие и к другим координатам, получается показатель специфичности: 0,523 у MoWE, 0,375 у плотной модели и 0,268 у стандартного MoE. Для сравнения, SAE типа TopK, обученные на тех же обновлениях FFN, дают 0,489 для плотной модели и 0,556 для MoWE. Поэтому «родной» MoWE по читаемости сопоставим с SAE поверх плотной модели.

Две демонстрации. В первой авторы нашли причинный путь из четырёх координат, поддерживающий предсказание слова degrees (градусов) после фразы «A right angle is 90». Затем они скопировали семь связанных с углами значений координат в модель, обрабатывающую «A rectangle has four», и следующее предсказанное слово сменилось с sides (сторон) на degrees. Во второй один фиксированный набор координат обнаруживает опасные запросы, а другой подталкивает следующий токен к отказу. Из этого авторы делают вывод: MoWE сочетает разреженные условные вычисления с координатами, которые можно находить и использовать напрямую.

Ключевые факты

MoWE, разреженный слой: выбранные эксперты читают всё представление, но пишут в фиксированные непересекающиеся выходные координаты.
Кросс-энтропия при примерно равных вычислениях (19 слоёв, 10 млрд токенов): MoWE 2,417, плотная модель 2,452, стандартный MoE 2,22.
Специфичность координат по оценке модели-судьи: MoWE 0,523, плотная 0,375, стандартный MoE 0,268; TopK SAE на тех же обновлениях FFN, 0,489 (плотная) и 0,556 (MoWE).
Копирование семи значений «угловых» координат изменило предсказание модели с sides на degrees после фразы «A rectangle has four».
Во втором примере один набор координат распознаёт опасные запросы, другой направляет следующий токен к отказу.

Почему это важно

Обычно интерпретируемость добывают постфактум: обучают SAE поверх готовой модели. MoWE пытается встроить читаемость в саму архитектуру, так что координаты выхода слоя можно изучать и использовать напрямую. По заявлению авторов, при этом сохраняются разреженные условные вычисления, как в MoE. Если подход масштабируется, он мог бы упростить поиск и редактирование отдельных механизмов модели, но это пока только заявка препринта.

Кому это важно

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

Как это применить

Практического применения в источнике не описано: нет данных о релизе кода или весов. Реалистичный шаг для исследователя, воспроизвести методику оценки (описание координаты моделью-судьёй и проверка на новых отрывках) и сравнить MoWE с плотной моделью и обычным MoE на собственных данных.

Можно ли доверять

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

Риски и подводные камни

По кросс-энтропии MoWE (2,417) лишь немного опережает плотную модель (2,452) и заметно уступает стандартному MoE (2,22), так что выигрыш в читаемости достаётся не бесплатно относительно качества MoE. Формулировка «примерно совпадающие вычисления» в аннотации не раскрыта количественно. Специфичность MoWE (0,523) сопоставима с SAE на плотной модели (0,489), но ниже, чем у SAE на самом MoWE (0,556). Масштабируемость за пределы 19 слоёв и 10 млрд токенов в источнике не показана.

«Языковые модели часто размазывают одно понятие по множеству внутренних измерений, из-за чего их признаки трудно интерпретировать.»

— Из аннотации препринта «Mixture of Wide Experts: Sparse Routing for Interpretable Native Coordinates», arXiv