EditMGT: Unleashing Potentials of Masked Generative Transformers in Image Editing

Probleemstelling

Recente vooruitgangen in beeldbewerking worden gedomineerd door Diffusiemodellen (DMs). Hoewel deze modellen uitstekende visuele kwaliteit leveren, hebben ze een fundamenteel nadeel: hun globale denoising-dynamiek. Omdat DMs het volledige beeld iteratief verfijnen, worden lokale bewerkingsdoelen vaak verward met de volledige beeldcontext. Dit leidt tot onbedoelde wijzigingen in niet-doelgebieden, een fenomeen dat bekend staat als "editing leakage" (bewerkingslekkage). Bestaande oplossingen, zoals het gebruik van handmatige masks of inversietechnieken, zijn vaak onvoldoende flexibel, traag of kunnen de lekkage niet volledig voorkomen.

Methodologie: EditMGT

De auteurs introduceren EditMGT, het eerste beeldbewerkingskader dat gebruikmaakt van Masked Generative Transformers (MGTs) in plaats van diffusiemodellen. MGTs genereren beelden door meerdere gemaskerde tokens parallel te voorspellen, wat een lokaal decoderingsparadigma biedt dat inherent beter is in het behouden van niet-relevante gebieden.

De kern van de methode bestaat uit drie pijlers:

1. Architectuur en Attention Injection:

   • EditMGT bouwt voort op de bestaande Meissonic-architectuur (een MGT voor tekst-naar-beeld).
   • In plaats van extra parameters toe te voegen, wordt het originele beeld ingebracht als een extra conditionering via attention injection. Het model gebruikt de attention-mechanismen om de relatie tussen de tekstuele instructie, het originele beeld en het gegenereerde beeld te beheren.
   • Een bias-term ($\gamma$) in de attention-weighten stelt het model in staat om de invloed van het originele beeld tijdens de inferentie te regelen zonder de modelgrootte te vergroten.

2. Multi-layer Attention Consolidation:

   • De auteurs stellen vast dat de cross-attention-kaarten van MGTs waardevolle signalen bevatten voor het lokaliseren van de te bewerken gebieden, maar dat deze vaak onduidelijk of verspreid zijn.
   • Ze ontwikkelen een multi-layer attention consolidatie waarbij attention-weights van meerdere lagen (specifiek lagen 28-36) worden geaggregeerd.
   • Dit wordt aangevuld met Adaptive Filtering om de scherpte en ruimtelijke precisie van de attention-kaarten te verbeteren, waardoor de doelgebieden duidelijker worden geïdentificeerd.

3. Region-Hold Sampling:

   • Om de lekkageproblematiek op te lossen, introduceren ze Region-Hold Sampling.
   • Tijdens de iteratieve token-flipping (het generatieproces) worden tokens in gebieden met een lage attention-score (niet-doelgebieden) niet gewijzigd. In plaats daarvan worden ze teruggeschakeld naar hun oorspronkelijke tokens uit het invoerbeeld.
   • Dit wordt gestuurd door een drempelwaarde ($\lambda$). Tokens onder deze drempel worden "gehouden" (niet bewerkt), wat garandeert dat de integriteit van het oorspronkelijke beeld behouden blijft.

Dataset: CrispEdit-2M

Vanwege het gebrek aan hoogwaardige datasets voor beeldbewerking met hoge resolutie, hebben de auteurs CrispEdit-2M gecreëerd.

• Omvang: 2 miljoen samples (gefilterd uit een totale collectie van 5 miljoen).
• Resolutie: Alle afbeeldingen hebben een korte zijde van $\ge$ 1024 pixels.
• Categorieën: De dataset beslaat 7 diverse bewerkingscategorieën: toevoegen, vervangen, verwijderen, kleurverandering, achtergrondverandering, stijltransformatie en bewegingswijziging.
• Kwaliteit: De data is gegenereerd met open-source modellen (zoals FLUX.1 en Step1X-Edit) en streng gefilterd op esthetische kwaliteit en instructievolging.

Resultaten

EditMGT werd getest op vier standaard benchmarks (Emu Edit, MagicBrush, AnyBench, en GEdit-EN-full) en presteerde opmerkelijk goed, ondanks het kleine formaat van het model.

• Efficiëntie: Het model heeft slechts 960M parameters (2x tot 8x kleiner dan concurrenten zoals VAREdit-8B of GoT-6B).
• Snelheid: EditMGT is 6x sneller dan vergelijkbare modellen, met een bewerkingsduur van slechts 2 seconden voor een 1024x1024 afbeelding.
• Kwaliteit:
  • Beeldgelijkheid: State-of-the-art (SOTA) prestaties op CLIP-im en DINO-metrics voor Emu Edit en MagicBrush.
  • Stijlverandering: Een verbetering van 3,6% ten opzichte van de tweede beste methode.
  • Stijloverdracht: Een opvallende verbetering van 17,6% ten opzichte van het grote 12B-parameter model FluxKontext.dev.
  • Implicitie: Uitstekende prestaties bij impliciete instructies (1,7% verbetering).
• Kwalitatieve bevindingen: Het model begrijpt instructies beter (bijv. het correct "verwijderen" van warmte in plaats van toevoegen) en behoudt de oorspronkelijke structuur van het beeld beter dan DM-gebaseerde modellen.

Bijdragen en Significantie

De belangrijkste bijdragen van dit werk zijn:

1. Paradigmawissel: Het introduceert MGTs als een krachtig alternatief voor Diffusiemodellen in beeldbewerking, wat fundamenteel het probleem van "editing leakage" oplost door gebruik te maken van het lokale token-flipping-mechanisme.
2. Technische Innovatie: De combinatie van multi-layer attention consolidatie en region-hold sampling biedt een nieuwe manier om bewerkingsgebieden nauwkeurig te lokaliseren zonder handmatige masks.
3. Efficiëntie en Toegankelijkheid: Met een compact model van 960M parameters dat 6x sneller werkt, maakt EditMGT hoogwaardige beeldbewerking toegankelijker voor toepassingen met beperkte rekenkracht.
4. Dataset: De release van CrispEdit-2M vult een belangrijke lacune in de gemeenschap door een grote, hoogresolutie dataset voor instructiegebaseerde beeldbewerking beschikbaar te stellen.

Conclusie: EditMGT bewijst dat Masked Generative Transformers een veelbelovende en efficiëntere route zijn voor beeldbewerking dan de huidige diffusie-standaard, met name door hun inherente vermogen om niet-doelgebieden te beschermen en snellere inferentie te bieden.