S-WITTI: Проектирование масштабируемых трансформеров для синтеза текстов в изображения

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

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

Масштабные AR модели

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

Преимущества масштабных моделей

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

Архитектура S-WITTI

Основная архитектура

Основная архитектура S-WITTI включает три основных компонента: RQ-VAE в качестве токенизатора изображений, предварительно обученный текстовый энкодер и блоки трансформеров с причинным вниманием. Мы используем предварительно обученный RQ-VAE, который представляет изображение в виде последовательности масштабов с различными разрешениями. Это позволяет модели эффективно генерировать изображения, начиная с низких разрешений и постепенно увеличивая их.

Модификации для улучшения производительности

В процессе обучения мы сталкиваемся с проблемами стабильности, которые приводят к не оптимальной сходимости модели. Для решения этих проблем мы внедряем несколько архитектурных модификаций, которые значительно улучшают стабильность и производительность модели. Например, мы заменяем активацию GELU на SwiGLU, что позволяет модели динамически контролировать поток информации через обучаемый механизм управления.

Анализ карт самовнимания

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

Влияние текстового кондиционирования

Мы также исследуем влияние текстового кондиционирования на различных масштабах модели. Высокие масштабы показывают минимальную зависимость от текстовой информации, что позволяет отключить классификаторно-свободное руководство (CFG) на последних масштабах. Это не только ускоряет процесс извлечения, но и иногда помогает смягчить артефакты генерации.

Обучение модели

Данные для предобучения

Мы собираем набор данных из 100 миллионов пар изображений и текстов, отобранных по эстетическим критериям. Это позволяет модели учиться на высококачественных изображениях с подробными текстовыми описаниями. Мы используем предварительно обученные текстовые энкодеры, такие как CLIP и OpenCLIP, для извлечения текстовых эмбеддингов, которые затем комбинируются и подаются в модель.

Процесс обучения

Процесс обучения включает два этапа: предобучение и тонкая настройка. На первом этапе мы обучаем модель на изображениях разрешения 256×256, а затем переходим к 512×512. Мы используем смешанную точность (BF16/FP32) и различные методы оптимизации, такие как Adam, для достижения стабильности и высокой производительности.

Оценка производительности

Автоматизированные метрики

Для оценки производительности S-WITTI мы используем комбинацию автоматизированных метрик, таких как CLIPScore, FID и другие. Эти метрики помогают количественно оценить качество сгенерированных изображений и их соответствие текстовым описаниям.

Оценка людьми

Мы также проводим исследование предпочтений среди людей, чтобы оценить качество генерации изображений по нескольким критериям: наличие дефектов, соответствие тексту, эстетика и сложность изображений. Результаты показывают, что S-WITTI превосходит другие AR модели по большинству критериев, конкурируя со сложными моделями диффузии.

Заключение

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

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