Источник: arXiv (препринт)
arXiv ID: 1910.01108v4 [cs.CL]
Оригинальное название: DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter
Ссылка на препринт: https://arxiv.org/abs/1910.01108
Лицензия оригинала: CC BY-SA 4.0 (Creative Commons Attribution-ShareAlike 4.0 International).
Данный HTML-файл представляет собой неофициальный перевод статьи на русский язык, выполненный для учебных/научно-образовательных целей. Перевод не является публикацией авторов оригинала.

DISTILBERT: ДИСТИЛЛИРОВАННАЯ ВЕРСИЯ BERT — МЕНЬШЕ, БЫСТРЕЕ, ДЕШЕВЛЕ И ЛЕГЧЕ

Первоисточник: Sanh V., Debut L., Chaumond J., Wolf T. DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter. arXiv:1910.01108v4. https://arxiv.org/abs/1910.01108 .

Victor Sanh, Lysandre Debut, Julien Chaumond, Thomas Wolf
Hugging Face
EMC²: 5th Edition Co-located with NeurIPS’19

АННОТАЦИЯ

По мере того как перенос обучения (Transfer Learning) на основе крупномасштабных предобученных моделей становится всё более распространённым в обработке естественного языка (NLP), эксплуатация таких больших моделей «на краю» (on-the-edge) и/или при ограниченных вычислительных бюджетах обучения или инференса остаётся сложной задачей. В данной работе мы предлагаем метод предобучения меньшей по размеру универсальной модели представлений языка — DistilBERT, — которую затем можно дообучать (fine-tune) и получать хорошие результаты на широком спектре задач, как и у более крупных аналогов. В то время как большинство предыдущих работ изучало дистилляцию для построения моделей, ориентированных на конкретные задачи, мы используем дистилляцию знаний на этапе предобучения и показываем, что можно уменьшить размер модели BERT на 40%, сохранив 97% её возможностей понимания языка и получив ускорение инференса на 60%. Чтобы задействовать индуктивные смещения, усваиваемые большими моделями при предобучении, мы вводим «тройную» функцию потерь, объединяющую потери языкового моделирования, дистилляции и косинусного расстояния. Наша более компактная, быстрая и лёгкая модель дешевле в предобучении; также мы демонстрируем её пригодность для вычислений на устройстве (on-device) в proof-of-concept эксперименте и в сравнительном on-device исследовании.

СОДЕРЖАНИЕ

  1. 1. Введение
  2. 2. Дистилляция знаний
  3. 3. DistilBERT: дистиллированная версия BERT
  4. 4. Эксперименты
  5. 5. Связанные работы
  6. 6. Заключение и дальнейшая работа
  7. Литература
  8. Сноски

1. Введение

Последние два года ознаменовались ростом подходов Transfer Learning в обработке естественного языка (NLP), в рамках которых крупномасштабные предобученные языковые модели стали базовым инструментом для многих NLP-задач [Devlin et al., 2018, Radford et al., 2019, Liu et al., 2019]. Хотя эти модели приводят к значительному улучшению, они зачастую имеют несколько сотен миллионов параметров, а текущие исследования по предобученным моделям1 показывают, что обучение ещё более крупных моделей по-прежнему приводит к лучшим результатам на последующих (downstream) задачах.

Тренд к увеличению моделей вызывает ряд опасений. Во-первых, это экологическая стоимость экспоненциального масштабирования вычислительных требований таких моделей, о чём говорится в [Schwartz et al., 2019], [Strubell et al., 2019]. Во-вторых, хотя запуск этих моделей на устройстве в реальном времени способен открыть новые и интересные приложения для обработки языка, растущие вычислительные требования и требования к памяти могут препятствовать их широкому распространению.

Число параметров нескольких недавно представленных предобученных языковых моделей
Рисунок 1. Число параметров нескольких недавно представленных предобученных языковых моделей.

