Как стать автором
Обновить

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

Семантика *Программирование *Data Mining *Машинное обучение *Natural Language Processing *
✏️ Технотекст 2021

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)

Вес на диске

cointegrated/rubert-tiny

6 мс

3 мс

45 мб

bert-base-multilingual-cased

125 мс

8 мс

680 мб

DeepPavlov/rubert-base-cased-sentence

110 мс

8 мс

680 мб

sentence-transformers/LaBSE

120 мс

8 мс

1.8 гб

sberbank-ai/sbert_large_nlu_ru

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]},
}
Теги:
Хабы:
Всего голосов 57: ↑57 и ↓0 +57
Просмотры 21K
Комментарии Комментарии 17

Работа

Data Scientist
123 вакансии