Обучение репрезентаций в обработке естественного языка тема диссертации и автореферата по ВАК РФ 00.00.00, кандидат наук Балаганский Никита Николаевич

• Балаганский Никита Николаевич
• кандидат наук
• 2026

Оглавление диссертации кандидат наук Балаганский Никита Николаевич

Введение

Глава 1. Общие сведения о архитектрах современных

языковых моделей

1.1 Устройство трансформера

1.1.1 Токенизация

1.1.2 Эмбединги

1.1.3 Слой трансформера

1.1.4 Механизм внимания

1.1.5 Полносвязная сеть

1.1.6 Последний слой

1.2 Виды языковых моделей

1.2.1 Маскированная языеовая модели (MLM)

1.2.2 Последовательная языковыя модель (Causal LM)

1.3 Выводы

Глава 2. Ранний выход в трансформерах

2.1 Связанные исследования

2.2 Мотивация

2.3 Базовые подходы

2.4 PALBERT: Ранний выход для ALBERT

2.4.1 Критерий выхода Q-exit

2.4.2 Архитектура Lambda-слоёв

2.5 Эксперименты и результаты

2.5.1 Абляционное исследование

2.5.2 Зависимость от порога выхода

2.5.3 Скорость против качества

2.5.4 Распределение индексов выхода

2.6 Выводы

Глава 3. ReBased: улучшенные линейные трансформеры

Стр.

3.1 Линейные трансформеры

3.2 Порядок операций в обычном и линейном трансформере ... 34 3.2.1 Линейное внимание как RNN: накопительные

состояния Si

3.3 Связанные исследования

3.4 Модель Based

3.5 ReBased: улучшение ядра

3.6 Эксперименты

3.6.1 Почему измеряют Associative Recall

3.6.2 Синтетические задачи

3.6.3 Моделирование языка

3.6.4 Задачи с несколькими примерами

3.7 Анализ и интерпретация

3.8 Робастность к выбору гиперпараметров

3.9 Выводы

Глава 4. Mechanistic Permutability: сопоставление признаков

между слоями

4.1 Введение

4.2 Предпосылки

4.2.1 Полисемантичность и суперпозиция

4.3 Sparse Autoencoders для интерпретации

4.3.1 Предпосылки использования SAE

4.3.2 Принцип работы SAE

4.3.3 История и применение

4.3.4 Современные модификации

4.3.5 Ограничения SAE

4.4 Связанные исследования

4.5 Метод

4.5.1 Сопоставление признаков

4.5.2 Folding SAE-весов

4.5.3 Композиция перестановок

4.6 Эксперименты

4.6.1 Используемые SAE и протокол сопоставления

Стр.

4.6.2 Датасеты и стратегия валидации

4.6.3 Метрики качества

4.6.4 Сопоставление признаков: количественно и качественно

4.6.5 Схожесть как метрика качества сопоставления

4.6.6 Персистентность признаков и композиция перестановок

4.6.7 SAE Match как «прореживание» слоёв

4.7 Выводы

Заключение

Список литературы

Список рисунков

Список таблиц

Приложение А. Вариационный вывод для механизма

раннего выхода

А.1 Постановка

А.2 Маргинальное правдоподобие и вариационная оценка

А.3 Практические замечания

Введение диссертации (часть автореферата) на тему «Обучение репрезентаций в обработке естественного языка»

Введение

Актуальность. Большие языковые модели (LLM), построенные на архитектуре трансформера, стремительно продвинули уровень технологий в понимании и генерации языка, однако их широкое внедрение ограничено четырьмя тесно связанными вызовами: вычислительная эффективность, масштабируемость по длине контекста, механистическая интерпретируемость.

Оригинальный трансформер обеспечивает параллельное обучение, но масштабируется квадратично по длине последовательности из-за механизма внимания (attention), что приводит к значительным затратам по времени и памяти на инференсе и ограничивает практическое применение на длинных последовятельностях.

Степень разработанности. Поскольку качество языковых моделей на основе трансформеров во многом зависит от ресурсов, доступных во время обучения и применения, ускорение и уменьшение моделей играет важную роль. Так как на обучение уходит большая часть ресурсов методы, которые позволяют ускорить модель без дополнительного обучения особенно распространены. Одним из таких методов является ранний выход из модели, который уменьшает число вычислений за счет динамического уменьшения числа слоев. Из-за архитектуры трансформера, состоящего из последовательности одинаковых слоев, такой подход особенно эффективен. В целом идея адаптивного числа вычислений для токенов была описана в [1], для трансформеров подобная идея представлена в FastBERT [2]. Однако, с улучшением генеративных авторегрессионных моделей и обучения внутри контекста, роль максимального размера контекста выходит на первый план. Начинают появлятся альтернативы, которые сохраняя возможность праллельного обучения демаонстрируют линейное или окололинейное время относительно длинны, основной работой здесь является Transformers are RNNs [3], а так же альтернативные верианты на основе SSM [4—7]. Для быстрого прогресса в области альтернативных трансформеру архитектур сейчас используют синтетические задачи. Было установлено [8], что для возможности обучения внутри контекста необходима способность модели копировать элементы последовательности, поэтому перед тестами на

реальных данных модель обучают на задачу "assotiative recall'^ только в случае успеха продолжают эксперименты. В целом раскрытие внутренних механизмов языковых моделей является одним из перспективных направлений. Одной из первых работ, предложившей использовать обучение разреженных автоэнкодеров для интерпретиации языковых моделей была [9]. Сейчас главным интересом в области интерпретируемости является ужу не выявление базовых компонент, а их изменение и взаимодействие с другими компонентами.

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

Для достижения поставленной цели необходимо решить следующие задачи:

1. Разобрать методы раннего выхода. Сравнить различные стретегии, изучив вариационный вывод разработать новую. Провести анализ полученной стратегии

2. Провести исследование альтернативных архитектур позволяющих добиться линейного прироста времени работы в зависимости от длины контекста. Опробировать свой метод сначала на задачи "assotiative recall затем обучить языковую модель и проверить на задачах связаных с текстом.

3. Изучить современные подходы к интерпретируемости языковых моделей. Изучить эволюцию представлений от слоя к слою для улучшения интепретации вычислений внутри языковой модели.

Научная новизна:

1. Предложен метод PALBERT (Pondering ALBERT) [10]. Метод позволяет адаптивно подбирать необходимое число вычислений для каждого примера. Кромет того, в отличии от предыдущих методов, позволяет совместно обучать такую модель, не основываясь на эвристиках. Подход продемонстрировал наилучшие метрики на наборе датасетов GLUE, он был применен к двум типам моделей: ALBERT и RoBERTa. Проведен анализ полученных моделей.

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

