Soft-Di[M]O: Improving One-Step Discrete Image Generation with Soft Embeddings

2509.22925v1 cs.CV, cs.AI, cs.LG 2025-10-01
Авторы:

Yuanzhi Zhu, Xi Wang, Stéphane Lathuilière, Vicky Kalogeiton

Резюме на русском

#### Контекст Современные технологии генерации изображений столкнулись с вызовами в области эффективности и точности. Особенно актуальной является проблема создания одношаговых генераторов, которые могут эффективно генерировать изображения за один проход, сохраняя высокое качество и точность. Однако существующие подходы страдают от моделирования биаса, заложенного в учительских моделях, и от невозможности применения постобучения оптимизации, таких как гауссовский тренинг или тест-тайм эмбеддинг оптимизация. Эти ограничения способствуют снижению качества и регрессии в процессе развития. #### Метод Мы предлагаем способ решения этих проблем с помощью **soft embeddings** — метода, который заменяет дискретные токены на ожидаемые эмбеддинги под управлением распределения генератора. Эта техника позволяет сохранить высокую точность дискретного генератора, при этом делая его дифференцируемым и подходящим для последовательного улучшения. В рамках этого подхода интегрируется существующий фреймворк Di[M]O, чтобы сделать генераторы одношаговыми и функционально эффективными. Это решение устраняет барьеры для таких методов, как GAN-обучение, тест-тайм оптимизация и другие методы, которые требуют дифференцируемости. #### Результаты Мы провели эксперименты с широким спектром моделей-учителей, таких как MaskBit и MaskGen, используя датасеты ImageNet-256 и другие. Результаты показали, что Soft-Di[M]O (сочетание метода soft embeddings с Di[M]O) улучшает классо-изображения, получает фидбэк-фри коллектив за 1.56 в ImageNet-256 и повышает оценки GenEval и HPS при применении тестовых оптимизаций. Эти результаты доказывают высокую эффективность и гибкость метода, которая позволяет улучшить качество генерации изображений за один проход. #### Значимость Предложенный подход имеет широкое применение в области генеративных моделей, в том числе для тексто-изображения и других задач. Он обеспечивает значительные преимущества, такие как улучшение точности и качества изображений, возможность применения различных методов последовательного оптимизации, таких как GAN и тест-тайм эмбеддинг оптимизация. Это может привести к укреплению приложений в глубоком обучении, где эффективность и качество изображений являются ключевыми факторами. #### Выводы Мы представили Soft-Di[M]O — новый подход к одношаговой генерации изображений, который улучшает точность и эффективность дискретных генераторов. Наш метод доказывает свою эффективность на разных моделях-учителях и датасетах, определяя новые показатели качества. Мы планируем продолжить развитие этого подхода, исследуя

Abstract

One-step generators distilled from Masked Diffusion Models (MDMs) compress multiple sampling steps into a single forward pass, enabling efficient text and image synthesis. However, they suffer two key limitations: they inherit modeling bias from the teacher, and their discrete token outputs block gradient flow, preventing post-distillation refinements such as adversarial training, reward-based fine-tuning, and Test-Time Embedding Optimization (TTEO). In this work, we introduce soft embeddings, a simple relaxation that replaces discrete tokens with the expected embeddings under the generator's output distribution. Soft embeddings preserve representation fidelity for one-step discrete generator while providing a fully differentiable continuous surrogate that is compatible with teacher backbones and tokenizer decoders. Integrating soft embeddings into the Di[M]O distillation framework (denoted Soft-Di[M]O) makes one-step generators end-to-end trainable and enables straightforward application of GAN-based refinement, differentiable reward fine-tuning, and TTEO. Empirically, across multiple MDM teachers (e.g., MaskBit, MaskGen), Soft-Di[M]O achieves state-of-the-art one-step results: improved class-to-image performance, a one-step FID of 1.56 on ImageNet-256 with GAN-based refinement, along with higher GenEval and HPS scores on text-to-image with reward fine-tuning, and further gains from TTEO.

Ссылки и действия