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

Казалось бы, что между ними нет ничего общего, но это не так. Абсолютно все эти компоненты объединяет одно — языковая модель. Чем выше её качество, тем меньше ошибок будет допущено при вводе текста, а значит, пользователь будет чуточку счастливее.

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

Немного истории

Исторически языковая модель и большинство компонент (например, подсказки) нам достались из Яндекс Спеллера. Эта технология позволяет предлагать наиболее вероятные исправления для слов с опечатками и ошибками. Спеллер использует CatBoost и благодаря ему, а также большой языковой модели, он может расшифровывать искажённые до неузнаваемости слова и учитывать контекст при исправлении опечаток.

Но эту технологию нельзя было применить напрямую в Клавиатуре из-за ограничений, связанных с её спецификой:

• Размер моделей: модель должна занимать несколько десятков мегабайт, потому что ресурсы мобильного устройства ограничены.

• Задержка: мы должны успевать обрабатывать запросы в строго отведённое время. Подсказки должны появляться за 200 мс, а реакция на нажатие клавиши должна быть меньше 10 мс.

• Энергопотребление: нам нельзя сильно нагружать ЦПУ, часто обращаться к сети или тратить много времени на отрисовку.

• Универсальность: все предыдущие условия должны выполняться на любых смартфонах под всевозможными архитектурами.

Для Спеллера, который работает на наших серверах, ограничения куда более щадящие. В его случае модель в несколько десятков, а то и сотни гигабайт, выглядит совершенно уместно. Поэтому нам пришлось проделать большую работу по сжатию и ускорению моделей, при этом сохранив качество.

Со временем стало понятно, что мы начали потихоньку упираться в предел технологий, некогда породивших Клавиатуру. Спеллер основан на классическом подходе в машинном обучении, который включает в себя n-граммные языковые модели. N-граммные языковые модели строятся на основе статистик, собранных на больших корпусах текстов. Для этого все тексты нарезаются на n-граммы длиной от 1 до n, где n — максимальная длина сниппета или порядок модели. Дальше все одинаковые n-граммы объединяются и им выставляется общее количество. На основе такой большой таблички (n-грамма и её количество), можно аппроксимировать вероятность слова в данном контексте. Например, для триграммной модели вероятность слова в контексте приближается вот такой формулой:

Качество таких моделей растёт с увеличением размера обучающих корпусов и порядка самой модели. Но это также означает, что и сама модель очень сильно разрастается. Как бы нам ни хотелось, мы не можем позволить себе использовать более крупную модель в самой Клавиатуре. Частично эту проблему решают облачные подсказки, которые показали хорошие результаты в экспериментах.

Но на этом проблемы не исчерпываются, и появляется всё больше вопросов:

• Как передать дополнительную информацию, например, названия приложений: следует ли дополнительно хранить множество маленьких n-граммных моделей и смешивать их с основной?

• Как сделать персонализацию: следует ли держать дополнительную персональную n-граммную модель и так же подмешивать её к основной?

• Что делать с незнакомыми словами в контексте?

• Как обрабатывать более длинные контексты?

Ответ на большинство этих вопросов уже известен: необходимо использовать более выразительную языковую модель, а именно — нейронную.

Нейроязыковая модель 

Изучаем подходы

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

Для языковой модели Samsung используют обычную LSTM, а для эмбеддингов и финального слоя с softmax — полносвязанные матрицы. Чтобы уменьшить размер полученных моделей, они делают SVD-разложение, используют общие веса для эмбеддингов и для слоя softmax, а затем полученную модель дополнительно квантизуют. В результате размер модели уменьшается почти в 8 раз — до 7,4 Мб. Качество после таких сжатий незначительно проседает, но всё равно остаётся значительно выше предыдущей n-граммной модели. А генерация подсказок укладывается в 10 мс, что очень хорошо.

Статья Samsung ценна ещё тем, что она обозначает основные проблемные места в таких моделях: большой размер эмбеддингов и слоя softmax. Это наблюдение пригодится нам в дальнейшем.