2. Изучены современные архитектуры, работающие за линейное время относительно длины последовательности [11]. В частности SSM-модели, которые, несмотря на преимущества в скорости, обладают недостатком в памяти. Разработан новый метод ReBASED. Устранены недостатки предыдущих моделей. Метод опрабирован на синтетической задаче assotiative recall, где продемонстрировал доминирующие характеристики предложенного метода. Обученная на текстовом датасете модель превосходит как обычный трансформер, так и безлайны работающие за линейное время.

3. Предложен новый метод исследования предобученных разреженных автоэнкодеров, он позволяет без дополнительного дообучения соединить латентные представления между собой [12]. Это позволяет отслеживать эволюцию признаков внутри модели и понять какие операции происходят внутри модели.

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

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

Основные положения, выносимые на защиту: диссертации заключаются в следующем:

1. PALBERT [10] - метод основанный на вариационном выводе превосходит как базовую модель, так и альтернативные методы раннего выхода. Метод позволяет адаптивно подбирать необходимое число вычислений для каждого примера. Кроме того,

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

2. ReBASED [11] - новая архитектура, работающая за линейное время относительно длины последовательности. Метод опробирован на синтетической задаче assotiative recall, где продемонстрировал доминирующие характеристики предложенного метода. Обученная на текстовом датасете модель превосходит как обычный трансформер, так и безлайны работающие за линейное время.

3. Предложен новый метод исследования предобученных разреженных автоэнкодеров, он позволяет без дополнительного дообучения соединить латентные представления между слоями [12]. Это позволяет отслеживать эволюцию признаков внутри модели и понять какие операции происходят внутри модели.

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

Исходный код методов опубликован в соответствующих репозиториях.

— https://github.com/tinkoff-ai/palbert

— https://github.com/corl-team/rebased

Апробация работы. Статьи представлены на международных конференциях:

— «Neural Information Processing Systems», 2022, Новый Орлеан, США.

— «Meeting of the Association for Computational Linguistics», 2024, Бангкок, Таиланд.

— «International Conference on Learning Representations», 2025, Сингапур, Сингапур.

Публикации. Основные научные результаты были опубликованы в четырех рецензируемых изданиях: три работы опубликованы на конференциях уровня Core A*, одна опубликована в журнале уровня К1, которые также входят в собственный перечень журналов МФТИ. Список публикаций:

1. PALBERT: Teaching ALBERT to Ponder / N. Balagansky, D. Gavrilov // Advances in Neural Information Processing Systems. Vol. 35 - 2022. - P. 14002-14012.

2. Linear Transformers with Learnable Kernel Functions are Better In-Context Models / Y.Aksenov, N.Balagansky, S. Lo Cicero Vaina, B. Shaposhnikov, A. Gorbatovski, D.Gavrilov // Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics. Vol. 1. - 2024. - P. 9584-9597.

3. Mechanistic Permutability: Match Features Across Layers / N. Balagansky, I. Maksimov, D. Gavrilov // International Conference on Learning Representations - 2025 - P. 84109-84126.

4. Обзор методов механистической интерпретируемости. / Н.Н.Балаганский // Труды МФТИ. - 2025. - Т. 17, № 3 - С. 95-103.

Личный вклад. Личное участие автора в представленных исследованиях заключается в следующем. В работе PALBERT [10] автору принадлежит идея использовать вариационный вывод для обучения модели раннего выхода, а также идея критерия раннего выхода и основная часть экспериментов и анализа. В работе ReBASED [11] авторы принадлежит идея модификации BASED, а так же был написан код экспериментов для языкового моделирования. В статье Mechanistic Permutability [12] автору принадлежит основная идея и часть экспериментов, не затрягивающую валидацию результатов с помощью внешней языковой модели. Автором полностью написана статья [13].

Объем и структура работы. Диссертация состоит из введения, 4 глав, заключения и 1 приложения. Полный объём диссертации составляет 97 страниц, включая 21 рисунок и 6 таблиц. Список литературы содержит 56 наименований.

Глава 1. Общие сведения о архитектрах современных языковых

моделей

Ниже фиксируется «общая картина» трансформера в той форме, на которую я буду ссылаться далее в диссертации. В основе лежат два источника: оригинальная статья Attention Is All You Need [14], в которой впервые была представлена архитектура трансформер, и взгляд с точки зрения изменения residual stream из серии Transformer Circuits [15], удобная для обсуждения внутренних механизмов больших языковых моделей. Автор будут рассматривать прежде всего авторегрессионный (decoder-only) вариант, характерный для LLM: модель читает префикс последовательности и предсказывает следующий токен (казуальный вариант), однако все выводы справедливы и для неказуальных языковых моделей.

1.1 Устройство трансформера

1.1.1 Токенизация

В дальнейшем под токеном подразумевается элемент последовательности после токенизации (например ВРЕ [16]). Токенизация управляет компромиссом между длиной последовательности и размером словаря; пробел, как правило, «приклеивается» к следующему токену; при необходимости используются служебные маркеры такие как <BOS> -начало последовательности, <EOS> - конец посладовательности, <PAD> - пад-токен для правильного вычисления для последовательностей разной длины, <UNK> - неизвестный символ.

Например, предложение «Трансформеры эффективны.» — может токенизироваться примерно так: ['Транс', 'форм', 'еры', ' э', 'ффектив', 'ны', '.']. Во всех формулах ниже длина контекста — это число токенов, а не символов или слов. Далее будем обозначать последовательность токенов как

