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'ами, понимающими русский язык. Скорость указана в расчёте на одно предложение из Лейпцигского веб-корпуса русского языка.
Модель | Скорость (CPU) | Скорость (GPU) | Вес на диске |
6 мс | 3 мс | 45 мб | |
125 мс | 8 мс | 680 мб | |
110 мс | 8 мс | 680 мб | |
120 мс | 8 мс | 1.8 гб | |
420 мс | 16 мс | 1.6 гб |
Все расчёты я выполнял на 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-токена оказывались лучше (ибо в целом у него скорее плохая репутация). Полный код моих экспериментов с оценкой можно посмотреть в блокноте, но он грязноват, так что лучше дождаться, пока я причешу этот бенчмарк.
Таблица с численными результатами
Задача \ Модель | rubert-tiny | bert-base-multilingual | rubert by DeepPavlov | LaBSE | sbert-large by Sber |
stsb | 0.6520 | 0.6216 | 0.7345 | 0.7728 | 0.6536 |
Paraphraser | 0.4345 | 0.4189 | 0.4204 | 0.4033 | 0.4194 |
XNLI | 0.3558 | 0.3481 | 0.3908 | 0.3355 | 0.3433 |
SentiRuEval2016 | 0.9026 | 0.9176 | 0.8260 | 0.9386 | 0.8656 |
Toxicity | 0.8610 | 0.8502 | 0.8939 | 0.9445 | 0.9687 |
Inappropriateness | 0.6834 | 0.6855 | 0.7455 | 0.7659 | 0.7911 |
Intents | 0.5848 | 0.5632 | 0.6054 | 0.7478 | 0.6786 |
Intets cross-lingual | 0.5394 | 0.2326 | 0.3638 | 0.7444 | 0.3658 |
FactRuEval | 0.4355 | 0.6829 | 0.6766 | 0.6963 | 0.6952 |
RuDReC | 0.5838 | 0.6520 | 0.6520 | 0.6705 | 0.6380 |
По итогам оценки оказалось, что модель 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]},
}