Distilling Specialized Orders for Visual Generation

Titel: Distilling Specialized Orders for Visual Generation

Auteurs: Rishav Pramanik et al. (Stony Brook University, ILLS, Mila, ServiceNow Research, etc.)

---

1. Het Probleem

Autoregressieve (AR) beeldgeneratiemodellen worden steeds populairder vanwege hun schaalbaarheid en hoge kwaliteit. Echter, traditionele AR-modellen zijn vastgezet op een specifieke generatievolgorde, meestal een raster-scan (van linksboven naar rechtsonder). Dit leidt tot twee belangrijke beperkingen:

1. Gebrek aan flexibiliteit: Modellen met een vaste volgorde kunnen taken zoals inpainting (invullen van gemaskeerde gebieden), outpainting (uitbreiden van een afbeelding) of bewerken niet uitvoeren zonder hertraining. Ze zijn structureel geblokkeerd in één generatietraject.
2. Compromis bij Any-Order modellen: Bestaande "Any-Order" AR-modellen lossen dit op door te leren genereren onder willekeurige volgorde van patches. Dit biedt flexibiliteit, maar gaat ten koste van de prestaties. Het model moet zijn beperkte capaciteit verdelen over het leren van $N!$ mogelijke permutaties (waarbij $N$ het aantal patches is), wat leidt tot een lagere kwaliteit van de gegenereerde afbeeldingen vergeleken met modellen die op één geoptimaliseerde volgorde zijn getraind.

De kernvraag van het paper is: Kan een enkel autoregressief model zowel de flexibiliteit van "Any-Order" behouden als de kwaliteit verbeteren door een gespecialiseerde generatievolgorde te ontdekken?

2. Methodologie: Ordered Autoregressive (OAR)

De auteurs stellen Ordered Autoregressive (OAR) voor, een zelf-distillatiepijplijn die de beste van beide werelden combineert. Het proces verloopt in drie fasen:

Fase 1: Training van een Any-Order AR-model (AO-AR)

Eerst wordt een standaard autoregressief model getraind om conditionele verdelingen te leren onder willekeurige volgorde van patches.

• Het model ontvangt een afbeelding (omgezet in discrete tokens via een VQ-VAE encoder) en een conditionele input (bijv. een label of tekst).
• Om willekeurige volgorde mogelijk te maken, wordt de positie van de volgende token expliciet aan het model gegeven (via positionele tokens of dubbele positionele encoding).
• Het doel is het maximaliseren van de log-waarschijnlijkheid over alle mogelijke permutaties.

Fase 2: Extractie van Gespecialiseerde Volgorde (Self-Distillation)

Na het trainen van het Any-Order model, wordt er een specifieke, optimale generatievolgorde afgeleid voor elke trainingsafbeelding.

• Greedy Order Selection: In plaats van complexe zoekalgoritmen (zoals beam search), gebruikt het model zijn eigen vertrouwensscores (predictie-waarschijnlijkheid).
• Bij elke stap $i$ kijkt het model naar alle mogelijke volgende patch-locaties en selecteert het de locatie die de hoogste waarschijnlijkheid heeft voor het voorspellen van de juiste token, gegeven de reeds gegenereerde context.
• Dit creëert een "pseudo-label" voor de volgorde ($l^*$) die specifiek is voor de inhoud van de afbeelding. Het model leert zo dat het eerst de meest "zeker" te voorspellen gebieden genereert (bijv. homogene achtergronden of duidelijke objecten) en pas later de moeilijkere, fijnere details.

Fase 3: Fine-tuning met Gespecialiseerde Volgorde

Het oorspronkelijke Any-Order model wordt vervolgens fine-ge-tuned op de dataset van (afbeelding, gespecialiseerde volgorde)-paren.

• Capaciteit Herallocatie: Door te focussen op één specifieke, data-gedreven volgorde per afbeelding, kan het model zijn capaciteit herverdelen van het leren van alle $N!$ permutaties naar het verfijnen van de kwaliteit langs deze ene optimale pad.
• Behoud van Flexibiliteit: Om te voorkomen dat het model de Any-Order vaardigheden vergeet (catastrophic forgetting), wordt er gefine-tuned met een lage leerfrequentie en een order-annealing schema. Dit schema gaat geleidelijk over van gemengde volgorde naar puur gespecialiseerde volgorde.

3. Belangrijkste Bijdragen

1. OAR Pijplijn: Een nieuwe zelf-distillatiestrategie die gespecialiseerde generatievolgorde extrahert uit de vertrouwensscores van een Any-Order model en het model daarop fijntuned.
2. Theoretisch Inzicht: Het paper toont aan dat "order specialisatie" werkt als een herallocatie van modelcapaciteit. Door te focussen op één traject per afbeelding, verbetert de kwaliteit zonder de onderliggende conditionele verdelingen die multi-task flexibiliteit mogelijk maken, te wissen.
3. Zero-Shot Flexibiliteit: In tegenstelling tot raster-scan modellen, behoudt OAR het vermogen om zero-shot inpainting, outpainting en bewerking uit te voeren zonder extra training.
4. Empirische Validatie: Uitgebreide experimenten op ImageNet, Fashion Products en CelebA-HQ tonen consistente verbeteringen.

4. Resultaten

De resultaten tonen aan dat OAR de kwaliteit van raster-scan en willekeurige Any-Order baselines overtreft:

• ImageNet (256x256):
  • De FID (Fréchet Inception Distance) verbeterde van 2.39 (Any-Order baseline) naar 2.17 (OAR met annealing).
  • Dit is een prestatie die concurrerend is met state-of-the-art diffusion-modellen en andere AR-modellen, maar dan met behoud van de Any-Order flexibiliteit.
  • Menselijke evaluatie toonde een 64% voorkeur voor OAR-afbeeldingen ten opzichte van de baseline.
• Zero-Shot Taken: Het model kan succesvol afbeeldingen invullen (inpainting) of uitbreiden (outpainting) zonder hertraining, omdat het de Any-Order condities behoudt.
• Andere Datasets: Op de Fashion Product dataset en CelebA-HQ (met tekst-conditionering) werden vergelijkbare verbeteringen gezien (bijv. FID verbetering van 1.94 naar 1.41 op CelebA-HQ).
• Inference Efficiency: Hoewel het selecteren van de beste volgorde theoretisch meer berekening vereist, gebruiken de auteurs geoptimaliseerde KV-caching en parallelle query-evaluatie om de inferentietijd vergelijkbaar te houden met standaard raster-scan AR-modellen.

5. Betekenis en Conclusie

Dit paper biedt een belangrijke doorbraak in autoregressieve beeldgeneratie. Het lost het fundamentele compromis op tussen flexibiliteit (kunnen doen wat je wilt, zoals inpainting) en kwaliteit (de beste mogelijke afbeelding genereren).

Door te leren dat er voor elke afbeelding een "natuurlijke" volgorde bestaat die gebaseerd is op de inhoudelijke complexiteit en de zekerheid van het model, kunnen AR-modellen nu zowel hoogwaardige afbeeldingen produceren als als universele tools fungeren voor diverse visuele taken. De methode vereist geen architecturale wijzigingen of extra annotaties, wat het een zeer efficiënte en toepasbare oplossing maakt voor de huidige generatie visuele AI-modellen.