В этой статье мы показываем, что на многих последующих задачах можно достичь сопоставимого качества, используя гораздо меньшие языковые модели, предобученные с применением дистилляции знаний; в результате получаются модели, которые легче и быстрее при инференсе, а также требуют меньшего вычислительного бюджета при обучении. Наши универсальные предобученные модели можно дообучать с хорошими результатами на широком спектре последующих задач, сохраняя гибкость более крупных моделей. Мы также показываем, что наши сжатые модели достаточно компактны, чтобы работать «на краю», например на мобильных устройствах.

Используя тройную функцию потерь, мы показываем, что Transformer, уменьшенный на 40% [Vaswani et al., 2017], предобученный посредством дистилляции под надзором более крупной Transformer-языковой модели, может достигать сопоставимого качества на различных последующих задачах, при этом будучи на 60% быстрее при инференсе. Дополнительные абляционные исследования показывают, что для наилучших результатов важны все компоненты тройной функции потерь.

Мы сделали доступными обученные веса вместе с кодом обучения в библиотеке Transformers2 от HuggingFace [Wolf et al., 2019].

2. Дистилляция знаний

Дистилляция знаний [Bucila et al., 2006; Hinton et al., 2015] — это техника сжатия, при которой компактная модель (студент) обучается воспроизводить поведение более крупной модели (учителя) или ансамбля моделей.

В обучении с учителем модель-классификатор обычно обучают предсказывать класс примера, максимизируя оценённую вероятность «истинных» меток. Стандартная целевая функция обучения, таким образом, включает минимизацию кросс-энтропии между предсказанным распределением модели и one-hot эмпирическим распределением обучающих меток. Модель, хорошо работающая на обучающем наборе, будет предсказывать распределение с высокой вероятностью на корректном классе и с близкими к нулю вероятностями на других классах. Однако некоторые из этих «почти нулевых» вероятностей больше других и отражают, в частности, обобщающую способность модели и то, насколько хорошо она будет работать на тестовом наборе. 3

Функция потерь обучения. Студент обучается с использованием дистилляционной функции потерь по «мягким» целевым вероятностям учителя: \(L_{ce} = \sum_i t_i \cdot \log(s_i)\), где \(t_i\) (соответственно, \(s_i\)) — оценка вероятности, выдаваемая учителем (соответственно, студентом). Такая цель даёт богатый обучающий сигнал, используя полное распределение учителя. Следуя Hinton et al. [2015], мы используем softmax с температурой: \(p_i = \frac{\exp(z_i/T)}{\sum_j \exp(z_j/T)}\), где \(T\) управляет «плавностью» выходного распределения, а \(z_i\) — оценка (логит) модели для класса \(i\). Одна и та же температура \(T\) применяется к студенту и учителю во время обучения, тогда как при инференсе \(T\) устанавливается равной 1, чтобы восстановить стандартный softmax.

$$L_{ce}=\sum_i t_i \cdot \log(s_i)$$
$$p_i=\frac{\exp(z_i/T)}{\sum_j \exp(z_j/T)}$$

Итоговая целевая функция обучения — линейная комбинация дистилляционной потери \(L_{ce}\) и потери обучения с учителем, в нашем случае — потери masked language modeling \(L_{mlm}\) [Devlin et al., 2018]. Мы обнаружили полезным добавить косинусную «эмбеддинговую» потерю (\(L_{cos}\)), которая стремится выровнять направления векторов скрытых состояний студента и учителя.

3. DistilBERT: дистиллированная версия BERT

Архитектура студента. В настоящей работе студент — DistilBERT — имеет ту же общую архитектуру, что и BERT. При этом удаляются эмбеддинги типа сегмента (token-type embeddings) и pooler, а число слоёв уменьшается в 2 раза. Большинство операций, используемых в архитектуре Transformer (линейные слои и нормализация слоя), хорошо оптимизированы в современных фреймворках линейной алгебры, и наши исследования показали, что вариации по последней размерности тензора (размерность скрытого слоя) оказывают меньший эффект на вычислительную эффективность (при фиксированном бюджете параметров), чем вариации по другим факторам, таким как число слоёв. Поэтому мы фокусируемся на уменьшении количества слоёв.

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

Таблица 1 — DistilBERT сохраняет 97% качества BERT. Сравнение на dev-наборах бенчмарка GLUE. Результаты ELMo приведены по данным авторов. Результаты BERT и DistilBERT — медианы 5 запусков с разными seed.

