Продолжаем разбирать работу ByteDance на тему авторегрессионной генерации изображений. В первой части рассказали, как авторы пытались увеличить размер кодбука и боролись с «коллапсом» с помощью lookup-free-квантизации.
Однако полечив одно, неизбежно столкнулись с другой проблемой: в авторегрессионных языковых моделях для предсказания очередного токена используют линейную классификационную голову, выходная размерность которой равна размеру словаря токенизатора. Чтобы применить такой способ к генерации картиночного токена, где размер кодбука VQ-VAE равен 2^D, нужно хранить в памяти матрицу размера h x 2^D, где h — размерность скрытого слоя трансформерного бэкбона. Для сколь-нибудь больших D и h порядка 10^3 это будет приводить к OOM.
Можно попробовать предположить независимость каждого из D каналов бинарного токена, как было сделано в Infinity — одной из прошлых работ ByteDance по авторегрессионной генерации. Это снизит рост матрицы головы с экспоненциального до линейного: теперь нам потребуется матрица размера h x 2D. Однако при таком предположении качество картиночной генерации заметно падает.
Именно в этом месте исследователям на помощь приходит диффузия. Для предсказания D-мерного бинарного токена (он же — вершина D-мерного куба) предлагают выучить отображение распределения гауссовского шума в вершины D-мерного куба при условии эмбеддинга z, для которого мы и предсказываем бинарный токен. Дизайн денойзера в статье особо не обсуждается и не проверяется: берут «небольшой» трансформер и обучают его по замшумлённому токену x_t, времени t и эмбеддингу z предсказывать x. Причём диффузионный лосс считают во flow-matching-формате, то есть предсказание сети перепараметризовывается в предсказание скорости. После финального шага диффузии от предсказания берётся знак. По сравнению с диффузией на стандартных VAE-латентах, значения которых никак не ограничены, такая жёсткая схема помогает избежать накопления ошибки, которое происходит как в самой диффузии, так и в итеративном авторегрессионном процессе, а также приводит к более быстрой сходимости обучения.
Наконец, авторы пытаются выжать максимум из предсказания картиночных токенов, переходя от потокенной генерации к параллельной генерации патчей из PxP токенов. Во-первых, это ускоряет процесс: предсказывать диффузией сразу целый патч оказывается эффективнее, чем генерировать токены по одному. Во-вторых, это лучше соответствует локальной структуре изображения. Во время генерации патча эмбеддингов LLM-бэкбон использует bidirectional-маску в аттеншне-механизме, тогда как для текстовых токенов остается каузальная. Это положительно влияет на качество генерации.
Так авторегрессия сохраняется на уровне патчей, а внутри патча используется диффузионная генерация.
В экспериментах авторы показывают, что увеличение размера кодбука действительно улучшает реконструкцию. При этом наши тесты показывают, что метрики не совсем честно отражают визуальное качество. Хотя по сравнению с другими дискретными токенизаторами реконструкции выглядят неплохо, особенно на высоких разрешениях.
В последней части разбора посмотрим на результаты масштабирования архитектуры в t2i-сетапе генерации и обсудим аблейшены, проведённые авторами.
Разбор подготовил
CV Time