Компания Snapchat представила компактный диффузионный трансформер, который позволяет создавать изображения по тексту прямо на смартфонах с качеством серверных моделей. Модель с 0,4 миллиарда параметров обходит соперников в 30 раз крупнее.
Snap Inc. разработала эффективный диффузионный трансформер для генерации изображений высокого разрешения непосредственно на смартфонах. Как указано в исследовательской работе, SnapGen++ формирует изображения размером 1024 × 1024 пикселя всего за 1,8 секунды на iPhone 16 Pro Max.
Ранние мобильные модели вроде SnapGen опирались на архитектуру U-Net, но SnapGen++ впервые переносит на смартфоны более мощную архитектуру диффузионных трансформеров. Именно она лежит в основе крупных серверных систем, таких как Flux и Stable Diffusion 3, однако раньше ее огромные вычислительные аппетиты делали использование на мобильных устройствах невозможным.
Диффузионные трансформеры стали прорывом в создании изображений. Они объединяют сильные стороны трансформеров — в частности, умение разбираться в сложных текстовых запросах и эффективно масштабироваться — с проверенным методом диффузии. В итоге получаются более цельные и детализированные картинки по сравнению с моделями на U-Net.
Новый подход к вниманию сильно экономит ресурсы
Главное препятствие для диффузионных трансформеров — колоссальные вычисления, которые растут квадратично с увеличением разрешения. Разработчики преодолели это с помощью свежего метода внимания, который резко снижает нагрузку.
Вместо обработки всего изображения целиком модель сочетает общее представление с точной проработкой локальных участков. Такой прием сокращает задержку на шаг вывода с 2000 миллисекунд до менее 300 миллисекунд, не теряя в качестве генерации.
Команда придумала Elastic Training. Одна сессия обучения выдает три варианта модели: Tiny с 0,3 миллиарда параметров для недорогих Android-смартфонов, Small с 0,4 миллиарда для топовых устройств и Full с 1,6 миллиарда для серверов или квантованного использования на устройствах.
Все варианты делят параметры и обучаются совместно. Авторы отмечают, что это позволяет подстраиваться под разное железо без отдельных тренировок.
Компактная модель обходит гигантов
SnapGen++ показывает отличные результаты в тестах. Вариант Small с 0,4 миллиарда параметров стабильно превосходит Flux.1-dev с 12 миллиардами параметров — в 30 раз больше — по качеству изображений и соответствию тексту. SD3.5-Large с 8,1 миллиарда параметров уступает даже самой большой модели Snap.
Для быстрого вывода на устройствах создали специальный метод дистилляции K-DMD. Он уменьшает шаги инференса с 28 до четырех при почти неизменном качестве. Общая задержка для Small-версии на iPhone 16 Pro Max — около 1,8 секунды при четырех шагах.
Соревнование за ИИ-генерацию на смартфонах
Snap раньше выпустил SnapGen — подход на U-Net, который давал изображения 1024 пикселя на смартфонах, но не дотягивал до серверных гигантов. Другие фирмы, например Google, тоже трудятся над компактными диффузионными моделями для мобильных. Однако по данным статьи, SnapGen++ — первая эффективная диффузионная трансформерная модель для высококачественной генерации прямо на устройстве.
Snap активно вкладывается в ИИ-функции для своего мессенджера. Помимо встроенного чат-бота My AI компания в ноябре 2024 года объявила о сделке на 400 миллионов долларов с Perplexity AI. Поисковик на базе ИИ интегрируют в Snapchat по умолчанию в этом году.