Модель Score CoLA MNLI MRPC QNLI QQP RTE SST-2 STS-B WNLI
ELMo 68.7 44.1 68.6 76.6 71.1 86.2 53.4 91.5 70.4 56.3
BERT-base 79.5 56.3 86.7 88.6 91.8 89.6 69.3 92.7 89.0 53.5
DistilBERT 77.0 51.3 82.2 87.5 89.2 88.5 59.9 91.3 86.9 56.3

Таблица 2 — DistilBERT демонстрирует сопоставимую эффективность на прикладных задачах. Сравнение на downstream задачах: IMDb (точность на тесте) и SQuAD 1.1 (EM/F1 на dev-наборе). D: второй шаг дистилляции на этапе дообучения.

Модель IMDb (acc.) SQuAD (EM/F1)
BERT-base 93.46 81.2/88.5
DistilBERT 92.82 77.7/85.8
DistilBERT (D) - 79.1/86.9

Таблица 3 — DistilBERT значительно меньше, оставаясь при этом стабильно быстрее. Время инференса полного прохода по задаче GLUE STS-B (анализ тональности) на CPU при размере батча 1.

Модель # param. (Millions) Inf. time (seconds)
ELMo 180 895
BERT-base 110 668
DistilBERT 66 410

Дистилляция. Мы применили лучшие практики обучения BERT, недавно предложенные в Liu et al. [2019]. Таким образом, DistilBERT дистиллируется на очень больших батчах с использованием накопления градиента (до 4K примеров на батч), с динамическим маскированием и без цели next sentence prediction.

Данные и вычислительные ресурсы. Мы обучаем DistilBERT на том же корпусе, что и оригинальную модель BERT: на конкатенации English Wikipedia и Toronto Book Corpus [Zhu et al., 2015]. DistilBERT обучался на 8 GPU V100 (16 GB) примерно 90 часов. Для сравнения, модель RoBERTa [Liu et al., 2019] требовала 1 день обучения на 1024 GPU V100 (32 GB).

4. Эксперименты

Общее понимание языка. Мы оцениваем способности DistilBERT к пониманию языка и обобщению на бенчмарке General Language Understanding Evaluation (GLUE) [Wang et al., 2018], представляющем собой коллекцию из 9 наборов данных для оценки систем понимания естественного языка. Мы приводим результаты на dev-наборах для каждой задачи, дообучая (fine-tuning) DistilBERT без использования ансамблирования или многозадачной схемы дообучения (эти приёмы в основном ортогональны рассматриваемой работе). Мы сравниваем результаты с базовой линией, предложенной авторами GLUE: энкодер ELMo [Peters et al., 2018], за которым следуют два BiLSTM. 4

Результаты по каждой из 9 задач приведены в таблице 1 вместе с макро-оценкой (средним значением отдельных метрик). Среди 9 задач DistilBERT всегда не уступает базовой линии ELMo или превосходит её (вплоть до 19 пунктов точности на STS-B). DistilBERT также неожиданно близок к BERT, сохраняя 97% качества при 40% меньшем числе параметров.

4.1. Бенчмарк прикладных задач

Прикладные задачи. Мы дополнительно исследуем качество DistilBERT на нескольких прикладных задачах при ограничениях, связанных с эффективным инференсом: задача классификации (анализ тональности IMDb — Maas et al. [2011]) и задача вопрос-ответ (SQuAD v1.1 — Rajpurkar et al. [2016]).

Как показано в таблице 2, DistilBERT отстаёт от BERT всего на 0.6 процентного пункта по точности на тесте IMDb, будучи при этом на 40% меньше. На SQuAD DistilBERT находится в пределах 3.9 пункта от полной версии BERT. Мы также изучили, можно ли добавить ещё один шаг дистилляции на этапе адаптации, дообучая DistilBERT на SQuAD с использованием BERT-модели, предварительно дообученной на SQuAD, в роли учителя как дополнительного слагаемого в функции потерь (knowledge distillation). В этом режиме получаются два последовательных шага дистилляции: один на этапе предобучения и один на этапе адаптации. В таком случае нам удалось добиться интересных результатов с учётом размера модели: 79.8 по F1 и 70.4 по EM, то есть в пределах 3 пунктов от полной модели.