{¿¡У^о1, где N - длина последовательности в токенах, при этом £ € ,

где V - размер вокабуляра модели. Типичный его размер колеблится от 30000 до 100000.

1.1.2 Эмбединги

После токенизации каждому токену сопоставляется свой вектор e = E(t) e G Rd, где d - размерность эмбедингов. Заметим что размер матрицы E G RdxV достаточно большой. Поэтому, если речь идет о максимально маленькой модели, то встает вопрос об уменьшении влияния этих параметров. Например, в модели ALBERT вместо одной операции происходит две: e = E(t) = E2E1 one-hot(t), где Ei G RrxV, E2 G Rdxr. Таким образом общее число параметров эмбедингов становится V • r + r • d против оригинальной V • d, при r << d такая параметризация эффктивнее.

1.1.3 Слой трансформера

После преобразования дискретных токенов в набор векторов {е;}^-1 эмбединги подаются на вход слою трансформера. Работу каждого слоя можно представить в виде рекуренты:

ht+1 = ht + L*(h), (1.1)

ho = (е0,...,бя-i), (1.2)

Набор (h0,..., hL-1) называют residual stream (дословно поток остатков, однако остаток здесь явно неверное слово в русском языке, это как раз ht как раз несет всю информацию в следующий слой). ht называют скрытым состоянием слоя t. Важно понимать, что вместо того, чтобы каждый раз перезаписывать информацию в поток, трансформер медленно сдвигает скрытое состояние в нужную сторону, читая и записывая части информации.

В слое трансформера содержится два блока - механизм внимания и полносвязная сеть.

1.1.4 Механизм внимания

Пусть на вход слою внимания поступает скрытое состояние И € . Из него с помощью трёх матриц весов WQк, € строятся

запросы, ключи и значения:

Я = К = htWк, V = htWv, (1.3)

где Я, К, V € х3'н, а ^ — размерность проекций одной головы внимания.

Далее вычисляется матрица «оценок внимания» как скалярные произведения запросов и ключей:

5 = , 5 € ^хЖ. (1.4)

Элемент отражает степень соответствия токена г (в роли запроса) токену ] (в роли ключа).

После нормировки через softmax по каждой строке получаем вероятностное распределение:

ехр(Б,) о ехр(Б^ш)

А = воЙ;тах(Б), А, = -. (1.5)

Матрица внимания А € хМ задаёт веса, с которыми каждый токен обращается ко всем остальным токенам последовательности.

Итоговое преобразование скрытых состояний выражается как

И = AV. (1.6)

Таким образом, каждое новое представление для позиции г есть линейная комбинация векторов-значений V, с весами а, = А,.

Многоголовое внимание. Чтобы захватывать разные типы зависимостей (например, грамматические связи и тематическую близость), используется несколько голов внимания Н. Для каждой головы И €

{1,..., Н} берутся свои матрицы весов Ж^,^,^. После вычисления

7 (Н)

щ. все результаты конкатенируются и проектируются в исходное

пространство:

MultiHead(ht) = Concat^^.., h(H)) Wq, (1.7)

где Wq G RHdhXd — матрица выходной проекции.

Интерпретация. Матрица A может рассматриваться как распределение вероятностей:

Кг = [ ]. (1.8)

В этом смысле внимание реализует механизм «взвешенного копирования» релевантных фрагментов контекста. В духе анализа Transformer Circuits [15] такие распределения позволяют выделять специализированные головы внимания — например, индукционные (induction heads), отвечающие за копирование шаблонов через несколько позиций, или головы, обеспечивающие синтаксические связи.

1.1.5 Полносвязная сеть

После блока внимания каждая позиция последовательности независимо обрабатывается двухслойной полносвязной сетью (feed-forward network, FFN). Если на вход подаётся ht G RNxd, то преобразование имеет вид:

FFN(ht) = о(htWi + bi) W2 + b2, (1.9)

где W1 G Rdxdff, W2 G Rdffxd, смещение b1 G Rdf, b2 G Rd, а a(^) — нелинейная активация. Обычно выбирается ReLU или GELU, причём размер скрытого слоя dff существенно больше d, например dff ~ 3d.

Gated-активации. В ряде современных моделей (Gated Linear Units, SwiGLU, GeGLU) используются gated-варианты FFN, где вместо одной активации вводится дополнительный управляющий множитель. Простейший вариант можно записать так:

FFNGLu(ht) = (htWa + ba) 0 О(htWb + bb), (1.10)

где 0 обозначает поэлементное произведение, а Wa,Wb € Шdхdff — разные матрицы весов.

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

Роль FFN. Если механизм внимания отвечает за сбор информации из всей последовательности, то полносвязная сеть служит для нелинейного преобразования и повышения размерности признакового пространства на уровне отдельного токена. Вместе они обеспечивают баланс между глобальными зависимостями и локальной обработкой признаков.

1.1.6 Последний слой

Финальным блоком языковой модели является LM-head — линейный слой, который преобразует скрытые состояния последнего слоя hi £ RNxd в распределение вероятностей по словарю модели.

На практике LM-голова представляет собой матрицу Wlm £ RdxV, где V — размер словаря:

Z = hiWm + bim, Z £ . (1.11)

После этого вычисляется вероятностное распределение для следующего токена через softmax:

p(t | t0,..., ti-i) = softmax(zi). (1.12)

В авторегрессионной постановке на шаге i используется только скрытое состояние hij, а маскирование в механизме внимания гарантирует, что предсказание зависит только от предыдущих токенов to,... ,ti-i.

Weight tying. Часто применяют технику weight tying, когда матрица Wlm идентична матрице эмбеддингов Eт. В этом случае LM-head

фактически «проверяет», какой токен наиболее близок в скрытом пространстве к текущему представлению h^:

z = hwET. (1.13)

Такой приём уменьшает число параметров и эмпирически улучшает качество, так как входные и выходные представления словаря разделяют одно пространство.

Роль LM-head. LM-head является звеном, связывающим непрерывные представления в residual stream с дискретным пространством токенов. Таким образом, именно через LM-head модель осуществляет основную задачу — генерацию последовательности токенов, приближая распределение данных pe(t) к распределению в обучающем корпусе.

1.2 Виды языковых моделей

Архитектура трансформера может использоваться в различных режимах обучения. Наиболее распространены два типа языковых моделей: Masked Language Models (MLM) и Causal Language Models (CausalLM). Они различаются по тому, как формируется задача предсказания токенов.

1.2.1 Маскированная языеовая модели (MLM)

Модели данного типа обучаются восстанавливать токены, случайным образом замаскированные в последовательности. Пусть вход состоит из последовательности токенов (t0,...,tN-i), в которой некоторые позиции заменены специальным маркером <MASK>. MLM обучается максимизировать правдоподобие:

Lmlm = - ^ logp0(ti | to,..., <MASK>, ...,tN-i), (1.14)

i£M

где M — множество замаскированных позиций.

Таким образом, модель видит как левые, так и правые контексты токена, что делает её двунаправленной. Классическим примером является BERT [17]. В текущей работе будут использованы модели ALBERT [18] и RoBERTa [19], которые также относятся к этому типу. Такой режим обучения хорошо подходит для задач понимания текста (классификация, поиск, извлечение информации), но менее естественен для генерации последовательностей.

1.2.2 Последовательная языковыя модель (Causal LM)

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

N-1

LcausalLM = - ^ l0gР0^ I ¿0,-- -,ti-l). (1.15)

