Маленький и быстрый 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, скачать и запустить дистиллированный BERT просто. Например, вот так вы можете получить 312-мерный CLS-эмбеддинг предложения.

# pip install transformers sentencepiece
import torch
from transformers import AutoTokenizer, AutoModel
tokenizer = AutoTokenizer.from_pretrained("cointegrated/rubert-tiny")
model = AutoModel.from_pretrained("cointegrated/rubert-tiny")
# model.cuda()  # uncomment it if you have a GPU

def embed_bert_cls(text, model, tokenizer):
    t = tokenizer(text, padding=True, truncation=True, return_tensors='pt')
    with torch.no_grad():
        model_output = model(**{k: v.to(model.device) for k, v in t.items()})
    embeddings = model_output.last_hidden_state[:, 0, :]
    embeddings = torch.nn.functional.normalize(embeddings)
    return embeddings[0].cpu().numpy()

print(embed_bert_cls('привет мир', model, tokenizer).shape)
# (312,)

Этот эмбеддинг вы можете использовать как признак для любой модели классификации или кластеризации текстов. Как (рекомендуемый) вариант, вы можете дообучить мою модель на собственной задаче. Про это будет отдельный пост, а пока отсылаю вас к обучающим материалам от 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%.

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

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

Заключение

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

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

P.S. Если вдруг вы хотите упомянуть этот пост или модель в научной публикации, оформить цитирование можно так:

@misc{dale_tiny_and_fast_bert_2021, 
   author = "Dale, David",
   title  = "Маленький и быстрый BERT для русского языка", 
   editor = "habr.com", 
   url    = "https://habr.com/ru/post/562064/", 
   month  = {June},
   year   = {2021},   
   note = {[Online; posted 10-June-2021]},
}