Размер и скорость инференса. Чтобы подробнее изучить компромисс «ускорение/размер» для DistilBERT, мы сравниваем (в таблице 3) число параметров каждой модели и время инференса, необходимое для полного прохода по dev-набору STS-B на CPU (Intel Xeon E5-2690 v3 Haswell @2.9GHz) при размере батча 1. DistilBERT имеет на 40% меньше параметров, чем BERT, и на 60% быстрее BERT.

Вычисления на устройстве. Мы изучили, можно ли использовать DistilBERT для edge-приложений, создав мобильное приложение для вопрос-ответа. Мы сравниваем среднее время инференса на современном смартфоне (iPhone 7 Plus) с ранее обученной нами моделью вопрос-ответа на базе BERT-base. Если исключить шаг токенизации, DistilBERT на 71% быстрее BERT, а весь набор весов модели занимает 207 MB (что можно дополнительно уменьшить с помощью квантизации). Наш код доступен. 5

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

В этом разделе мы исследуем влияние различных компонентов тройной функции потерь и инициализации студента на качество дистиллированной модели. Мы приводим макро-оценку на GLUE. В таблице 4 представлены приращения относительно полной конфигурации с тройной потерей: удаление потери masked language modeling оказывает небольшое влияние, тогда как две потери дистилляции объясняют значительную часть качества.

Таблица 4 — Абляционное исследование. Изменения приведены относительно модели, обученной с тройной потерей и инициализацией весов студента весами учителя.

Абляция Изменение макро-оценки GLUE
∅ - Lcos - Lmlm -2.96
Lce - ∅ - Lmlm -1.46
Lce - Lcos - ∅ -0.31
Тройная потеря + случайная инициализация весов -3.69

5. Связанные работы

Дистилляция, специфичная для задачи. Большинство предыдущих работ сосредоточено на построении схем дистилляции, ориентированных на конкретную задачу. Tang et al. [2019] переносят знания, дообучая классификационную модель BERT в LSTM-классификатор. Chatterjee [2019] дистиллирует модель BERT, дообученную на SQuAD, в более компактную трансформерную модель, предварительно инициализированную из BERT. В настоящей работе мы обнаружили, что полезнее использовать дистилляцию для модели общего назначения на этапе предобучения, а не дистилляцию, завязанную на конкретную задачу. Turc et al. [2019] используют исходную цель предобучения, чтобы обучить меньшую student-модель, а затем дообучают её с применением дистилляции. Как показано в абляционном исследовании, мы также выяснили, что полезно использовать знания учителя при предобучении, добавляя дополнительный дистилляционный сигнал.

Мульти-дистилляция. Yang et al. [2019] объединяют знания ансамбля учителей, применяя multi-task learning для регуляризации дистилляции. Авторы используют Multi-Task Knowledge Distillation, чтобы обучить компактную модель вопрос-ответа на основе набора больших моделей вопрос-ответа. Одним из применений мульти-дистилляции является многоязычность: Tsai et al. [2019] применяют схожий подход, предобучая многоязычную модель с нуля исключительно с помощью дистилляции. Однако, как показано в абляционном исследовании, использование знаний учителя вместе с инициализацией и дополнительными потерями приводит к существенному выигрышу.

Другие методы компрессии. Для сжатия больших моделей также изучались иные техники. Недавние исследования по прореживанию весов (weight pruning) показывают, что можно удалять некоторые attention-головы на этапе инференса без существенного ухудшения качества [Michel et al., 2019]. В некоторых слоях число attention-голов можно сократить до одной. Отдельная линия исследований использует квантизацию для получения более компактных моделей [Gupta et al., 2015]. Прореживание и квантизация являются ортогональными по отношению к настоящей работе.

6. Заключение и дальнейшая работа

