Маленький и быстрый BERT для русского языка

BERT нейросеть, способная весьма неплохо понимать смысл текстов на человеческом языке. Впервые появивишись в 2018 году, эта модель совершила переворот в компьютерной лингвистике. Базовая версия модели долго предобучается, читая миллионы текстов и постепенно осваивая язык, а потом её можно дообучить на собственной прикладной задаче, например, классификации комментариев или выделении в тексте имён, названий и адресов. Стандартная версия BERT довольно большая: весит больше 600 мегабайт, обрабатывает предложение около 120 миллисекунд (на CPU). В этом посте я предлагаю уменьшенную версию BERT для русского языка 45 мегабайт, 6 мс на предложение. Уже есть tinybert для английского от Хуавея, есть моя уменьшалка FastText'а, а вот маленький (англо-)русский BERT, кажется, появился впервые. Но насколько он хорош?

Дистилляция путь к маленькости

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

Если очень коротко, то BERT работает так: сначала токенизатор разбивает текст на токены (кусочки размером от одной буквы до целого слова), от них берутся эмбеддинги из таблицы, и эти эмбеддинги несколько раз обновляются, используя механизм self-attention для учёта контекста (соседних токенов). При предобучении классический BERT выполняет две задачи: угадывает, какие токены в предложении были заменены на специальный токен [MASK], и шли ли два предложения следом друг за другом в тексте. Как потом показали, вторая задача не очень нужна. Но токен [CLS], который ставится перед началом текста и эмбеддинг которого использовался для этой второй задаче, употреблять продолжают, и я тоже сделал на него ставку.

Дистилляция способ перекладывания знаний из одной модели в другую. Это быстрее, чем учить модель только на текстах. Например, в тексте [CLS] Ехал Грека [MASK] реку "верное" решение поставить на место маски токен через, но большая модель знает, что токены на, в, за в этом контексте тоже уместны, и это знание полезно для обучения маленькой модели. Его можно передать, заставляя маленькую модель не только предсказывать высокую вероятность правильного токена через, а воспроизводить всё вероятностное распределение возможных замаскированных токенов в данном тексте.

В качестве основы для модели я взял классический bert-multilingual (веса), ибо хочу, чтобы модель понимала и русский, и английский, и его же использую на ранних стадиях дистилляции как учителя по распределению токенов. Словарь этой модели содержит 120К токенов, но я отобрал только те, которые часто встречаются в русском и английском языках, оставив 30К. Размер эмбеддинга я сократил с 768 до 312, число слоёв с 12 до 3. Эмбеддинги я инициализировал из bert-multilingual, все остальные веса случайным образом.

Поскольку я собираюсь использовать маленький BERT в первую очередь для классификации коротких текстов, мне надо, чтобы он мог построить хорошее векторное представление предложения. Поэтому в качестве учителей для дистилляции я выбрал модели, которые с этим здорово справляются: RuBERT (статья, веса), LaBSE (статья, веса), Laser (статья, пакет) и USE (статья, код). А именно, я требую, чтобы [CLS] эмбеддинг моей модели позволял предсказать эмбеддинги предложений, полученные из этих трёх моделей. Дополнительно я обучаю модель на задачу translation ranking (как LaBSE). Наконец, я решил, что неплохо было бы уметь полностью расшифровывать предложение назад из CLS-эмбеддингов, причём делать это одинаково для русских и английских предложений как в Laser. Для этих целей я примотал изолентой к своей модели декодер от уменьшенного русского T5. Таким образом, у меня получилась многозадачная модель о восьми лоссах:

• Обычное предсказание замаскированных токенов (я использую full word masks).

• Translation ranking по рецепту LaBSE: эмбеддинг фразы на русском должен быть ближе к эмбеддингу её перевода на английский, чем к эмбеддингу остальных примеров в батче. Пробовал добавлять наивные hard negatives, но заметной пользы они не дали.

• Дистилляция распределения всех токенов из bert-base-multilingual-cased (через несколько эпох я отключил её, т.к. она начала мешать).

• Приближение CLS-эмбеддингов (после линейной проекции) к эмбеддингам DeepPavlov/rubert-base-cased-sentence (усреднённым по токенам).

