Метод

За основу мы взяли архитектуру модели из STAR и последовательно анализируя и улучшая ее пришли к конечной модели.

Первое, с чем мы столкнулись, это с тем, что оригинальная архитектура становилась крайне нестабильной в процессе обучения. Внутренние активации модели вырастали до очень больших значений, которые уже нельзя было представить в машинной точности, и обучение разваливалось.

Следуя работе Lumina, мы добавили дополнительные нормализационные слои в модель, и данная архитектурная модификация стабилизировала обучение. Кроме того, качество самой модели тоже улучшилось.

Далее, мы проанализировали карты внимания авторегрессионной модели, и обнаружили, что текущее разрешение почти не смотрит на прошлые, поэтому разрешения можно генерировать независимо друг от друга. Убрав авторегрессию на прошлые разрешения мы нисколько не потеряли в качестве, и при этом ускорили модель примерно на 20-30%.

И последняя ключевая находка оказалась в том, что технику classifier-free-guidance (CFG), улучшающую качество генераций и соответствие текстовому запросу, но требующую два прогона через модель вместо одного, можно отключить на высоких разрешениях без ухудшения конечного результата. За счет этого, можно добиться почти двухкратного ускорения по сравнению со стандартной процедурой генерации с CFG.

Полученную модель мы назвали Switti - Scale-wise transformer for text-to-image synthesis, так как она генерирует изображение по тексту разрешение за разрешением.