Мы представили DistilBERT — универсальную предобученную версию BERT, на 40% меньшую и на 60% более быструю, при этом сохраняющую 97% возможностей понимания языка. Мы показали, что универсальная языковая модель может быть успешно обучена с использованием дистилляции, и проанализировали различные компоненты с помощью абляционного исследования. Далее мы продемонстрировали, что DistilBERT является привлекательным вариантом для edge-приложений.

References

  1. Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In NAACL-HLT, 2018.
  2. Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019.
  3. Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar S. Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke S. Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. ArXiv, abs/1907.11692, 2019.
  4. Roy Schwartz, Jesse Dodge, Noah A. Smith, and Oren Etzioni. Green ai. ArXiv, abs/1907.10597, 2019.
  5. Emma Strubell, Ananya Ganesh, and Andrew McCallum. Energy and policy considerations for deep learning in nlp. In ACL, 2019.
  6. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In NIPS, 2017.
  7. Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, and Jamie Brew. Transformers: State-of-the-art natural language processing, 2019.
  8. Cristian Bucila, Rich Caruana, and Alexandru Niculescu-Mizil. Model compression. In KDD, 2006.
  9. Geoffrey E. Hinton, Oriol Vinyals, and Jeffrey Dean. Distilling the knowledge in a neural network. ArXiv, abs/1503.02531, 2015.
  10. Yukun Zhu, Ryan Kiros, Richard S. Zemel, Ruslan Salakhutdinov, Raquel Urtasun, Antonio Torralba, and Sanja Fidler. Aligning books and movies: Towards story-like visual explanations by watching movies and reading books. 2015 IEEE International Conference on Computer Vision (ICCV), pages 19–27, 2015.
  11. Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R. Bowman. Glue: A multi-task benchmark and analysis platform for natural language understanding. In ICLR, 2018.
  12. Matthew E. Peters, Mark Neumann, Mohit Iyyer, Matt Gardner, Christopher Clark, Kenton Lee, and Luke Zettlemoyer. Deep contextualized word representations. In NAACL, 2018.
  13. Alex Wang, Jan F. Tenny, Yada Pruksachatkun, Katheryn Yu, Jun Hula, Patrick Xia, Raghu Pappagari, Shuning Jin, R. Thomas McCoy, Roma Patel, Yinqui Huang, Jason Phang, Edouard Grave, Najoung Kim, Phu Mon Htut, Thibault F’evry, Berlin Chen, Nikita Nangia, Hakun Liu, Ahmad Mohananey, Shikha Bordia, Nicolas Patry, Ellie Pavlick, and Samuel R. Bowman. jiant 1.1: A software toolkit for research on general-purpose text understanding models. http://jiant.info/, 2019.
  14. Andrew L. Maas, Raymond E. Daly, Peter T. Pham, Dan Huang, Andrew Y. Ng, and Christopher Potts. Learning word vectors for sentiment analysis. In ACL, 2011.
  15. Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. Squad: 100, 000+ questions for machine comprehension of text. In EMNLP, 2016.
  16. Raphael Tang, Yao Lu, Linqing Liu, Lili Mou, Olga Vechtomova, and Jimmy Lin. Distilling task-specific knowledge from bert into simple neural networks. ArXiv, abs/1903.12136, 2019.
  17. Debajyoti Chatterjee. Making neural machine reading comprehension faster. ArXiv, abs/1904.00796, 2019.
  18. Iulia Turc, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Well-read students learn better: The impact of student initialization on knowledge distillation. ArXiv, abs/1908.08962, 2019.
  19. Ze Yang, Linjun Shou, Ming Gong, Wutao Lin, and Daxin Jiang. Model compression with multi-task knowledge distillation for web-scale question answering system. ArXiv, abs/1904.09636, 2019.
  20. Henry Tsai, Jason Riesa, Melvin Johnson, Naveen Arivazhagan, Xin Li, and Amelia Archer. Small and practical bert models for sequence labeling. In EMNLP-IJCNLP, 2019.
  21. Paul Michel, Omer Levy, and Graham Neubig. Are sixteen heads really better than one? In NeurIPS, 2019.
  22. Suyog Gupta, Ankur Agrawal, Kailash Gopalakrishnan, and Pritish Narayanan. Deep learning with limited numerical precision. In ICML, 2015.