Начало

Языковые модели, основанные на архитектуре трансформеров, такие как Llama, Mistral и прочие, показывают впечатляющие результаты на английском языке. Однако их эффективность на других языках, включая казахский, может страдать. Дообучение на отдельный домен, даже при наличии хорошего датасета, может не давать значительного прироста в качестве. И дело не столько в том, что базовая модель при обучении видела мало текста на казахском, сколько в неэффективной токенизации. Этот недостаток приводит к тому, что модели не могут в полной мере использовать свой потенциал на языках, отличных от английского.

В качестве примера возьмем Mistral (без паники, он использует тот же BPE токенизатор со словарем в 32к токенов), а так же Gemma (чей словарь адаптирован для работы с разными языками, и потому больше в 8 раз) и попробуем разбить текст на токены. К слову, токен - это минимальная значимая единица текста (может быть словом, частью слова или даже отдельной морфемой), которую обрабатывают трансформеры. В итоге вот что мы получаем:

В то время как тот же текст на английском занимал бы 10 токенов, на казахском он все еще в 2-3 раза больше. Отсюда вытекает не только слабая эффективность обучения, но и низкая скорость генерации ответа (учитывая то, что у одинаковых моделей токены предсказываются с одной и той же скоростью, на генерацию текста целиком нужно потратить больше времени той из них, у которой токенизатор разбивает слово на несколько токенов, вместо 1-2). Плюс не стоит забывать о контекстном окне, которое будет заполняться гораздо быстрее, если не исправить эту проблему. 

С одной стороны токенайзер у Gemma выглядит более привлекательнее, судя по примеру выше, а с другой, учитывая что при предсказании модель должна определить вероятность для каждого из 256к токенов из словаря (вместо 32к у Mistral’а) интерес к ней медленно угасает, ввиду лишнего усложнения. 

Токенизатор

Словарь токенизатора у Mistral содержит 32к токенов, преимущественно английских лексем (есть и на казахском, но их очень мало). Было решено расширять словарь не более чем в 2 раза. Для этого мы собрали около 20 Гб сырых текстов всяких новостей, статей и прочего (преимущественно на казахском и процентов 5-10 на русском), определили размер нового словаря (те же 32к) и поучили токенизатор. Результаты выглядели заметно лучше даже при таком маленьком датасете:

Итоговый словарь токенизатора (после объединения исходного и нашего нового, а также после удаления дубликатов) содержит чуть более 60к токенов. Далее нас ждет сложный, долгий и дорогой этап, а именно предварительное обучение.

Пре-трейн

Так как моделька с нашим новым токенизатором не знакома и такие размеры ей в новинку, нужно проинициализировать те слои, которые будут с ним “взаимодействовать”, а именно входной (embed_tokens, переводит токены со скриншотов выше в понятные модели векторные представления) и выходной (lm_head, переводит вектора в вероятность генерации того или иного токена). После этого новые слои будут плохо “дружить” с тем, что находится между ними (само “тело” модели-трансформера), и вместо осмысленного текста генерировать мусор из случайных токенов, поэтому следующим шагом мы будем их учить. 

Учитывая, что на этом этапе основная задача натренировать модель вразумительно продолжать текст, мы используем тот же общий датасет необработанных текстов, на котором мы учили токенизатор. Он маленький (всего на 2B токенов) и не самый чистенький, но в качестве эксперимента для построения опытной модели должно быть достаточно.

Из мощностей в распоряжении у нас были 2хH100, которые мы в течение недели нещадно эксплуатировали. Модель училась ровно, лоссы (и трейн, и эвал) падали (всем видом показывая, что 7-ми дней ей мало), сэмплы демонстрировали все лучшие и лучшие результаты. По окончании эвал лосс встал на отметке 1.75 (при этом на старте он был аж 32.46, так как с расширенным словарем токенизатора мы пытались заставить ЛЛМку писать текст на мало знакомом ей языке). Вот несколько удачных и неудачных примеров в задаче ответов на вопросы (красным выделил то, что генерировала наша новая модель):

Хоть результаты еще далеки от совершенства, но это уже значительно лучше того, как сейчас генерирует текст на казахском та же Llama-2-70b или прочие опенсорсные ЛЛМки. Подобные модели либо пытаются ответить на английском (на столько, на сколько они поняли вопрос), либо просто сыпят случайными словами на казахском: 

И вот мы пришли к тому, с чего начинали, однако теперь у нас есть модель, которая худо бедно работает с казахским языком на нашем новом токенизаторе. Иными словами, мы “прирастили” новый язык к модели. Теперь казахский переводится в такие вектора, которые для модели несут правильный и понятный смысл — как это раньше происходило для нативного ей английского языка. Дальнейшим правилом хорошего тона является провести файн-тюнинг.

Файн-тюнинг

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

Преимущественно, такие датасеты составлены на английском языке. Все остальное, чаще всего, является просто переводом через Deepl или Google Translate. Казахский язык тут не является исключением, поэтому мы насобирали то, что уже есть, а также дополнительно сами попереводили несколько небольших датасетов. В итоге наш датасет содержал переведенные Alpaca, Dolly, и много других мелки (включая датасеты для DPO, например belebele). После небольшой ручной чистки, у нас получилось около 200к сэмплов, для которых мы четко определили шаблон промпта и раскидали специальные токены (bos, eos, и т.д.) по его краям, чтобы модель училась понимать и то, когда ответ стоит заканчивать, а не продолжать генерацию бесконечно. 

На этом этапе (как впрочем и на предыдущем) использовали библиотеки HuggingFace, такие как transformers, peft, trl и datasets. Помимо них были еще deepspeed, bitsandbytes и flash-attention — эти отвечали за эффективность дообучения, и позволили получить результаты лучше в те же сроки. Предобученную модель квантизовали  до 4-х бит и тюнили LoRA для всех слоев (спасибо за это QLoRA, и в частности тому, что можно заморозить модель, понизить ее точность и потом на лету восстанавливать оригинальные веса через двойную деквантизацию). По окончании тренировки, а именно спустя примерно 12 часов, наша модель выглядела разительно лучше предобученной (и базового Mistral-7B, само собой). За это время эвал лосс упал с 3.39 до финальных 0.74, а модель научилась неплохо извлекать информацию из небольшого контекста:

К сожалению другие операции (например, суммаризация или перефразирование) работают хуже (в датасете было мало хороших сэмплов, поэтому пока она может извлекать только простую информацию — определенно есть задел на будущее! Мы уже начали готовить датасеты), но на вопросы без контекста отвечает достаточно развернуто и зачастую правильно. Вот несколько примеров:

В завершение

Итоговая модель продемонстрировала отличный потенциал для дальнейшего обучения (метрики будут представлены позже). Она хорошо отвечает на простые вопросы, способна работать с контекстом (пусть пока и слабо) и, благодаря более эффективному токенизатору, скорость генерации текста на казахском языке увеличилась примерно в 3-4 раз по сравнению с базовой моделью — а вместе с этим и уменьшилась стоимость генерации ответов. Во столько же увеличилась вместительность контекстного окна — теперь в модель можно подать на несколько страниц текста больше.

Наш опыт адаптации языковой модели для казахского языка показывает, что улучшение токенизации и предварительное обучение на целевом языке могут значительно повысить качество и скорость работы модели. Этот подход может быть применен и для других языков, которые в меньшей мере были представлены во всех популярных моделях из опенсорса.

Скачать 

• Предобученная модель

• ЛоРА

Разработку ведет команда Gen2B^{[ссылка удалена мод.]}

Авторы:
- Денис: @bomze
- Игорь: @seeyouall
- Армен: @armen