GenFirst: сквозное обучение генеративных моделей без коллапса латентного пространства

GenFirst: сквозное обучение генеративных моделей без коллапса латентного пространства

Латентные генеративные модели изображений обычно обучают в два независимых этапа: сначала вариационный автоэнкодер (VAE) учат сжимать изображения в латентное пространство и точно восстанавливать их обратно, а затем поверх уже готового, замороженного латентного пространства отдельно обучают генеративную модель. Проблема в том, что латенты, оптимизированные под точное восстановление, не обязательно удобны для генерации. Совместное сквозное обучение обеих частей сразу выглядит логичной альтернативой, но на практике почти всегда приводит к «коллапсу латентного пространства» (вырождению) и к конфликту между целями генерации и восстановления, раньше прямое сквозное обучение без коллапса не удавалось.

Авторы разобрали, как разные функции потерь формируют латентное пространство, и сформулировали два ключевых наблюдения. Во-первых, энтропийный член в дивергенции Кульбака-Лейблера критичен для предотвращения коллапса: восстановление и подгонка под априорное распределение стремятся сжать апостериорное распределение, тогда как энтропия сохраняет неустранённую (невырожденную) неопределённость латентов. Во-вторых, восстановление и генерация обучаются асимметрично, восстановление происходит быстро и находится под сильным сигналом обучения, тогда как генерация медленнее и хуже поддаётся оптимизации.

На основе этих наблюдений авторы предлагают GenFirst, простую стратегию «сначала генерация, потом восстановление». Вместо того чтобы с самого начала одинаково давить на обе цели, генеративная цель первой формирует латентное пространство при слабом давлении со стороны восстановления, и лишь затем давление восстановления постепенно усиливают, чтобы вернуть визуальные детали. По утверждению авторов, это позволило впервые провести прямое сквозное обучение без коллапса латентного пространства. Стратегию проверили на двух типах генеративных приоров: непрерывных авторегрессионных с точным правдоподобием и на приорах SiT с неявным правдоподобием. При сквозном обучении с GenFirst модель SiT на ImageNet-256 показывает gFID 0,97 с classifier-free guidance (CFG) и 1,45 без CFG, а модель MMDiT достигает оценки GenEval 0,90 на генерации изображений по тексту.

Помимо генерации изображений, авторы распространили подход на общие визуальные латенты для одновременного обучения генерации и распознавания образов (representation learning), а также на непрерывную объединённую генерацию текста и изображений. По их словам, это показывает, что предложенный способ стабильного сквозного обучения латентного пространства работает независимо от конкретного типа генеративного приора и модальности.

Ключевые факты

  • GenFirst, стратегия сквозного обучения латентных генеративных моделей, в которой генеративная цель формирует латентное пространство раньше цели восстановления, а не одновременно с ней.
  • Авторы утверждают, что впервые добились прямого сквозного (end-to-end) обучения без коллапса латентного пространства, типичной проблемы при совместном обучении VAE и генератора.
  • Две находки лежат в основе метода: энтропийный член в KL-дивергенции не даёт латентам вырождаться, а восстановление и генерация обучаются с разной скоростью (восстановление быстрее и стабильнее генерации).
  • На ImageNet-256 модель SiT с GenFirst даёт gFID 0,97 (с CFG) и 1,45 (без CFG); модель MMDiT показывает GenEval 0,90 на генерации изображений по тексту.
  • Подход также распространили на общие визуальные латенты для генерации и распознавания образов и на непрерывную объединённую генерацию текста и изображений.

Почему это важно

Стандартный способ строить латентные генеративные модели, обучать VAE и генератор раздельно, друг за другом. Из-за этого латентное пространство получается оптимизированным под восстановление, а не под генерацию, и совместное (сквозное) обучение обеих частей сразу до сих пор почти всегда заканчивалось коллапсом латентного пространства или конфликтом между целями генерации и восстановления. Авторы GenFirst объясняют это тем, что энтропийный член в дивергенции Кульбака-Лейблера не даёт латентам вырождаться, а восстановление и генерация обучаются с разной скоростью, восстановление быстрое и хорошо направляемое сигналом, генерация медленнее и сложнее в оптимизации. Меняя порядок давления этих двух целей друг на друга, они впервые, по их утверждению, добились прямого сквозного обучения без коллапса.

Кому это важно

Исследователям и инженерам, которые строят или дообучают латентные генеративные модели изображений, диффузионные (в том числе SiT), авторегрессионные и мультимодальные (MMDiT) архитектуры, и сегодня вынуждены обучать автоэнкодер и генератор раздельно, мирясь с латентным пространством, не заточенным под генерацию. Метод также полезен командам, которые расширяют латентные генеративные модели на смежные задачи: совместное обучение генерации и распознавания образов на общих латентах и объединённую генерацию текста и изображений.

Как это применить

GenFirst, это порядок обучения, а не новая архитектура: сначала генеративная цель формирует латентное пространство при слабом давлении со стороны восстановления, затем давление восстановления постепенно усиливают, чтобы вернуть визуальные детали. Стратегию проверили на двух типах приоров, непрерывных авторегрессионных с точным правдоподобием и SiT с неявным правдоподобием, и получили gFID 0,97 (с CFG) и 1,45 (без CFG) для SiT на ImageNet-256, а также GenEval 0,90 для MMDiT на генерации изображений по тексту. Команды, работающие со схожими латентными архитектурами, могут применить тот же принцип смены порядка давления целей к своим пайплайнам обучения.

Можно ли доверять

Материал, это абстракт научной статьи о GenFirst на Hugging Face Papers. В самом абстракте не указаны ни имена авторов и их организации, ни дата публикации или площадка, ни статус выхода кода или весов модели, ни вычислительные затраты на обучение. В тексте также нет чисел для прямого сравнения с базовыми двухэтапными методами, поэтому оценить масштаб улучшения по абстракту нельзя, метрики gFID 0,97/1,45 и GenEval 0,90 приведены как самостоятельные значения, без сопоставления «было/стало». Утверждение о том, что это первое сквозное обучение без коллапса, формулировка самих авторов, а не результат независимой проверки.

Риски и подводные камни

Без раскрытых кода, весов и деталей обучения (вычислительный бюджет, время) независимо воспроизвести и проверить результаты по одному абстракту невозможно. Метрики даны без базовой линии сравнения с существующими двухэтапными методами, поэтому неизвестно, насколько именно GenFirst превосходит обучение VAE и генератора по отдельности, известны только абсолютные значения. Расширения на общие визуальные латенты и на объединённую генерацию текста и изображений упомянуты одной фразой без деталей архитектуры и отдельных метрик, так что судить об их зрелости преждевременно.