Архитектура SpecSlot снизила FID на 2,84 против JiT-CM и сэкономила 3,9% GFLOPs

Генеративные модели, которые работают прямо в пикселях, обходятся без отдельно обученного автоэнкодера и его узкого места сжатия. Цена за это — длинная последовательность: изображение целиком приходится превращать в токены трансформера. Архитектура SpecSlot разделяет две вещи, которые обычно жёстко связаны: насколько детально модель выдаёт картинку и сколько токенов участвуют в глобальном внимании. На ImageNet-512 конфигурация W256/K4 снизила FID (метрику расстояния между сгенерированными и реальными изображениями, чем ниже, тем лучше) на 2,84 относительно JiT-CM. Это контрольная модель JiT, подобранная авторами под те же FLOPs. При этом SpecSlot потратила на 3,9% меньше GFLOPs.

Почему патч-сетка упирается в квадратичное внимание

В большинстве пиксельных трансформеров один токен соответствует одному пространственному патчу. Сетка патчей одновременно определяет, что модель читает, что предсказывает каждый токен и сколько токенов попадает в глобальное внимание. Если на изображении 256×256 уменьшить патч с 16 до 8, число токенов вырастет с 256 до 1024, а квадратичная часть стоимости внимания вырастет вместе с ним.

Существующие подходы решают проблему с разных сторон. JiT пропускает каждый патч через глобальную последовательность. Многомасштабные системы выносят детали на дополнительные стадии или в «детализирующие» головы. Методы эффективного внимания сокращают последовательность окнами, выборочными глобальными токенами или обученной разреженностью. Spectral Forcing использует частотную структуру, но только чтобы фильтровать зашумлённый вход, а его токены остаются пространственными патчами. Во всех этих схемах детализация выхода привязана к сетке патчей, даже если обмен между токенами сжат или ограничен.

Read, Allocate, Route: якоря и спектральные слоты

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

Read: изображение читается на сетке пространственных якорей, которые собирают входные данные. Небольшая свёртка с нулевой инициализацией позволяет соседним ячейкам обменяться информацией до разделения ролей. Allocate: каждый якорь разворачивается в индексированные слоты, и каждому слоту достаётся своя непересекающаяся радиальная группа коэффициентов DCT (дискретного косинусного преобразования). Route: роль определяет коммуникацию. В глобальную стадию от каждого якоря попадает только слот с самой низкочастотной группой, а внутри якоря взаимодействуют все слоты.

Расписание внимания называется GGLL (global–global–local–local): по одному низкочастотному носителю на якорь идёт в глобальные блоки, а все слоты работают в локальных блоках внутри якоря. Финальные состояния слотов декодируют спектральную основу на базе lapped-DCT и пространственную поправку, выровненную по якорям. Их складывают в пиксельном пространстве. Поправка не добавляет новых токенов. Состояние, цель и функция потерь остаются в пространстве пикселей, а модель обучается в рамках rectified flow.

В итоге появляются две отдельные переменные распределения. Одна задаёт плотность пространственных якорей, другая — спектральное деление внутри якоря. Общее число слотов растёт, а глобальная последовательность содержит только носителей. Фиксированный бюджет токенов можно тратить либо на более плотные якоря, либо на более тонкие спектральные роли. Выбор размера патча превращается в явную задачу компромисса между качеством и вычислениями.

Что показали замеры

На ImageNet-256 и ImageNet-512 SpecSlot даёт две измеренные рабочие точки между дешёвой и дорогой конфигурациями JiT. На ImageNet-512 вариант W256/K4 уступает JiT-B/16 по FID всего на 0,69, но обходится на 38,2% дешевле. Более крупные выходные окна покрывают больше пикселей при неизменной глобальной последовательности. Авторы связывают с этим то, что выигрыш относительно JiT-CM оказался больше.

В генерации изображений по тексту SpecSlot превзошла JiT-B/8 на 1,06 пункта в DPG-Bench (тесте на соответствие картинки сложному описанию) при стоимости прямого прохода ниже на 40,6%.

Есть и оговорка самих авторов: равные доли коэффициентов у слотов не означают равной стоимости, поскольку коммуникация и проекции масштабируются с разбиением по-разному. Поэтому сравнивать нужно измеренные затраты, а не только число токенов. Код и веса авторы выложили на GitHub и Hugging Face.