Привет, Хабр! Если вы интересуетесь NLP или просто современными DL моделями, то приглашаю вас узнать, как можно, имея всего лишь одну A100, около 30 гигабайтов текста и несколько дней обучения, решить проблему ограниченного окна контекста для русскоязычных трансформеров. А ещё сделаем несколько оптимизаций и добьёмся почти лучших метрик в бенчмарке encodechka.
Получившиеся русскоязычные модели доступны в открытом доступе на HuggingFace 🤗:
Размер контекста | Classic BERTs | Sentence BERTs |
512+ | ||
2048+ |
Откуда взялась идея и мотивация
Методы позиционного кодирования сильно продвинулись за последние пару лет. Сейчас практически каждая появляющаяся LLM имеет контекст в 8k токенов и больше, к тому же появилось большое количество методов для расширения контекста, некоторые из них даже позволяют не дообучать предобученную модель. Только вот, к сожалению, все эти улучшения касаются в основном causal моделей, т. е. декодеров. Кажется, даже для англоязычных моделей, исследование SOTA методов позиционного кодирования в энкодерах застыло где-то на этапе DeBERTa, с её disentangled attention.
Однако ещё в середине 2021 года выходила модель RoFormer, которая использовала Rotary Position Embedding, именно этот метод сейчас применяется во всех LLM (Mistral, Llama, GPT-NeoX и тд). К сожалению, RoFormer не был адаптирован для русского языка, но потенциал, особенно сейчас, он имеет значительный. Большинство моделей на русском имеют стандартный размер контекста в 512 токенов, исключениями являются cointegrated/rubert-tiny2 от Дэвида Дале и kazzand/ru-longformer-tiny-16384 от MTS. Поэтому целью данного проекта и стала адаптация обычных моделей для возможности работы с RoPE, увеличение контекстного окна и улучшение их производительности. Как бонус — получившиеся в результате модели, по качеству превосходят существующие модели длинного контекста (и не только) для русского языка.
Описание метода обучения
Как же адаптировать модели с выученным позиционным кодированием для работы с RoPE? Очевидно, что нам нужно просто заменить одно на другое в архитектуре модели. Для этого мы можем удалить nn.Embedding слой, отвечающий за старое позиционное кодирование, а далее модифицировать SelfAttention блок в трансформере так, чтобы position_ids передавались на каждый слой и там же применялся RoPE к query и key проекциям, ведь именно так он и устроен.
Применив такое преобразование, мы сломаем модель почти полностью. Но это не мешает нам восстановить её способности, так как большинство весов мы не изменяли. Для этого, как простой вариант, нам потребуется просто предобучить модель в режиме RoBERTы — то есть учить MLM задачу, применив поверх модели соответствующую голову.
Но также, мы хотим именно восстановить способности модели, а не просто обучить её на новом датасете — это значит, дистиллировать оригинальную модель-учитель, в нашего RoPE-ученика. Мне кажется, сейчас уместнее это называть клонированием, так как размеры и архитектура остаются одинаковыми. Для того чтобы это осуществить, достаточно к MLM-лоссу добавить дистилляционный лосс, который может быть, по большому счёту, любым, кроме KL-дивергенции, так как MLM-головы никто не кладёт с весами модели, и распределения учителя на токены у нас не будет. Я решил использовать 2 лосса, по отдельности: Cosine Similarity усредненый по токенам, для каждого слоя, и InfoNCE, в симметричном режиме, только для последнего слоя, именно он используется в CLIP для сближения модальностей текста и картинок. В обоих лоссах также подбиралась и своя температура для увеличения важности близости.
Данные
В качестве данных я взял ~32 гигабайта из первых 10 русскоязычных фолдов датасета CulturaX. В семплах это составило около 5 миллионов. Данный датасет хорош тем, что довольно хорошо подготовлен и отфильтрован, в русских семплах действительно мало других языков. А ещё он имеет хорошее распределение длин, тут найдутся примеры для обучения модели с практически любым контекстным окном до 16к.
ruRoPEBert-classic-base
Эта первая модель проекта является клоном ai-forever/ruBert-base. Для её клонирования использовался послойный cosine similarity лосс с температурой 0.5 между скрытыми представлениями ученика и учителя.
ruRoPEBert-e5-base
Эта же модель является клоном hivaze/ru-e5-base, а та в свою очередь русифицированной версией intfloat/multilingual-e5-base, полученой с помощью удаления большинства лишних токенов из других языков, по методу Дэвида Дале, снижая таким образом размер словаря с 250k до 69k, что существенно облегчило её тренировку. Кстати, модель, в связи с изначальной мультиязычностью, сохраняет свойства и на английском языке, но он не являлся основным.
Из-за своей особенности представления похожести текстов, e5 клонировалась с симметричным InfoNCE лоссом (как в CLIP) поверх последнего слоя и температурой 0.1. Выбор этого лосса обоснован тем, что модель должна гораздо лучше понимать разницу между разными по смыслу текстами, а не просто сближать косинусную близость со своим оригиналом. Кстати, этот же лосс использовался авторами и при обучении оригинальной e5.
Метод расширения контекста
Я не буду подробно описывать как работает RoPE, а также его формулы и почему работают методы его интерполяции/экстраполяции. Для этого я предлагаю изучить статью Scaling Laws of RoPE-based Extrapolation или пост на хабре О методах позиционного кодирования в Transformer, а ещё оригинальную статью Roformer.
Итак, использовать клонирование оригинальной модели, очевидно, можно только на длинах до 512 токенов. Но нам нужно больше, поэтому вторым этапом в пайплайне подготовки моделей стало расширение контекста — дообучение моделей только с MLM лоссом и только на текстах с длинами 512-2048 токенов.
В ходе экспериментов, я пришёл к выводу, что тут лучше не использовать какие-либо техники интерполяции или экстраполяции RoPE эмбедингов при обучении. Дело в том, что модель и так довольно легко сходилась на увеличенном контексте к качеству, как на 512 токенах, ещё где-то на ~200к семплах. Техники, где мы как-то влияем на эмбединги, конечно, помогают, но они нужны скорее для значительного увеличения контекста, например с 4к до 16к, так как там потребовалось бы больше времени и данных, если ничего не использовать.
Тут вы могли заметить: если мы будем дообучать модель на длинный контекст без клонирующего лосса, то как мы убедимся, что наша модель будет всё ещё похожа на оригинальную? На самом деле, очень просто — нам достаточно обучать только attention блоки, не трогая последующий MLP блок. RoPE затрагивает именно query и key, соответственно, только их проекции нам и нужно дообучать, MLP будет мапить получившиеся вектора примерно в то же место в пространстве, что и раньше, а значит характеристики вроде косинусной близости между скрытыми представлениями с оригиналом, почти не изменятся.
Для дообучения мы не используем никакие трюки со скейлингом RoPE. Но вот для инференса вполне можем, есть как минимум метод Dynamic-NTK, который позволяет применять его без дообучения. Это означает, что потенциально, модели могут использоваться и на более длинном контексте чем 2k и не очень сильно терять в качестве.
Ниже приведена таблица с различными оценками модели ruRoPEBert-e5-base-2k, которая была дополнительно дообучена на 512-2048 токенах, на разных длинах контекста с Dynamic-NTK скейлингом и без него.
orig 0-512 | 512-2k | 2k-3k (no scale) | 2k-3k (with scale) | 2k-4k (no scale) | 2k-4k (with scale) | |
mlm_loss | 1.2157 | 1.2788 | 2.217 | 2.1488 | 2.6534 | 2.5895 |
mlm_accuracy | 0.7476 | 0.7358 | 0.6166 | 0.6253 | 0.5633 | 0.5746 |
Как видно, скейлинг помогает адаптироваться на длины больше, чем модель видела при обучении, но тем не менее, качество всё ещё ощутимо деградирует, критичность такой деградации будет зависеть от конкретных задач. Мне кажется, для улучшения результатов, тут стоит попробовать другие мощные методы, вроде xPos или Alibi, подробнее о них можно узнать в статье A Length-Extrapolatable Transformer.
Замена механизма вычисления внимания
Так как ручное вычисление внимания (eager) является абсолютно неоптимизированным, квадратичным по вычислительной сложности и слабым местом по памяти, его нельзя просто использовать как он есть, если мы хотим увеличить контекст модели в несколько раз.
Существует много вариантов на что можно заменить механизм вычисления внимания с ручного, например, модули из xformers или FlashAttention v1/v2 и тд. В качестве простого варианта, я обратил внимание на нативный метод, который активно продвигает PyTorch — torch.nn.functional.scaled_dot_product_attention (sdpa). Этот метод поддерживает сразу 3 способа оптимизации аттеншена, которые можно включать и отключать по собственному желанию. Среди них: Memory-Efficient Attention, FlashAttention, Fused C++ имплементация от PyTorch.
В недавнем релизе 2.2.0 авторы PyTorch добавили поддержку FlashAttention второй версии в эту функцию, тем самым она стала, наверное, самым удобным способом добавления оптимизации вычисления внимания в модель, так как использование нативного Flash не очень удобное.
В таблице ниже я привожу бенчмарки инференса, произведённые с помощью Pytorch Profiler 12-ти слойного ruRoPEBert в конфигурации base. Замеры производились на ноутбучной RTX 3070 Ti с процессором Intel Core i7 12700H. Размер батча во всех тестах = 1.
Attention type | Torch version | Context length | CPU-bound time | CUDA-bound time | Total run time | Self CUDA Mem |
eager | 2.1.0 | 512 | 18.774ms | 8.613ms | 27.387ms | -674.01 Mb |
sdpa | 2.1.0 | 512 | 15.920ms | 7.411ms | 23.331ms | -298.31 Mb |
eager | 2.1.0 | 2048 | 20.388ms | 57.924ms | 78.312ms | -7.39 Gb |
sdpa | 2.1.0 | 2048 | 19.234ms | 38.830ms | 58.064ms | -1.19 Gb |
eager | 2.2.0 | 2048 | 15.270ms | 57.735ms | 73.005ms | -7.53 Gb |
sdpa | 2.2.0 | 2048 | 13.500ms | 32.869ms | 46.369ms | -1.18 Gb |
Эффективность использования памяти при выборе sdpa вырастает в несколько раз, также модель получает ускорение относительно eager, пропорциональное увеличению размера контекста. При обучении также получается достичь ускорения в 2-3 раза и кратно сэкономить память.
Использовать модель, выбирая соответствующий способ вычисления внимания можно следующим образом:
model_name = 'Tochka-AI/ruRoPEBert-e5-base-2k'
model = AutoModel.from_pretrained(model_name,
trust_remote_code=True,
attn_implementation='sdpa') # или 'eager'Считаем метрики и сравниваем
За основу бенчмарка качества был выбран encodechka от Дэвида Дале, так как он обладает единым и простым интерфейсом запуска тестов, в нём есть немало современных моделей, а главное, что оценка производится исключительно на русскоязычных датасетах. К тому же он даёт возможность сравнить модели по их способности в семантическое представление предложений, что является важной целью данного проекта.
ruRoPEBert-e5-base и лучшие sentence-модели
Model name | STS | PI | NLI | SA | TI | IA | IC | ICX | NE1 | NE2 | Avg S (no NE) | Avg S+W (with NE) |
intfloat/multilingual-e5-large | 0.86 | 0.73 | 0.47 | 0.81 | 0.98 | 0.8 | 0.82 | 0.77 | 0.24 | 0.37 | 0.78 | 0.686 |
ruRoPEBert-e5-base-512 | 0.793 | 0.704 | 0.457 | 0.803 | 0.970 | 0.788 | 0.802 | 0.749 | 0.328 | 0.396 | 0.758 | 0.679 |
ruRoPEBert-e5-base-2k | 0.787 | 0.708 | 0.460 | 0.804 | 0.970 | 0.792 | 0.803 | 0.749 | 0.402 | 0.423 | 0.759 | 0.689 |
intfloat/multilingual-e5-base | 0.834 | 0.704 | 0.458 | 0.795 | 0.964 | 0.782 | 0.803 | 0.740 | 0.234 | 0.373 | 0.76 | 0.668 |
text-embedding-ada-002 | 0.78 | 0.66 | 0.44 | 0.77 | 0.96 | 0.77 | 0.75 | 0.73 | - | - | 0.734 | - |
cointegrated/rubert-tiny2 | 0.75 | 0.65 | 0.42 | 0.74 | 0.94 | 0.75 | 0.76 | 0.64 | 0.36 | 0.39 | 0.704 | 0.638 |
sberbank-ai/sbert_large_mt_nlu_ru | 0.78 | 0.65 | 0.4 | 0.8 | 0.98 | 0.8 | 0.76 | 0.45 | 0.3 | 0.34 | 0.703 | 0.626 |
mts longformer | 0.760 | 0.647 | 0.416 | 0.739 | 0.938 | 0.749 | 0.755 | 0.626 | 0.396 | 0.387 | 0.703 | 0.641 |
Как видно, ruRoPEBert-e5-base-2k является лучшей моделью на всем бенчмарке encodechka по скору S+W (sentence + word embeddings). Однако несколько отстаёт от моделей e5 по STS, возможно, это связано с тем, что во время обучения следовало взять температуру ниже, чем 0.1 или использовать дополнительные датасеты. Тем не менее во всех остальных тестах ruRoPEBert оказалась лучше, чем оригинал e5-base, с этим и связано общее первенство. Также обе модели обходят text-embedding-ada-002 от OpenAI.
ruRoPEBert-classic-base и лучшие BERT-энкодеры
Model name | STS | PI | NLI | SA | TI | IA | IC | ICX | NE1 | NE2 | Avg S (no NE) | Avg S+W (with NE) |
ruRoPEBert-classic-base-512 | 0.695 | 0.605 | 0.396 | 0.794 | 0.975 | 0.797 | 0.769 | 0.386 | 0.410 | 0.609 | 0.677 | 0.630 |
ruRoPEBert-classic-base-2k | 0.684 | 0.601 | 0.396 | 0.777 | 0.974 | 0.794 | 0.769 | 0.381 | 0.609 | 0.470 | 0.672 | 0.631 |
sberbank-ai/ruRoberta-large | 0.7 | 0.6 | 0.35 | 0.78 | 0.98 | 0.8 | 0.78 | 0.32 | 0.3 | 0.46 | 0.666 | 0.609 |
ai-forever/ruBert-base | 0.670 | 0.533 | 0.391 | 0.773 | 0.975 | 0.783 | 0.765 | 0.384 | - | - | 0.659 | - |
Получившиеся RoPEBert модели во всех тестах превосходят свой оригинал - ai-forever/ruBert-base, и даже более крупную ai-forever/ruRoberta-large. Судя по скорам NE1 и NE2, classic модели обладают большим потенциалом в решении NER задач, чем sentence клоны e5.
Идеи по дальнейшим улучшениям и выводы
Не всё получилось, как было задумано, например, в бенчмарке STS клонированная e5 модель проигрывает оригиналу, а также модели, на мой взгляд, не очень хорошо обобщаются на длины большие, чем видели при обучении.
Эти проблемы можно решить усовершенствовав пайплайн тренировки или некоторые гиперпараметры. Также можно заменить механизм экстраполяции, например, на xPos, Alibi или совсем недавний LongRoPE, а можно сделать спарсификацию аттеншена, как это предлагается в статье Blockwise Self-Attention for Long Document Understanding, ещё от 2019 года. Кроме того, в самой статье RoFormer упоминается, что авторы попробовали объединить RoPE с архитектурой Performer и работать посимвольно и линейно по сложности.
Удивляет, что существует огромное количество архитектур (от старого Reformer до новейших Biderectional SSM), которые могли бы эффективно использоваться уже долгое время в энкодерах, но стандартный вариант с квадратичным вниманием — всё ещё остаётся единственным путём, которым идёт большинство SOTA моделей, особенно на русском языке. В конечном итоге этот проект делает совсем небольшое усилие в сторону современных подходов в NLP и доказывает их относительную простоту реализации.
Веса всех моделей, их код и инструкции по запуску доступны в нашем аккаунте на HuggingFace (ссылки в начале статьи). В дальнейшем, у нас есть планы по выпуску large версий моделей и дотренировке с контекстом до 8k токенов, а также продолжение экспериментов с оптимизациями.