i=1

Здесь используется каузальное внимание: матрица A маскируется так, чтобы токен i мог обращаться только к токенам с индексами ^ i, но не к будущим позициям. Благодаря этому CausalLM естественным образом применяются для генерации текста. Примеры: GPT-семейство моделей [20].

Сравнение. MLM обеспечивает богатые двусторонние представления и лучше подходит для задач понимания текста, тогда как CausalLM ближе к реальной задаче генерации и является основой большинства современных LLM. Интересно отметить, что обе парадигмы могут быть объединены: например, T5 [21] использует span-masking (разновидность MLM), а модели типа GPT [20] применяют исключительно авторегрессию.

1.3 Выводы

Трансформер является основной архитектурой в обработке естественного языка. Он состоит из последовательных одинаковых блоков

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

Глава 2. Ранний выход в трансформерах

Механизм раннего выхода (Early Exit) позволяет динамически регулировать глубину вычислений в языковой модели в зависимости от конкретного входного примера. Это направление исследований связано с задачей повышения вычислительной эффективности LLM и уменьшения их задержки на инференсе. В данной главе мы опишем основные идеи раннего выхода, а также рассмотрим их реализацию на примере архитектуры PALBERT [10].

2.1 Связанные исследования

Проблема избыточных вычислений в языковых моделях на основе трансформеров привела к активному развитию методов раннего выхода (early exit), которые позволяют адаптировать глубину вычислений к сложности входного примера. Ранние подходы, такие как BranchyNet [22], DeeBERT [23] и FastBERT [2], основывались на энтропийных критериях — модель завершала вычисления, когда распределение вероятностей на выходе становилось достаточно уверенным. Однако впоследствии было показано, что такие методы подвержены явлению overthinking (избыточного размышления), при котором энтропия продолжает снижаться, даже если предсказание уже неверно [24]. Кроме того, энтропийные критерии плохо переносятся на регрессионные задачи.

Для решения этих проблем был предложен метод Patient BERT [24], использующий консенсус между предсказаниями различных классификаторов на промежуточных слоях модели. Вычисления останавливаются, если несколько соседних слоёв выдают одинаковый результат. Этот подход оказался более устойчивым, особенно для задач классификации, и стал основой для дальнейших модификаций, таких как LeeBERT [25] и GloPFE [26], которые расширили идею консенсусного выхода за счёт дополнительных эвристик.

Параллельно развивались вариационные методы адаптивного времени вычислений. Наиболее заметным из них является PonderNet [27], в котором индекс выходного слоя рассматривается как латентная переменная. Модель обучается максимизацией нижней оценки правдоподобия, а во время инференса выбор момента выхода осуществляется стохастически путём сэмплирования из апостериорного распределения вероятностей выхода. Несмотря на теоретическую изящность, стохастическая природа критерия порождает большую дисперсию индекса выходного слоя и ухудшает стабильность результатов.

В работе PALBERT предложено устранить этот недостаток за счёт детерминированного критерия выхода Q-exit, который вычисляет кумулятивную функцию распределения вероятности выхода и завершает вычисления, когда она превышает заданный порог. В отличие от PonderNet, данный подход полностью устраняет случайность в процессе вывода и обеспечивает более предсказуемое поведение модели. Дополнительно модифицирована архитектура вспомогательных Lambda-слоёв, предсказывающих вероятность выхода, что позволило учитывать динамику скрытых состояний и количество уже выполненных итераций. В результате PALBERT показал улучшение производительности по сравнению с PABEE и PonderNet при сохранении сопоставимого качества с базовой моделью ALBERT [18].

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

2.2 Мотивация

Современные языковые модели строятся на десятках и сотнях слоёв трансформера. Однако практические наблюдения показывают, что сложность входных примеров сильно различается: для одних задач достаточно нескольких слоёв, в то время как для других необходима полная глубина модели. При фиксированном числе слоёв модель всегда выполняет

одинаковый объём вычислений, что ведёт к избыточным затратам и феномену overthinking: использование излишнего числа слоёв может даже снижать качество предсказаний.

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

2.3 Базовые подходы

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

— Энтропийные критерии. Методы FastBERT [2], DeeBERT [23]

и др. анализируют энтропию распределения предсказаний: если она достаточно низкая (модель уверена в ответе), выполнение останавливается.

— Консенсусные методы. Подход PABEE [24] использует несколько классификаторов на разных уровнях. Если несколько слоёв подряд дают одинаковый ответ, модель выходит преждевременно.

— Вариационные методы. Архитектура PonderNet [27] рассматривает индекс выходного слоя как скрытую переменную и обучает распределение вероятностей выхода на каждом слое. Однако в изначальной версии для инференса применялось стохастическое семплирование, что приводило к высокой дисперсии в индексе выхода.

Метод PABEE. Один из первых практических подходов к раннему выходу — метод PABEE ("Patient Early Exit") [24]. Его идея состоит в том, что на каждой из L слоёв трансформера обучается свой классификатор модель прекращает вычисления, если несколько соседних классификаторов дают одинаковый предсказанный класс. Формально, если

Ci(hi) = Cj+i (hj+i) = ••• = Ci+t(hi+t),

то на слое i +t выполняется выход, где t — гиперпараметр "терпеливости" (patience). Таким образом, PABEE снижает переобучение на поздних слоях ("overthinking") и экономит вычисления. Однако данный метод чувствителен к качеству промежуточных классификаторов, которые не всегда хорошо обучаются на малых датасетах.

Метод PonderNet. Другой подход, PonderNet [27], трактует номер слоя выхода как латентную переменную i. Для каждого слоя i вводится дополнительный модуль Л (hi), оценивающий вероятность Л^ того, что на этом слое нужно завершить вычисления. Это порождает геометрическое распределение вероятностей выхода:

i-i

p(i | x) = ^ П(1 - Л,-). j=i

При обучении оптимизируется вариационная нижняя оценка маргинального правдоподобия:

L(x,y) = logp(y | x,i)] - в KLp(i | x) || p(i | Л)),

где p(i | Л) — априорное геометрическое распределение, а в — коэффициент регуляризации.

На инференсе в оригинальной версии PonderNet индекс выхода i семплируется из p(i | x), что вносит существенную дисперсию в результаты: одна и та же последовательность может завершиться на разных слоях. Этот недостаток позже был решён в работах PALBERT/PRoBERTa с помощью детерминированного критерия Q-exit.

Список литературы диссертационного исследования кандидат наук Балаганский Никита Николаевич, 2026 год

Список литературы

