Большие языковые модели TinyBERT — дистилляция BERT для NLP

Масштабные модели TinyBERT - улучшенные версии BERT для обработки естественного языка

Разблокирование потенциала дистилляции трансформатора в LLM-моделях

Введение

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

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

Принимая во внимание эту проблему, были разработаны специальные методы сжатия LLM. В этой статье мы сосредоточимся на дистилляции трансформатора, которая привела к созданию уменьшенной версии BERT, называемой TinyBERT. Кроме того, мы разберем процесс обучения в TinyBERT и несколько тонкостей, которые делают TinyBERT таким надежным. Эта статья основана на официальной статье о TinyBERT.

Основная идея

Недавно мы уже рассмотрели, как работает дистилляция в DistilBERT: в краткой форме, целевая функция потери модифицируется таким образом, чтобы предсказания студента и учителя были похожи. В случае с TinyBERT, изменения идут еще дальше: целевая функция потери учитывает не только ЧТО производят оба моделя, но и КАК получаются предсказания. Согласно статье, целевая функция потери TinyBERT состоит из трех компонентов, которые охватывают различные аспекты обоих моделей:

  1. вывод слоя вложений
  2. скрытые состояния и матрицы внимания, полученные из слоя трансформатора

3. выход логитов от слоя предсказания

Потери дистилляции трансформера

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

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

Отображение слоев

Представляя собой меньшую версию BERT, TinyBERT имеет меньшее количество слоев энкодера. Пусть количество слоев BERT равно N, а количество слоев TinyBERT равно M. Учитывая, что количество слоев различается, не очевидно, каким образом можно рассчитать функцию потери дистилляции.

Для этой цели вводится специальная функция n = g(m), которая определяет, какой слой BERT n используется для дистилляции его знаний в соответствующий слой m в TinyBERT. Выбранные слои BERT используются для расчета потерь во время обучения.

Введенная функция n = g(m) имеет два ограничения в рассуждении:

  • g(0) = 0. Это означает, что слой эмбеддингов в BERT отображается непосредственно на слой эмбеддингов в TinyBERT, что имеет смысл.
  • g(M + 1) = N + 1. Уравнение показывает, что слой предсказаний в BERT отображается на слой предсказаний в TinyBERT.

Для всех остальных слоев TinyBERT 1 ≤ m ≤ M необходимо отображение соответствующих значений функции n = g(m). Предположим, что такая функция уже определена. Настройки TinyBERT будут рассмотрены позже в этой статье.

Дистилляция трансформера

1. Дистилляция слоя эмбеддинга

Перед подачей входных данных на модель они сначала токенизируются, а затем отображаются на изученные эмбеддинги. Эти эмбеддинги затем используются в качестве первого слоя модели. Все возможные эмбеддинги можно представить в виде матрицы. Для сравнения разницы между эмбеддингами студента и преподавателя можно использовать стандартную метрику регрессии, применяемую к их соответствующим матрицам эмбеддингов E. Например, дистилляция трансформера использует среднеквадратичную ошибку (MSE) в качестве метрики регрессии.

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

Потери дистилляции слоя эмбеддинга

Так как пространства эмбеддингов студента и преподавателя различаются, матрица W также имеет важное значение для линейного преобразования пространства эмбеддинга студента в пространство эмбеддинга преподавателя.

2. Дистилляция слоя трансформера

Визуализация потерь дистилляции слоя трансформера

2A. Дистилляция слоя внимания

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

В TinyBERT рассматриваются все слои внимания, и результирующее значение потерь для каждого слоя равно сумме значений MSE между соответствующими матрицами весов внимания студента и преподавателя для всех голов.

Потери дистилляции слоя внимания

Матрицы внимания A, используемые для дистилляции слоя внимания, являются не нормализованными, в отличие от их softmax-выхода softmax(A). Согласно исследователям, эта тонкость приводит к более быстрой сходимости и улучшенной производительности.

2B. Дистилляция скрытого слоя

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

Потери дистилляции скрытого слоя

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

3. Дистилляция слоя предсказаний

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

Потери на уровне прогнозирования

Иногда логиты делятся на параметр температуры T, который контролирует гладкость распределения выходных данных. В TinyBERT температура T установлена на 1.

Уравнение потерь

В TinyBERT, в зависимости от типа, каждый слой имеет свою функцию потерь. Чтобы дать некоторым слоям большую или меньшую важность, соответствующие значения потерь умножаются на постоянную a. Окончательная функция потерь равняется взвешенной сумме значений потерь на всех слоях TinyBERT.

Функция потерь в TinyBERT

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

Обучение

Столь же важно отметить, что большинство моделей NLP (включая BERT) разрабатываются в два этапа:

  1. Модель предварительно обучается на большом корпусе данных, чтобы получить общую информацию о структуре языка.
  2. Модель настраивается на другом наборе данных для решения конкретной задачи.

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

  1. Генеральная дистилляция. TinyBERT получает общее представление о структуре языка от предварительно обученной модели BERT (без настройки). Используя меньше слоев и параметров, после этого этапа TinyBERT в целом показывает худшую производительность по сравнению с BERT.
  2. Задачно-специфичная дистилляция. В этот раз роль учителя выполняет настроенная версия BERT. Чтобы улучшить производительность, как предложили исследователи, на обучающей выборке применяется метод аугментации данных. Результаты показывают, что после задачно-специфичной дистилляции TinyBERT достигает сравнимой производительности с BERT.
Процесс обучения

Аугментация данных

Была разработана особая техника аугментации данных для задачно-специфичной дистилляции. Она заключается в замене процента слов в последовательности в двух случаях:

  • Если слово токенизировано так же, то это слово предсказывается моделью BERT, и предсказанное слово заменяет исходное слово в последовательности.
  • Если слово токенизировано на несколько подслов, то эти подслова заменяются на наиболее похожие векторы GloVe.

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

Пример аугментации

Настройки модели

С всего 14.5M параметрами, TinyBERT примерно в 7.5 раза меньше, чем BERT base. Их подробное сравнение показано на следующей диаграмме:

Сравнение BERT base и TinyBERT

При сопоставлении слоев авторы предлагают равномерную стратегию, согласно которой каждый слой TinyBERT отображается на каждый третий слой BERT: g(m) = 3 * m. Изучались и другие стратегии (например, использование всех самых нижних или верхних слоев BERT), но равномерная стратегия показала лучшие результаты, что кажется логичным, так как она позволяет передавать знания с разных слоев абстракции, делая передаваемую информацию более разнообразной.

Different layer mapping strategies. Performance results are shown for the GLUE dataset.

Говоря о процессе обучения, TinyBERT обучается на английской Википедии (2500M слов) и имеет большинство своих гиперпараметров такими же, как в базовой версии BERT.

Заключение

Дистилляция Transformer является большим шагом в обработке естественного языка. Учитывая, что модели на основе Transformer являются одними из самых мощных в настоящее время в машинном обучении, мы можем дальше развиваться, применяя дистилляцию Transformer, чтобы эффективно сжимать их. Одним из лучших примеров является TinyBERT, который сжат в 7,5 раза по сравнению с BERT базовой версии.

Несмотря на такое значительное сокращение параметров, эксперименты показывают, что TinyBERT демонстрирует сопоставимую производительность с BERT базовой версии: достигая 77,0% оценки в бенчмарке GLUE, TinyBERT незначительно уступает BERT с оценкой 79,5%. Явно, это удивительное достижение! Наконец, другие популярные техники сжатия, такие как квантизация или срезка, могут быть применены к TinyBERT, чтобы сделать его еще меньше.

Ресурсы

Все изображения, если не указано иное, принадлежат автору