• Приближение CLS-эмбеддингов (после другой линейной проекции) к CLS-эмбеддингам LaBSE.

• Приближение CLS-эмбеддингов (после третьей проекции) к эмбеддингам LASER.

• Приближение CLS-эмбеддингов (после ещё одной проекции) к эмбеддингам USE.

• Расшифровка декодером от T5 предложения (на русском) из последней проекции CLS-эмбеддинга.

Скорее всего, из этих лоссов больше половины можно было безболезненно выкинуть, но ресурсов на ablation study я пока не нашёл. Обучал я это всё в течении нескольких дней на Colab, по пути нащупывая learning rate и другие параметры. В общем, не очень научно, но дешево и результативно. В качестве обучающей выборки я взял три параллельных корпуса англо-русских предложений: от Яндекс.Переводчика, OPUS-100 и Tatoeba, суммарно 2.5 млн коротких текстов. Весь процесс создания модели, включая некоторые неудачные эксперименты, содержится в блокноте. Сама модель, названная мною rubert-tiny (или просто Энкодечка), выложена в репозитории Huggingface.

И как этим пользоваться?

Если у вас есть Python и установлены пакет transformers и sentencepiece, скачать и запустить модель просто. Например, вот так вы можете получить 312-мерный CLS-эмбеддинг предложения.

Этот эмбеддинг вы можете использовать как признак для любой модели классификации или кластеризации текстов. Как (рекомендуемый) вариант, вы можете дообучить мою модель на собственной задаче. Про это будет отдельный пост, а пока отсылаю вас к обучающи материалам от Huggingface (раз, два).

Насколько быстр и мал мой Энкодечка? Я сравил его с другими BERT'ами, понимающими русский язык. Скорость указана в расчёте на одно предложение из Лейпцигского веб-корпуса русского языка.

