Technische Samenvatting: MOSAIK

Probleemstelling

Pixelruimte-diffusiemodellen (PDM's) bieden een aanzienlijk voordeel ten opzichte van latente diffusiemodellen (LDM's) door het reconstructieplafond te vermijden dat wordt opgelegd door Variational Autoencoders (VAE's), waardoor fijne, hoogfrequente ruimtelijke details behouden blijven. Het genereren direct in de pixelruimte resulteert echter in een substantieel hoger aantal tokens. Omdat self-attention mechanismen in transformers een kwadratische complexiteit vertonen ten opzichte van het aantal tokens, brengen PDM's prohibitieve computationele kosten met zich mee.

Bestaande efficiëntiemethoden proberen dit te verzachten door de patchgroottes (het ruimtelijke gebied dat door een enkele token wordt gerepresenteerd) te variëren over de denoising-tijdstappen of netwerkdiepte. Hoewel deze benaderingen het totale aantal tokens verminderen, passen ze een ruimtelijk uniforme patch-indeling toe over de gehele afbeelding op elk gegeven tijdstip. Deze uniformiteit creëert een fundamentele inefficiëntie: het verspilt ofwel rekenkracht aan eenvoudige, regio's met weinig informatie, of beperkt daarentegen de representatiecapaciteit in gevoelige regio's waar fijne details cruciaal zijn. Het artikel stelt dat de visuele informatiedichtheid inherent heterogeen is, wat een dynamische, per-regio allocatiestrategie vereist die huidige methoden missen.

Methodologie: MOSAIK

MOSAIK (Multi-patch cOntent-aware Spatial Allocation of Image toKens) is een framework ontworpen om patchgroottes te variëren over verschillende beeldregio's en denoising-stappen op basis van de voorspelde "schade" (fidelity loss) veroorzaakt door vergroving. De aanpak is gebouwd op de PixelDiT backbone en maakt gebruik van een trainingspipeline in drie stadia:

Fase 1: Multi-Patch Vloeiendheid

Het doel is om het model in staat te stellen te opereren op meerdere uniforme patchgroottes ($P \in \{16, 32, 64\}$) zonder catastrofale vergetelheid van de oorspronkelijke generatieve prior.

• Architectuur: De gewichten van de backbone zijn bevroren. Low-Rank Adaptation (LoRA) modules worden geïnjecteerd in de DiT-blokken. Afzonderlijke tokenizers worden geïntroduceerd voor $p32$ en $p64$, geïnitialiseerd door de vooraf getrainde $p16$ projectie te componeren met bilineaire downsampling.
• Training: Het model ondergaat zelfdistillatie, waarbij wordt geoptimaliseerd om de velocity-voorspellingen van het bevroren $p16$ teacher-model te matchen over alle uniforme patchgroottes.

Fase 2: Mixed-Patch Adaptatie

Deze fase stelt het model in staat om willekeurige, heterogene patch-indelingen te verwerken (waarbij verschillende regio's verschillende patchgroottes gebruiken) binnen een enkele forward pass.

• Positionele Embeddings: Standaard Rotary Positional Embeddings (RoPE) gaan uit van een uniform rooster. MOSAIK adresseert de ruimtelijke mismatch van gemengde patches door elke token te verankeren aan het geometrische centrum van zijn patch binnen een genormaliseerd $p16$ coördinatenstelsel.
• Training: De generator wordt getraind op willekeurig gesamplede heterogene indelingen, waarbij de ruimtelijke ordening van patchgroottes bewust wordt losgekoppeld van de beeldinhoud. Dit dwingt het model om elke indeling uit te voeren, ongeacht de uitlijning met visuele kenmerken.

Fase age 3: Damage-Guided Layout Predictor

Deze fase introduceert een lichtgewicht voorspeller die dynamisch tokens toewijst op basis van de geschatte degradatie veroorzaakt door het vergroven van specifieke regio's.

• Definitie van Schade: "Schade" wordt gedef Definieerd als het kwadratische verschil tussen de voorspelling van een uniforme $p64$ indeling en de hoog-getrouwe $p16$ teacher, gemiddeld over niet-overlappende $64 \times 64$ regio's.
• Predictor Training: Een lichtgewicht convolutioneel netwerk met FiLM-conditionering wordt getraind om de schadekaart direct te regresseren vanuit tussenliggende backbone-features (specifiek de DiT block-6 activaties) en tekst-embeddings. Dit voorkomt de noodzaak van een secundaire forward pass tijdens inferentie.
• Inference Strategie: Bij elke denoising-stap schat de predictor de schadekaart in. Regio's met een hoge voorspelde schade krijgen fijnere patches ($p16$), terwijl regio's met lage schade grovere patches ($p32$ of $p64$) toegewezen krijgen. De allocatie wordt beperkt door een door de gebruiker gespecificeerd tokenbudget, waarbij prioriteit wordt gegeven aan verfijningen die de hoogste schadevermindering per toegevoegde token bieden.

Belangrijkste Bijdragen

1. Ruimtelijk Adaptieve Patchificatie: MOSAIK is het eerste framework dat heterogene patch-indelingen toepast binnen een pixelruimte-diffusiemodel, waarmee het voorbijgaat aan de ruimtelijke uniformiteitsbeperkingen van eerdere efficiëntiemethoden.
2. Damage-Guided Allocatie: De auteurs formuleren patch-allocatie als een schade-voorspellingsprobleem, waarbij een lichtgewicht predictor wordt gebruikt om regio's te identificeren waar vergroven de meeste fidelity loss zou veroorzaken.
3. Drie-fasen Trainingspipeline: Een nieuwe trainingsstrategie die eerst multi-patch vloeiendheid vestigt, vervolgens adapteert naar gemengde indelingen, en tot slot een aparte predictor traint om token-routing te sturen zonder de generator opnieuw te trainen.
4. Integratievermogen: Het framework is ontworpen om orthogonaal te zijn aan temporele efficiëntiemethoden (zoals feature caching), waardoor gecombineerde optimalisatie mogelijk is.

Resultaten

Experimenten werden uitgevoerd op 1024x1024 beeldgeneratie met de GenEval en DPG-Bench benchmarks, waarbij MOSAIK werd vergeleken met baselines inclusief uniforme patch-reductie, temporele patch-scheduling (DDiT) en feature caching (TeaCache, TaylorSeer).

• Efficiëntie vs. Kwaliteit: MOSAIK bereikt een reductie van 70% in FLOPs en een reductie van 83% in het aantal tokens, terwijl het de volledige-compute PixelDiT evenaart op GenEval (0.74 score) en slechts 1.0 punt daalt op DPG-Bench (82.5 vs. 83.5).
• Vergelijking met Temporele Scheduling: In agressieve compute-regimes (bijv. 60-70% reductie in FLOPs) presteert MOSAIK consistent beter dan DDiT. Terwijl DDiT een enkele patchgrootte globaal toepast, herricht MOSAIK de rekenkracht dynamisch naar gevoelige regio's (bijv. subjectgrenzen), waardoor de structurele integriteit en fijne details zoals texturen en randen behouden blijven.
• Vergelijking met Caching: MOSAIK presteert beter dan caching-gebaseerde baselines (TeaCache, TaylorSeer) bij gelijke budgetten, vooral in zeer beperkte regimes.
• Complementariteit: Wanneer gecombineerd met caching-methoden, behoudt MOSAIK de prestaties bij zelfs hogere compressieniveaus (tot 80% FLate FLOPs reductie), waar caching alleen leidt tot significante kwaliteitsafname.
• Ablatie: Het vervangen van de damage-guided predictor door een willekeurige of inhouds-blinde allocator resulteert in een significante prestatieverslechtering, wat valideert dat de voorspelde schade-ranking werkelijke informatie bevat over ruimtelijke sensitiviteit.

Betekenis en Claims

Het artikel claimt dat MOSAIK aantoont dat ruimtelijke heterogeniteit in visuele informatie een cruciale factor is voor efficiënte generatie die over het hoofd is gezien door methoden die zich uitsluitend richten op temporele of diepte-gebaseerde adaptatie.

• Optimale Resource Allocatie: De auteurs stellen dat uniforme fijne patchificatie rekenkracht verspilt aan eenvoudige regio's, terwijl uniforme grove patchificatie details in complexe regio's vernietigt. MOSAIK overbrugt deze kloof door middelen toe te wijzen waar ze het meest nodig zijn.
• Voordeel van de Pixelruimte: Het framework maakt gebruik van de inherente ruimtelijke redundantie die beschikbaar is in pixelruimte-modellen (in tegen tegenstelling tot latente modellen waar VAE-compressie de redundantie al heeft geconsumeerd) om dynamisch vergroving te gebruiken zonder de noodzaak van end-to-end gezamenlijke optimalisatie van een router en generator.
• Schaalbaarheid: De aanpak maakt stabiele prestaties mogelijk over een breed scala aan compute-budgetten, van gematigde reducties tot agressieve beperkingen, zonder dat hertraining voor verschillende budgetniveaus vereist is.

Het artikel concludeert dat MOSAIK een nieuwe standaard zet voor efficiëntie in pixelruimte-diffusie, door een zeer competitief alternatief te bieden voor temporele scheduling en caching, en suggereert dat toekomstig werk deze principes kan uitbreiden naar video-diffusie, waarbij redundantie over zowel ruimte als tijd wordt geëxploiteerd.