MobileLLM: Optimizing Sub-billion Parameter Language Models for On-Device Use Cases

Юзать LLM хотят все, но инференс в облаке стоит конских денег. Но есть простое решение - давайте юзеры будут запускать LLM у себя, а не в облаке. Однако не стоит забывать – большинство пользователей интернета владеют лишь телефоном. Оперативной памяти там не так уж и много - у iPhone это 6 гигов, у большинства андроид флагманов до 12, так что нужно оптимизировать параметры, но и не забывать про перформанс. Авторам пейпера удалось заметно улучшить качество инференса на телефонах, без потерь в перформансе.

Как это вышло?

➖ Авторы выяснили, что масштабирование модели вглубь при таком же количестве параметров работает лучше чем при масштабировании модели вширь. Таким образом модель на 125 миллионов параметров имеет 30 слоёв (у GPT-2, Bert и вообще всего в такой весовой категории обычно 12).
➖ Традиционный для "эффективных" моделей шеринг входного и выходного слоя эмбеддингов.
➖ Завезли Group Query Attention, который раньше в таких моделях не использовался.
➖ Убедились что все слои влезают в кэш телефонов, потому что оперативка значительно медленнее.

Самая интересная часть - layer sharing, так что её разберём поподробнее. Идея в том, чтобы использовать один и тот же слой несколько раз, и таким образом улучшить перформанс. Попробовали несколько стратегий:

➖ Повторять слои незамедлительно - то есть слой компьютит что-то, а потом его вывод кормится в него же.
➖ Повторять все слои в том же порядке - по сути моделька запускается два раза на одном и том же инпуте.
➖ Сначала считать слои в обычном порядке, а потом задом наперёд.

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

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

Пейпер
Код

@ai_newz