ML Underhood

Как устроена модель исправления ошибок в нейроредакторе Яндекс Браузера — часть II

Продолжаем говорить о модели исправления ошибок, которая работает «под капотом» нейроредактора в Яндекс Браузере. В прошлой части ML-разработчик Никита Авдосев рассказал о качестве исправления и работы с промптом, а сегодня речь пойдёт о перфомансе.

Для ускорения генерации в компании прибегли к методу спекулятивного декодирования. Суть его заключается в использовании компактной «черновой» (draft) модели, которая предлагает варианты продолжения цепочек токенов. Основная модель проверяет их и выбирает одну с помощью стохастического алгоритма выборки.

Существует несколько подходов к спекулятивному декодированию, а в Яндексе остановились на одном из самых популярных — EAGLE. Он предполагает дообучение небольших голов поверх основой модели. Гипотезы при этом генерируются в виде дерева, а не списка, благодаря чему повышается точность принятия токенов.

В качестве эксперимента инженеры Яндекса решили попробовать метод на модели исправления грамматических ошибок в русскоязычных текстах. Выбор был сделан не случайно — это одна из самых широко используемых моделей в Браузере, ежедневно к ней обращаются более 50 тысяч человек.

Чтобы обучить EAGLE предсказывать токены быстрее, чем их придумает оригинальная модель, требуется много текстов. В Яндексе использовали 250 тысяч текстов из логов пользователей. Позже попробовали обучить на более чем миллионе текстов, но, к сожалению, это не улучшило результат.

Затем начали тестировать новую конфигурацию для инференса: подбирать количество предсказываемых токенов перебором. Остановились на предсказании четырёх токенов. Это золотая середина, после которой ускорения не происходило, а местами даже увеличивалась задержка.

Благодаря EAGLE время генерации текста сократилось более чем в два раза. Теперь она в среднем занимает меньше секунды, что в контексте LLM — почти моментально.

Для ускорения моделей, которые работают с промптами пользователей, применяли FP8-квантизацию. Её отличительная особенность — квантизация не в целые, а в вещественные числа. Подход позволил добиться ускорения на 15% по сравнению с методом SmoothQuant, использованным ранее.

ML Underhood