1. Adaptive Computational Time For Recurrent Neural Networks/Graves, A. //arXiv preprint arXiv:1603.08983. — 2016. URL: https://arxiv.org/abs/ 1603.08983 (visited on September 9 2025).

2. FastBERT: a self-distilling BERT with adaptive inference time/Liu, W., Zhou, P., Wang, Z., Zhao, Z., Deng, H., Ju, Q.//Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics. — 2020.

— Vol. 1. — P. 6035—6044.

3. Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention / Katharopoulos, A., Vyas, A., Pappas, N., Fleuret, F. // Proceedings of the 37th International Conference on Machine Learning (ICML). — 2020. — P. 5156—5165.

4. Efficiently Modeling Long Sequences with Structured State Spaces/Gu, A., Goel, K., Re, C.//International Conference on Learning Representations.

— 2022. URL: https://openreview.net/forum?id=uYLFoz1vlAC (visited on September 9 2025).

5. Efficiently Modeling Long Sequences with Structured State Spaces/Gu, A., Goel, K., Re, C.//International Conference on Learning Representations.

— 2022. URL: https://openreview.net/forum?id=uYLFoz1vlAC (visited on September 9 2025).

6. Simplified State Space Layers for Sequence Modeling / Smith, J. T., Warrington, A., Linderman, S.//International Conference on Learning Representations. — 2023. URL: https : / / openreview . net / forum ? id = Ai8Hw3AXqks (visited on September 9 2025).

7. Mamba: Linear-Time Sequence Modeling with Selective State Spaces /Gu, A., Dao, T.//First Conference on Language Modeling. — 2024. URL: https://openreview.net/forum?id=tEYskw1VY2 (visited on September 9 2025).

8. In-context Learning and Induction Heads / Olsson, C., Elhage, N., Nanda, N., Joseph, N., DasSarma, N., Henighan, T., Mann, B., Askell, A., Bai, Y., Chen, A., Conerly, T., Drain, D., Ganguli, D., Hatfield-Dodds, Z., Hernandez, D., Johnston, S., Jones, A., Kernion, J., Lovitt, L., Ndousse, K., Amodei, D., Brown, T., Clark, J., Kaplan, J., McCandlish, S., Olah, C.// Transformer Circuits Thread. — 2022. URL: https://transformer-circuits. pub/2022/in-context-learning-and-induction-heads/index.html (visited on September 9 2025).

9. Transformer Visualization via Dictionary Learning: Contextualized Embedding as a Linear Superposition of Transformer Factors/Yun, Z., Chen, Y., Olshausen, B., LeCun, Y.//Proceedings of the 2nd Workshop on Knowledge Extraction and Integration for Deep Learning Architectures (DeeLIO). — 2021. — P. 1—10.

10. PALBERT: Teaching ALBERT to Ponder/Balagansky, N., Gavrilov, D.// Advances in Neural Information Processing Systems. — 2022. — Vol. 35.

— P. 14002—14012.

11. Linear Transformers with Learnable Kernel Functions are Better In-Context Models/Aksenov, Y., Balagansky, N., Lo Cicero Vaina, S., Shaposhnikov, B., Gorbatovski, A., Gavrilov, D.//Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics. — 2024.

— Vol. 1. — P. 9584-9597.

12. Mechanistic Permutability: Match Features Across Layers/Balagansky, N., Maksimov, I., Gavrilov, D.//International Conference on Learning Representations. — 2025. — P. 84109—84126.

13. Обзор области механистической интерпретируемости/Балаганский, Н. Н. //Труды МФТИ. — 2025. — Vol. 17. — P. 95—103.

14. Attention is All you Need/Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., Polosukhin, I.//Advances in Neural Information Processing Systems. — 2017. — Vol. 30. — P. 5998-6008.

15. A Mathematical Framework for Transformer Circuits / Elhage, N., Nanda, N., Olsson, C., Henighan, T., Joseph, N., Mann, B., Askell, A., Bai, Y., Chen, A., Conerly, T., DasSarma, N., Drain, D., Ganguli, D., Hatfield-Dodds, Z., Hernandez, D., Jones, A., Kernion, J., Lovitt, L., Ndousse, K., Amodei, D., Brown, T., Clark, J., Kaplan, J., McCandlish, S., Olah, C.//Transformer Circuits Thread. — 2021. URL: https : / / transformer - circuits . pub / 2021 / framework / index . html (visited on September 9 2025).

16. Neural Machine Translation of Rare Words with Subword Units /Sennrich, R., Haddow, B., Birch, A.//Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics. — 2015. — Vol. 1. — P. 1715—1725.

17. BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding/Devlin, J., Chang, M.-W., Lee, K., Toutanova, K.// Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics. — 2019. — Vol. 1. — P. 4171—4186.

18. ALBERT: A Lite BERT for Self-supervised Learning of Language Representations/Lan, Z., Chen, M., Goodman, S., Gimpel, K., Sharma, P., Soricut, R.//International Conference on Learning Representations. — 2020. URL: https://openreview.net/forum?id=H1eA7AEtvS (visited on September 9 2025).

19. RoBERTa: A Robustly Optimized BERT Pretraining Approach/Liu, Y., Ott, M., Goyal, N., Du, J., Joshi, M., Chen, D., Levy, O., Lewis, M., Zettlemoyer, L., Stoyanov, V.//arXiv preprint arXiv:1907.11692. — 2019. URL: https://arxiv.org/abs/1907.11692 (visited on September 9 2025).

20. Language Models are Unsupervised Multitask Learners/Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., Sutskever, I. — 2019. URL: https: / / cdn. openai. com/better- language- models /language _ models _ are _ unsupervised_multitask_learners.pdf (visited on September 9 2025).

21. Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer/Raffel, C., Shazeer, N. M., Roberts, A., Lee, K., Narang, S.,

Matena, M., Zhou, Y., Li, W., Liu, P. J.//Journal of machine learning research. — 2019. — Vol. 21. — P. 5485—5551.

22. BranchyNet: Fast inference via early exiting from deep neural networks /Teerapittayanon, S., McDanel, B., Kung, H. T.//Proceedings of the 23rd International Conference on Pattern Recognition (ICPR). — 2016. — P. 2464—2469.

23. DeeBERT: Dynamic early exiting for accelerating BERT inference/Xin, J., Tang, R., Lee, J., Yu, Y., Lin, J.//Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics. — 2020. — Vol. 1. — P. 2246—2251.

24. BERT Loses Patience: Fast and Robust Inference with Early Exit /Zhou, W., Xu, C., Ge, T., McAuley, J., Xu, K., Wei, F.//Advances in Neural Information Processing Systems. — 2020. — Vol. 33. — P. 18330-18341.

