Суть инновации: Простота вместо глубины
Лаборатория Ostris представила Primal VAE — вариационный автоэнкодер, который отказывается от традиционного глубокого сжатия (8x) в пользу минималистичного подхода. Модель выполняет только 2x пространственное уменьшение разрешения, сокращая каналы до 1 или 2. Для достижения стандартного для современных диффузионных моделей (как Flux) формата 8x сжатия используется встроенный в DiT (Diffusion Transformer) механизм patchify (патчинг 4x4).
Ключевое преимущество — отсутствие необходимости в масштабировании (scaling) и сдвиге (shift) латента. Латентное пространство обучено так, чтобы его норма соответствовала стандартному нормальному распределению шума. Это означает, что после кодирования encode() можно сразу передавать данные в диффузионную модель без промежуточных преобразований, упрощая пайплайн.
Архитектура и метрики
Обе версии модели (2x1 и 2x2) имеют одинаковую структуру: зеркальные сверточные U-Net с 4 уровнями (ширины 96/192/384/384) и общим объемом параметров около 80 миллионов. В отличие от классических VAE, здесь нет постериорного распределения, KL-члена или квантования. Нормализация осуществляется через RMSNorm по каналам, что позволяет кодировать изображения тайлами без артефактов на границах.
| Характеристика | Primal VAE (2x1) | Primal VAE (2x2) | Flux.1 VAE | Flux.2 VAE |
|---|---|---|---|---|
| Пространственное сжатие | 2x | 2x | 8x (через патчинг) | 8x (через патчинг) |
| Каналы латента | 1 | 2 | 16 | 32 |
| Итоговый формат (после patchify) | 16 ch @ 8x | 32 ch @ 8x | 16 ch @ 8x | 32 ch @ 8x |
| Параметры модели | ~80M | ~80M | Информация недостаточна | Информация недостаточна |
| Нормализация латента | Unit Normal (mean 0, std 1) | Unit Normal (mean 0, std 1) | Требует shift/scale | Требует batch-norm |
Результаты реконструкции (PSNR, dB)
Сравнение качества восстановления изображений (512px) показывает, что Primal VAE 2x2 конкурирует с Flux.2, в то время как 2x1 уступает на текстурированных изображениях из-за меньшего объема данных (1 канал против 16/32). Однако 2x1 требует значительно меньше шагов обучения (62k против 147k для 2x2).
| Изображение | Primal 2x1 | Flux.1 VAE | Primal 2x2 | Flux.2 VAE |
|---|---|---|---|---|
| Люди, студия | 32.1 | 35.7 | 37.5 | 37.2 |
| 22.0 | 24.2 | 27.1 | 26.5 | |
| Скриншот UI | 21.3 | 24.7 | 27.4 | 25.9 |
Технические детали и использование
Модель обучалась в ai-toolkit с использованием специфичных техник:
- Moment Loss: Принудительное приведение среднего к 0 и дисперсии к 1 для каждого изображения отдельно.
- Gaussian Order-Statistic Loss: Выравнивание формы распределения значений латента с нормальным распределением.
- Noise-perturbed decode: Добавление шума к латенту при декодировании для сглаживания пространства кодов.
Для использования в diffusers требуется загрузка с trust_remote_code=True и указанием варианта (variant="2x1" или "2x2"). Размеры изображений должны быть кратны 16.
Почему это важно
Primal VAE демонстрирует, что для совместимости с DiT-архитектурами не всегда требуется тяжелый энкодер с глубоким сжатием. Перенос части работы по сжатию на сам диффузионный процесс (через патчинг) позволяет создать легковесный, быстрый в кодировании VAE, который легко интегрируется в существующие пайплайны Flux без потери совместимости латентных пространств.
Источник: Huggingface ↗