В статье Google про Federated Learning есть очень краткое описание архитектуры: вместо обычной LSTM они используют CIFG, что позволяет сократить количество параметров на 25% и упростить обучение на устройствах. Так же, как и в Samsung, они шарят веса эмбеддингов и слоя softmax. В результате получается очень компактная модель размером всего 1,4 Мб. Но даже при таком маленьком размере, они с большим запасом выигрывают у n-граммной модели. 

Из этих статей становится ясно две вещи:

• В клавиатуре стоит использовать нейроязыковые модели: они значительно лучше по качеству n-граммных моделей и достаточно быстрые даже для мобильных устройств.

• Все основные оптимизации связаны с размером эмбеддингов и слоя softmax. Поэтому на них бы и обратим всё своё внимание.

Архитектура

Итоговая архитектура языковой модели в клавиатуре выглядит так:

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

Encoder

Здесь мы не стали выдумывать ничего хитрого и взяли стандартную LSTM, которая показала себя с лучшей стороны на фоне других вариантов RNN. Хотя её использование может показаться несколько устаревшим в 2023 году, но в данном случае это полностью оправдано. Пользователи вводят текст последовательно, и LSTM, в силу своей авторегрессионности, нам отлично подходит. 

Для каждого нового токена достаточно пересчитать только одно состояние LSTM. Это то же самое, что выполнить одну достаточно тяжёлую итерацию энкодера. Учитывая ограниченные вычислительные ресурсы мобильных устройств, это достаточно эффективный подход.                 

Embedding

Здесь мы рассмотрели несколько вариантов. 

Стандартный Word Embedding, где каждому слову из закрытого словаря ставится в соответствие свой вектор. Главный недостаток этого подхода — большой размер: размер эмбеддинга × количество слов. Кроме того, такой подход не умеет эффективно обрабатывать незнакомые слова, но его основное преимущество заключается в скорости работы. Именно такой эмбеддинг (с некоторыми оптимизациями) и используют Samsung и Google.

BPE (Byte Pair Encoding). Решает все проблемы стандартного подхода: большой размер и незнакомые слова. Такой подход позволяет эффективно обрабатывать незнакомые слова, так как любое слово может быть разбито на BPE-токены, и так как количество токенов значительно меньше количества слов, то и размер модели получается меньше. 

Однако у BPE есть свой недостаток — недетерминированность. Количество BPE-токенов может отличаться для разных слов, что приводит к разному количеству итераций для энкодера. Это сильно усложняет работу с кэшами и делает время работы крайне непредсказуемым.

CHAR-CNN (свёрточные эмбеддинги). Основная идея была взята из статьи Character-Aware Neural Language Models и звучит примерно так: «Давайте строить эмбеддинги на побуквенных свёртках». Этот подход позволяет модели учитывать информацию о структуре слова на уровне символов. 

Каждое слово разбивается на буквы, затем для каждой буквы считаются их эмбеддинги, после чего мы прогоняем эти побуквенные эмбеддинги через ряд свёрток с различными размерами фильтров. В такой постановке, фильтр размера 1 для свёрток представляет собой рассмотрение только побуквенных униграмм, а фильтр размера 2 — побуквенных биграмм. Полученные выходы объединяются для формирования финального эмбеддинга слова.

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

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

А это итоговая табличка по эмбеддингам. CNN-эмбеддинг — однозначный лидер, поэтому в дальнейшем будем использовать только его.

Softmax-слой

Помимо того, что Softmax-слой много весит, он и работает дольше всех компонент из-за большого размера словаря. Можно попробовать использовать BPE, но так же, как и в эмбеддингах, это портит свойство детерминированности: мы можем сгенерировать как один токен, так и пять. 

Оказывается, что у этой проблемы есть очень элегантное решение — Differentiated Softmax. Идея очень простая: самым частотным словам выделять проекции побольше, а менее частотным — поменьше. Это убивает сразу двух зайцев: размер и количество вычислений становятся заметно меньше. 

Влияние на другие компоненты клавиатуры