25. LeeBERT: Learned Early Exit for BERT with Cross-Level Optimization /Zhu, W.//Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics. — 2021. — Vol. 1. — P. 2968—2980.

26. A Global Past-Future Early Exit Method for Accelerating Inference of Pre-trained Language Models/Liao, K., Zhang, Y., Ren, X., Su, Q., Sun, X., He, B.//Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics. — 2021. — Vol. 1. — P. 2013—2023.

27. PonderNet: Learning to Ponder/Banino, A., Balaguer, J., Blundell, C.// 8th ICML Workshop on Automated Machine Learning (AutoML). — 2021. URL: https : / / openreview . net / forum ? id = 1EuxRTe0WN (visited on September 9 2025).

28. GLUE: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding/Wang, A., Singh, A., Michael, J., Hill, F., Levy, O., Bowman, S.//Proceedings of the 2018 EMNLP Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP. — 2018. — P. 353-355.

29. The Llama 3 Herd of Models/team, L.//arXiv preprint arXiv:2407.21783. — 2024. URL: https://arxiv.org/abs/2407.21783 (visited on September 2025).

30. Rethinking Attention with Performers/Choromanski, K., Likhosherstov, V., Dohan, D., Song, X., Gane, A., Sarlos, T., Hawkins, P., Davis, J., Belanger, D., Colwell, L., Weller, A.//International Conference on Learning Representations (ICLR). — 2021. URL: https://openreview.net/ forum?id=Ua6zuk0WRH (visited on September 9 2025).

31. Random Feature Attention / Peng, H., Pappas, N., Yogatama, D., Schwartz, R., Smith, N., Kong, L.//International Conference on Learning Representations. — 2021. URL: https ://openreview . net / forum ? id = QtTKTdVrFBB (visited on September 9 2025).

32. cosFormer: Rethinking Softmax In Attention/Qin, Z., Sun, W., Deng, H., Li, D., Wei, Y., Lv, B., Yan, J., Kong, L., Zhong, Y.//International Conference on Learning Representations. — 2022. URL: https : / / openreview.net/forum?id=Bl8CQrx2Up4 (visited on September 9 2025).

33. Linear Transformers Are Secretly Fast Weight Programmers/Schlag, I., Irie, K., Schmidhuber, J.//Proceedings of the 38th International Conference on Machine Learning. — 2021. — Vol. 139. — P. 9355—9366.

34. Long Range Arena : A Benchmark for Efficient Transformers/Tay, Y., Dehghani, M., Abnar, S., Shen, Y., Bahri, D., Pham, P., Rao, J., Yang, L., Ruder, S., Metzler, D.//International Conference on Learning Representations. — 2021. URL: https ://openreview . net / forum ? id = qVyeW-grC2k (visited on September 9 2025).

35. Never Train from Scratch: Fair Comparison of Long-Sequence Models Requires Data-Driven Priors / Amos, I., Berant, J., Gupta, A. // International Conference on Representation Learning. — 2024. — P. 1843—1861.

36. In-Context Learning and Induction Heads / Olsson, C., Elhage, N., Nanda, N., Joseph, N., DasSarma, N., Henighan, T., [et al.] — 2022. URL: https: //transformer-circuits.pub/2022/in-context-learning-and-induction-heads/index.html (visited on September 9 2025).

37. Zoology: Measuring and Improving Recall in Efficient Language Models /Arora, S., Eyuboglu, S., Timalsina, A., Johnson, I., Poli, M., Zou, J., Rudra, A., Re, C.//The Twelfth International Conference on Learning Representations. — 2024. URL: https : / / openreview . net / forum ? id = LY3ukUANko (visited on September 9 2025).

38. The Hedgehog and the Porcupine: Expressive Linear Attentions with Softmax Mimicry/Arora, S., Eyuboglu, S., Zhang, M., Timalsina, A., Alberti, S., Zou, J., Rudra, A., Re, C.//International Conference on Machine Learning. — 2024. — P. 53551—53580.

39. Layer Normalization/Ba, J. L., Kiros, J. R., Hinton, G. E.//arXiv preprint arXiv:1607.06450. — 2016. URL: https : / / arxiv . org / abs / 1607 . 06450 (visited on September 9 2025).

40. The Pile: An 800GB Dataset of Diverse Text for Language Modeling /Gao, L., Biderman, S., Black, S., Golding, L., Hoppe, T., Foster, C., Phang, J., [et al.]//arXiv preprint arXiv:2101.00027. — 2020. URL: https: //arxiv.org/abs/2101.00027 (visited on September 9 2025).

41. RWKV: Reinventing RNNs for the Transformer Era/Peng, B., Alcaide, E., Anthony, Q., Albalak, A., Biderman, S., Cao, H., [et al.]//Findings of the Association for Computational Linguistics: EMNLP 2023. — 2023. — P. 14048—14077.

42. Parallelizing Linear Transformers with the Delta Rule over Sequence Length/Yang, S., Wang, B., Zhang, Y., Shen, Y., Kim, Y.//Advances in Neural Information Processing Systems. — 2024. — Vol. 37. — P. 115491—115522.

43. Hyena Hierarchy: Towards Larger Convolutional Language Models /Poli, M., Massaroli, S., Nguyen, E., Fu, D. Y., Dao, T., Baccus, S., Bengio, Y., Ermon, S., Re, C.//Proceedings of the 40th International Conference on Machine Learning. — 2023. — Vol. 202. — P. 28043—28078.

44. Show Your Work: Improved Reporting of Experimental Results/Dodge, J., Gururangan, S., Card, D., Schwartz, R., Smith, N. A.//Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing. — 2019. — Vol. 1. — P. 2185—2194.

45. Toy Models of Superposition / Elhage, N., Hume, T., Olsson, C., Schiefer, N., Henighan, T., Kravec, S., Hatfield-Dodds, Z., Lasenby, R., Drain, D., Chen, C., Grosse, R., [et al.] — 2022. URL: https://transformer-circuits.pub/2022/toy_model/index.html (visited on September 9 2025).

46. Towards Monosemanticity: Decomposing Language Models with Dictionary Learning/Bricken, T., Templeton, A., Batson, J., Chen, B., Jermyn, A., Conerly, T., Turner, N., Denison, C., Askell, A., Lasenby, R., Olah, C., [et al.] — 2023. URL: https://transformer-circuits.pub/2023/ monosemantic-features/index.html (visited on September 9 2025).

