[Статья] [Код] [пост на Abstract DL]
Введение
Исследователи со всего мира уже многие годы ищут оптимальную по скорости и качеству парадигму генерации. Диффузионные модели нынче самая ходовая тема со своими достоинствами и недостатками, но на ней клином свет не сошелся, и целесобразно рассматривать и альтернативу.
Некоторые время назад на свет появились авторегрессивные модели (DALL-E, Parti) для генерации, мотивированные успехом больших языковых моделей. Однако, они не приобрели столь большой популярности, как диффузионки, из-за худшего качества и медленной генерации. Генерить токен за токеном долгую последовательность долго и утомительно.
(прим. вычислительная сложность
O(n^6) у тех, кто не знает, что можно кэшировать keys и values). И чуваки решили пересмотреть авторегрессионные модели, так чтобы они работали быстро 🏃♂️ и эффективно 💪.
Метод
Суть метода в следующем - изображение кодируют иерархическим автокодировщиком, где каждый уровень представлен в виде некоторой последовательности токенов. В некотором смысле, подобная схема напоминает лапласовскую пирамиду. Сначала пытаются приблизить одним токеном насколько возможно исходную картинку, затем увеличивают feature map и уже приближают residual, и так вплоть до максимального разрешения. Токены лежат в некотором обучаемом кодбуке и на этапе квантования, токен, полученный некоторой проекцией приближается к ближайшему вектору из кодовой книги. Словарь общий для всех уровней иерархии.
Автокодировщик обучается на взвешенную сумму l2 лосса, адверсариального лосса и LPIPS.
Генерация выглядит так - на каждом уровне за раз генерируются
h_k x w_k токенов, соответствующих данному уровню, обуславливаясь на все токены с прошлых уровней. Таким образом, нужно делать порядка log(n) проходов вместо n^2, что является безусловным прогрессом. В качестве модели берут обычный декодер-like трансформер без модных финтифлюшек типа SwiGLU, RoPE и прочих приблуд, а старый дедовский GPT-2-like c разницей лишь в том, что вместо обычных LayerNormов используются обусловленные (на класс).