Так как мы решили заменить такую важную составляющую клавиатуры как языковую модель, то это должно повлиять и на остальные компоненты. Однако на деле это затронуло только динамическую сетку. Для автокорректа и классического свайпа всё осталось, как прежде, потому что здесь мы используем Сatboost-классификатор. Языковая модель здесь выступает только в роли фичи.

А с динамической сеткой дела обстоят уже поинтереснее. Раз уж мы отказались от n-граммной модели, которая хранится в побуквенном префиксном дереве в пользу пословной нейромодели, то теперь предсказывать следующую букву становится невозможным. Чтобы исправить эту ситуацию, мы обучили ещё один небольшой декодер для побуквенного предсказания слов, используя те же эмбеддинги и энкодер, что и в пословной модели. 

Производительность и качество

Таким образом, на двух простых и красивых идеях (использовать побуквенные свёртки для эмбеддингов и выдавать более частотным словам большие проекции в Softmax-слое) мы решили все наши проблемы. 

Итоговые размеры и скорости у нас получились примерно следующими:

• Embedding: 1 МБ — 2 мс;

• Encoder: 5 МБ — 5 мс;

• Softmax: 12 МБ — 20 мс. 

По итоговым замерам скорость работы новой языковой модели оказалась хоть и медленнее классической n-граммной на 30%, но она укладывалась в наш порог в 200 мс с большим запасом.

В клавиатурах есть две основные оффлайн-метрики:

• Saved keys — процент сэкономленных нажатий клавиш пользователем с учётом подсказок.

• Predicted words — процент верно угаданных слов по пустому префиксу вводимого слова. 

С использованием нейроязыковой модели показатель saved keys вырос с 43,1% до 47,5% (+4.3 п.п), а predicted words — с 16,2% до 19,0% (+2,8 п.п).

Это много или мало? Ответить на этот вопрос можно так: наша большая облачная модель, которая превосходит по размеру нейромодель более чем в 200 раз, выдаёт качество хуже. Однозначно это хорошо, только не для облачной модели: от неё нам пришлось отказаться. Возможно, когда-нибудь мы снова вернёмся к ней.  

По онлайн-метрикам тоже всё хорошо:

• Доля символов из саджеста выросла на 7% (основная метрика подсказок).

• Количество отмен автокорректа упало на 2% (основная метрика автокорректа).

• Количество полезных символов на нажатие увеличилось на 1% (интегральная метрика хорошести ввода).

Итог очевиден: мы смогли добиться заметного прироста качества Клавиатуры при помощи новой нейроязыковой модели, попутно победив все сложности и ограничения мобильных устройств.

App-Specific подсказки

Как было сказано в самом начале, проблем у классической языковой модели было много:

• Незнакомые слова в контексте: например, опечатки при вводе предыдущих слов. Для обычной n-граммной модели опечатка в контексте это почти тоже самое, что и разрыв контекста.

• Длинные контексты: n-граммные модели ограничены порядком и им сложно держать в памяти начало длинной фразы.

• Сложность использования дополнительной информации, которая отличается от текстовой. Например, идентификатор приложения, где вводится текст, может очень сильно влиять на предсказания.

И если мы, внедрив новую нейромодель, решили первые две проблемы за счёт CNN-эмбеддингов и энкодера, то почему бы сразу не решить и третью проблему?

Начнём с того, что люди пишут совершенно разные тексты в разных приложениях. И исходя из этого хочется уметь показывать разные подсказки. Например, подсказка «привет» совершенно уместна по префиксу «п» в мессенджерах, но в приложениях, связанных с картами или контактами, мы хотим увидеть «проспект» или «Паша». 

Оказывается, что с новой моделью это сделать очень просто: достаточно подмешивать дополнительный эмбеддинг для приложений. Вот и всё! 

В итоге мы получили приличный прирост по метрике saved keys в разных группах приложениях:

• food: 37% → 52% (+15 п.п),

• contacts: 22% → 28% (+6 п.п),

• geo: 38% → 45% (+7 п.п),

• market: 40% → 48% (+8 п.п).

 А по онлайн-замерам мы получаем дополнительные 2% символов из саджеста. 

Подведём итог

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

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