47. Scaling Monosemanticity: Extracting Interpretable Features from Claude 3 Sonnet/Templeton, A., Conerly, T., Marcus, J., Lindsey, J., Bricken, T., Chen, B., Pearce, A., Citro, C., Freeman, C. D., Olah, C., Henighan, T., [et al.] — 2024. URL: https://transformer-circuits.pub/2024/scaling-monosemanticity/index.html (visited on September 9 2025).

48. Gemma Scope: Open Sparse Autoencoders Everywhere All at Once on Gemma 2/Lieberum, T., Rajamanoharan, S., Conmy, A., Smith, L., Sonnerat, N., Varma, V., Kramár, J., Dragan, A., Shah, R., Nanda, N.// arXiv preprint arXiv:2408.05147. — 2024. URL: https://arxiv.org/abs/ 2408.05147 (visited on September 9 2025).

49. Linear Algebraic Structure of Word Senses, with Applications to Polysemy /Arora, S., Li, Y., Liang, Y., Ma, T., Risteski, A.//Transactions of the Association for Computational Linguistics. — 2018. — Vol. 6. — P. 483-495.

50. Gemma 2: Improving Open Language Models at a Practical Size/Gemma Team, Riviere, M., Pathak, S., Sessa, P. G., Hardin, C., Bhupatiraju, S., Hussenot, L., Mesnard, T., Shahriari, B., Ramé, A., Ferret, J., Liu, P., Tafti, P., Friesen, A., [et al.]//arXiv preprint arXiv:2408.00118. — 2024. URL: https://arxiv.org/abs/2408.00118 (visited on September 9 2025).

51. Neuronpedia: Interactive Reference and Tooling for Analyzing Neural Networks/Lin, J. — 2023. URL: https://www.neuronpedia.org (visited on September 9 2025).

52. Jumping Ahead: Improving Reconstruction Fidelity with JumpReLU Sparse Autoencoders/Rajamanoharan, S., Lieberum, T., Sonnerat, N., Conmy, A., Varma, V., Kramar, J., Nanda, N.//arXiv preprint arXiv:2407.14435. — 2024. URL: https : / / arxiv . org / abs / 2407 . 14435 (visited on September 9 2025).

53. Git Re-Basin: Merging Models Modulo Permutation Symmetries /Ainsworth, S. K., Hayase, J., Srinivasa, S.//International Conference on Learning Representations (ICLR). — 2022. URL: https://openreview. net/forum?id=CQsmMYmlP5T (visited on September 9 2025).

54. OpenWebText Corpus/Gokaslan, A., Cohen, V., Pavlick, E., Tellex, S. — 2019. URL: http://Skylion007.github.io/OpenWebTextCorpus (visited on September 9 2025).

55. GitHub Small Near-Dedup (Code2) Dataset/Allal, L. B., collaborators — 2024. URL: https: / / huggingface. co / datasets / loubnabnl / github- small -near-dedup (visited on September 9 2025).

56. beta-VAE: Learning Basic Visual Concepts with a Constrained Variational Framework/Higgins, I., Matthey, L., Pal, A., Burgess, C., Glorot, X., Botvinick, M., Mohamed, S., Lerchner, A.//International Conference on Learning Representations. — 2017. URL: https://openreview.net/forum? id=Sy2fzU9gl (visited on September 9 2025).

Список рисунков

1. Рисунок 2.1. — Сравнение стохастического выхода (sampling) и

детерминированного Q-exit в PALBERT [10]. ... 22

2. Рисунок 2.2. — Зависимость качества от порога q. Рисунок из [10]. 26

3. Рисунок 2.3. — Компромисс скорость-качество для PALBERT и

PABEE на SST-2 и RTE [10; 24]............ 28

4. Рисунок 2.4. — Гистограммы индексов раннего выхода при

разных значениях Л априорного распределения.

Слева результаты для датасета MRPC, справа

для RTE. Рисунок из [10]................ 30

5. Рисунок 2.5. — Плотности постериорного распределения,

полученного с разными сидами. Рисунок из [10]. . 31

6. Рисунок 3.1. — Интуитивная разница между Based [38] и

ReBased [11]. В то время как Based пытается апроксимировать экспоненту, ReBased позволяет обучать нужную функцию схожести внутри модели, что дает большую гибкость. Рисунок из [11]............................. 38

7. Рисунок 3.2. — Разница архитектур Based (слева) и ReBased

(справа). Рисунок из [11]................ 41

8. Рисунок 3.3. — Качество различных модификаций модели

Based [38]. Пердложенный метод (ReBased [11]

является наилучшим вариантом). Рисунок из

[11]. [20pt]........................ 44

9. Рисунок 3.4. — Примеры матриц внимания для Based и ReBased.

ReBased демонстрирует более чёткие паттерны

внимания, приближаясь к поведению

классического softmax-внимания........... 46

10. Рисунок 3.5. — Ожидаемое качество модели в зависимомти от

размера сетки для перебора гиперпараметров. ... 48

11. Рисунок 4.1. — Сглаживание производной для обучения

порогового значения 6................. 54

12. Рисунок 4.2. — Для сопоставления признаков разных слоев

используеься матрица перестановки, найденная с

помощью алгоритма LAP, аналогично [53].

Рисунок из [12]...................... 58

13. Рисунок 4.3. — Динамика норм скрытых состояний и значений 6

в JumpReLU по слоям. Рисунок из [12]........ 60

14. Рисунок 4.4. — Демонстрация сопоставления признаков между

слоями. Рисунок из [12]................. 61

15. Рисунок 4.5. — Пример фичей которые были найдены с помощью

предложенного алгоритма. Фича на слое 20 отвечает за время, сопоставленные фичи на других слоях имеют похожие описания. Визуализация сделана с помощью Neuroпpedia [51]. 65

16. Рисунок 4.6. — Динамика схожести сопоставленных направлений

в зависимости от стартового слоя. Рисунок из [12]. 66

17. Рисунок 4.7. — Попарное сравнение сопоставлений в зависимости

от слоя. Рисунок из [12]................ 67

18. Рисунок 4.8. — Попарное сравнение сопоставлений в зависимости

от слоя. Рисунок из [12]................ 69

19. Рисунок 4.9. — Сравнение прямого сопоставления и композиции

перестановок. Рисунок из [12]............. 71

20. Рисунок 4.10. —Сравнение методов сопоставление для

автоэнкодеров. Рисунок из [12]............ 73

21. Рисунок 4.11. —Сравнение метода на различных доменах.

Рисунок из [12]...................... 75

