Часть 2, методы
Приступив к созданию Switti (Scale-wise transformer for text-to-image), мы взяли за основу архитектуру STAR и начали постепенно её улучшать. Оригинальная модель оказалась нестабильной при обучении, и, мотивируясь работами по диффузионным трансформерам, мы добавили в модель дополнительные нормализационные слои, что стабилизировало обучение и улучшило итоговое качество.
Затем мы обратили внимание, что модель на входе на текущем уровне уже получает информацию о всех предыдущих уровнях по построению и при этом дополнительно смотрит на них с помощью attention-слоев, что показалось нам лишним и неэффективным. Убрав авторегрессию на прошлые разрешения, удалось ускорить пайплайн генерации для 512х512 изображений на 11%, уменьшить потребление памяти и немного улучшить качество генерации.
Также мы заметили, что модель слабо опирается на текстовые описания на самых последних уровнях генерации. Это натолкнуло на мысль, что можно не использовать технику classifier-free-guidance (CFG) на уровнях высокого разрешения. Напомним, что CFG играет важную для повышения качества генерации и соответствия запросу в text-to-image моделях, но при этом требует дополнительный прогон модели на каждом шаге. Поэтому отключение CFG на последних уровнях значительно ускоряет генерацию. Более того, мы заметили, что помимо скорости, модель также продуцирует меньше артефактов при генерации мелких деталей и позволяет использовать более высокие значения CFG без ущерба качеству.
Обзор подготовил
CV Time