Все расчёты я выполнял на Colab (Intel(R) Xeon(R) CPU @ 2.00GHz и Tesla P100-PCIE c батчом размера 1 (если использовать крупные батчи, то ускорение на GPU ещё заметнее, т.к. с маленькой моделью можно собрать более большой батч).

Как видим, rubert-tiny на CPU работает раз в 20 быстрее своих ближайших соседей, и легко помещается на бюджетные хостинги типа Heroku (и даже, наверное, на мобильные устройства). Надеюсь, эта модель сделает предобученные нейросети для русского языка в целом более доступными для прикладных применений. Но надо ещё убедиться, что модель хоть чему-то научилась.

Оценка качества эмбеддингов

Рекомендованный и проверенный временем рецепт использования BERT дообучать его на конечную задачу. Но дообучение процесс небыстрый и наукоёмкий, а гипотезу об осмысленности выученных эмбеддингов хочется проверить побыстрее и попроще. Поэтому я не дообучаю модели, а использую их как готовые feature extractors, и обучаю поверх их эмбеддингов простые модели логистическую регрессию и KNN.

Для оценки русских моделей есть бенчмарк RussianSuperGLUE, но задачи там сложные, требующие подбора параметров и дообучения моделей, так что к нему я тоже перейду позже. Ещё есть бенчмарк RuSentEval, там задачи более простые, но они слишком абстрактные, ориентированные больше на лингвистические свойства, а хочется начать с чего-то прикладного. Поэтому я собрал свой маленький бенчмарк. Пока что он тоже живёт в блокноте, но надо будет допилить его и выложить в более удобном виде. И вот какие задачи туда вошли:

STS: бенчмарк по семантической близости предложений (переведённый с английского). Например, фразы "Кошка спит на фиолетовой простыне" и "Черно-белый кот спит на фиолетовом одеяле" из этого датасета оценены на 4 из 5 баллов сходства. Качество моделей на нём я мерял ранговой корреляций этих баллов с косинусной близостью эмбеддингов предложений. Для наилучшей модели, LaBSE, корреляция оказалась 77%, для моей 65%, на одном уровне с моделью от Сбера, которая в 40 раз больше моей.

Paraphraser: похожий бенчмарк по семантической близости, но с чисто русскими новостными заголовками и трёхбалльной шкалой. Эту задачу я решил так: обучил поверх косинусных близостей логистическую регрессию, и ей предсказывал, к какому из трёх классов пара заголовков относится. На этом бенчмарке моя модель выдала точность 43% и победила все остальные (впрочем, с небольшим отрывом).

XNLI: предсказание, следует ли одно предложение из другого. Например, из "это стреляющее пластиковое автоматическое оружие" не следует "это более надежно, чем металлическое оружие", но и не противоречит. Оценивал её я как предыдущую логрегом поверх косинусных близостей. Тут на первом месте оказалась модель от DeepPavlov (которая дообучалась ровно на этой задаче), а моя заняла второе.

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

OKMLCup: детекция токсичных комментариев из Одноклассников. Тут моя модель заняла четвёртое место по ROC AUC, обогнав только bert-base-cased-multilingual.

Inappropriateness: детекция сообщений, неприятных для собеседника или вредящих репутации. Тут моя модель оказалась на последнем месте, но таки набрала 68% AUC (у самой лучшей, Сберовской, вышло 79%).

Классификация интентов: накраудсоршенные обращения к голосовому помощнику, покрывающие 18 доменов и 68 интентов. Они собирались на английском языке, но я перевёл их на русский простой моделькой. Часть переводов получились странными, но для бенчмарка сойдёт. Оценивал я по точности логистической регрессии или KNN (что лучше). LaBSE набрала точность 75%, модель от Сбера 68%, от DeepPavlov 60%, моя 58%, мультиязычная 56%. Есть, куда расти.

Перенос классификации интентов с английского: та же классификация интентов, но обучающая выбора на английском языке, а тестовая на русском. Тут моя модель заняла второе место после LaBSE, а качество всех остальных сильно просело (что ожидаемо, ибо на параллельных корпусах они не обучались).

factRuEval-2016: задача распознавания классических именованных сущностей (адреса, организации, личности). Я обучал логистическую регрессию поверх эмбеддингов токенов, а качество мерял макро F1 скором (относительно токенов же, что не вполне корректно). Оказалось, что на таком NER моя модель работает откровенно плохо: она набрала скор 43%, остальные 67-69%.

RuDReC: распознавание медицинских именованных сущностей. Тут моя модель тоже проиграла остальным, но с меньшим отрывом: 58% против 62-67%.

Для всех задач, кроме NER, я пробовал для каждой модели два вида эмбеддингов: из CLS токена и усреднённые по всем токенам. Для меня было неожиданностью, что на многих задачах и моделях эмбеддинг СLS-токена оказывался даже лучше, чем средний (в целом, у него так себе репутация). Поэтому рекомендую не лениться и пробовать оба варианта. Код и результаты моей оценки можно взять из этого блокнота, но лучше подождать, пока я ещё причешу этот бенчмарк.

По итогам оценки оказалось, что модель LaBSE очень крутая: она заняла первое место на 6 из 10 задач. Поэтому я решил выложить LaBSE-en-ru, у которой я отрезал эмбеддинги всех 99 языков, кроме русского и английского. Модель похудела с 1.8 до 0.5 гигабайт, и, надеюсь, таким образом стала чуть более удобной для практического применения. Ну а rubert-tiny оказался по качеству в целом близок к моделям от DeepPavlov и Sber, будучи при этом на порядок меньше и быстрее.

Заключение

Я обещал сделать компактную модель для эмбеддингов русских предложений, и я это наконец сделал. Процесс дистилляции, скорее всего, я настроил неоптимально, и его ещё можно сильно улучшать, но уже сейчас маленькая модель на некоторых задачах приближается к уровню своих учителей и даже иногда обходит его. Так что если вам нужен маленький и быстрый BERT для русского языка, то пользуйтесь: https://huggingface.co/cointegrated/rubert-tiny.

Впереди работы много: с одной стороны, хочется обучить маленький BERT решать задачи из RussianSuperGLUE (и не только), с другой затащить в русский язык хорошие небольшие модели для контролируемой генерации текста (я уже начал делать это для T5). Посему лайкайте данный пост, подписывайтесь на мой канал про NLP, подкидывайте в комментариях и в личке интересные задачи, и, если у вас доведутся руки попробовать rubert-tiny, то обязательно оставляйте обратную связь!
Мне и самому интересно, что будет дальше.