Список таблиц

1. Таблица 1.

2. Таблица 2.

3. Таблица 3.

4. Таблица 4.

5. Таблица 5.

6. Таблица 6.

Сранвнение PALBERT с похожими методами на наборе задач GLUE [28]. В колонке macro мы показываем среднее значение по всем задачам. Лучшие результаты показаны жирным шрифтом. Мы не применяли жирный шрифт для базовой модели, результаты базовой модели приведены для

справки........................... 23

Сравнение модели PRoBERTa с современными подходами на бенчмарке GLUE [28]. В колонке Macro приведены усреднённые результаты по задачам. Жирным выделены наилучшие значения. Для строк RoBERTa [19] выделение не применялось, так как в этих случаях ранний выход не использовался, и результаты приведены только

для справки. Таблица из [10]. [10pt].......... 24

Абляционное исследование предложенной архитектуры PALBERT. «Lambda LR» означает настройку Lambda-слоя с использованием собственного learning rate; «3-layer Lambda» указывает на замену однослойного MLP в Lambda-слое на трёхслойный; «h concat.» обозначает конкатенацию двух скрытых состояний, подаваемых на вход в Lambda-слой. Таблица из [10]. [10pt]......................... 25

Перплексия на датасете Pile. ReBased улучшает

результат по AR-токенам по сравнению с Based,

хотя до классического Attention остаётся разрыв. . . 45

Результаты 1-shot на задачах LM Evaluation

Harness. ReBased стабильно превосходит Based. ... 46

Результаты 1-shot на задачах SuperGLUE. ReBased

показывает рост точности по сравнению с Based. . . 46

Приложение А Вариационный вывод для механизма раннего выхода

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

А.1 Постановка

Пусть задан вход x (контекст токенов) и целевой выход y (например, следующий токен в авторегрессионной постановке). Рассмотрим стек из n одинаковых по типу трансформерных слоёв.1 Параметризация практически идентична PonderNet [27]. Обозначим через i Е {1,... ,n} латентный индекс слоя, на котором выполняется предсказание.

Модель определяет:

— условные распределения на выходе каждого уровня: p(y | x, i);

— апостериорное распределение на индексе выхода p(i | x), индуцированное Lambda-слоями.

Для Lambda-слоёв вводится вероятность выхода на шаге i:

Лг = A([hi,hi_i]) Е (0,1), (А.1)

где hi — скрытое состояние (residual stream) после i-го слоя, а Л(-) — небольшая MLP (обычно с tanh), принимающая, например, конкатенацию [hi,hi_1]. Тогда индуцированное распределение выхода по слоям имеет обобщённо-геометрический вид:

i_ i

p(i | x) = ^ ^(1 _ Л,-), i = 1,... ,n _ 1, (А.2)

_i

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

а для последнего уровня применяется нормировка

и—1

р(п | х) = 1 — ^^р(г | х), (А.3)

¿=1

чтобы ^и=1 Р(г I х) = 1 при конечном п.

В качестве приора на индекс выхода используем геометрическое распределение

р(г | Л) = Л(1 — Л)г—1, г = 1,...,п — 1, (А.4)

а для п-го уровня — р(п | Л) = 1—^П=—1 Р(г I Л), что обеспечивает нормировку при конечном п. Здесь Л € (0,1) — гиперпараметр приора (или набор таковых).

А.2 Маргинальное правдоподобие и вариационная оценка

Нас интересует маргинальное правдоподобие данных:

n n

logp(y I x) = log£p(v,i I x) = p(y I x,i)p(i I x). (A.5)

i=1 i=1

Непосредственно максимизировать (A.5) неудобно и ведет к колапсу распределения; используем стандартный приём вариационного вывода. Введём произвольное вариационное распределение q(i | x) на {1,...,n} и запишем тождество:

_^ р( y I x i) p(i I A) _т p(i I x)

1ogp(y 1 x) = tog^j q(i 1 x)-- + tog^j q(i 1 x) рщ—)

n

q(i 1 x)

p(i I x)

x)

i=i P(i 1 A)'

Применяя неравенство Йенсена к первому логарифму, получаем нижнюю оценку (ELBO):

1ogP(y 1 x) > Eq(i|x)[1ogP(y 1 x,i)] - KL(q(i 1 x) || p(i 1 A)) , (A.6)

где KL(q || p) = qi 1og j — дивергенция Кульбака-Лейблера. Это классическая форма ELBO: первый член — средний по q лог-правдоподобия

«локальных» предсказаний, второй — регуляризатор, притягивающий q к приору.

Выбор q(i | x). В PonderNet-подобной постановке естественно брать q(i | x) равным модельному апостериору p(i | x) из (А.2):

q(i | x) = p(i | x). (А.7)

Тогда (А.6) принимает вид

logp(y | x) ^ [logp(y | x,i)] - KL(P( | x) II | A)). (А.8)

в-вес регуляризации. На практике полезно вводить коэффициент в > 0, регулирующий силу регуляризации (аналогично в-VAE [56]):

[logp(y | x,i)] - в KL(p( | x) I | A)) ^ logp(y | x).

(А.9)

Максимизация L по параметрам трансформера, голов предсказания p(y | x,i) и Lambda-слоёв, определяющих p(i | x), реализует совместное обучение: модель учится как предсказывать на разных глубинах, так и останавливаться там, где это уместно, при этом p(i | x) мягко притягивается к «геометрическому» приору p(i | A).

А.3 Практические замечания

Замкнутая форма первого слагаемого. Первый член в (А.9) удобно вычислять без сэмплирования индекса, как взвешенную сумму:

n

Ei~p(-|x)[l°g p(y| x,i)] = p(i | x) log p(y| x,i). (А.10)

i=i

Это избавляет от дисперсии Монте-Карло и стабилизирует обучение.

Техническая нормировка последнего уровня. Из-за конечного n приор и апостериор для i = n обычно доводят до нормировки:

n—1 n—1

p(n | A) = 1 — ^^p(i | A), p(n | x) = 1 — ^^p(i | x). (А.11)

i=i i=i

Также допускается «усечение» хвоста при вычислении KL, если кумулятивная масса p(i I x) уже близка к 1.

Детерминированный критерий выхода (Q-exit). Хотя сэмплирование i ~ р(; I x) согласовано с вариационной трактовкой, на инференсе можно использовать детерминированное правило по CDF: выполнить выход на наименьшем i, для которого Yj=1 P(j I x) ^ q (порог q £ (0,1)). Это устраняет дисперсию индекса выхода при сохранении обученной апостериорной структуры.