FOCUS: Optimal Control for Multi-Entity World Modeling in Text-to-Image Generation

Titel: FOCUS: Optimale Controle voor Multi-Entiteit Wereldmodellering in Tekst-naar-Afbeelding Generatie

1. Het Probleem

Moderne tekst-naar-afbeelding (T2I) modellen, zoals Stable Diffusion 3.5 en FLUX.1, presteren uitstekend bij prompts met één enkel onderwerp. Echter, bij prompts met meerdere entiteiten (bijv. "een astronaut, een viool en een zonnebloem") vertonen deze modellen vaak ernstige tekortkomingen:

• Attribuutlekkage (Attribute Leakage): Eigenschappen van het ene object worden onterecht aan een ander object toegeschreven (bijv. de "rode hoed" van de astronaut komt op de viool).
• Identiteitsverstrengeling (Identity Entanglement): Objecten smelten samen tot één vervormde vorm.
• Omissies: Sommige onderwerpen in de prompt verschijnen niet in de gegenereerde afbeelding.

Deze fouten wijzen op een zwakke "wereldmodellering": het model heeft geen goed begrip van welke entiteiten aanwezig zijn en welke attributen bij welke entiteit horen. Bestaande oplossingen zijn vaak heuristisch, vereisen veel handmatige tuning, en zijn niet goed portabel naar moderne Flow Matching (FM) architecturen.

2. Methodologie

De auteurs introduceren een principieel theoretisch raamwerk dat generatie omvormt tot een Stochastische Optimale Controle (SOC) probleem binnen het Flow Matching kader.

A. Theoretisch Kader (Flow Matching als SOC)
In plaats van het generatieproces als een statisch proces te zien, wordt het behandeld als een dynamisch systeem waarbij een kleine controle $u$ wordt toegepast op de drift (snelheid) van het model. Het doel is om een compromis te vinden tussen:

1. Fideliteit: Dichtbij blijven bij het oorspronkelijke base model (stijl en kwaliteit behouden).
2. Disentanglement: Het maximaliseren van de scheiding tussen verschillende onderwerpen.

Dit wordt geformuleerd als een kwadratisch optimalisatieprobleem waarbij de kostenfunctie bestaat uit een "running cost" (die verstrengeling meet) en een controlekosten-term.

B. De Loss Functie: FOCUS
Om de "running cost" te definiëren, introduceren de auteurs FOCUS (Flow Optimal Control for Unentangled Subjects).

• Principe: Cross-attention kaarten in T2I-modellen corresponderen met de ruimtelijke plaatsing van entiteiten.
• Berekening: FOCUS behandelt deze attention kaarten als kansverdelingen. Het berekent de Jensen-Shannon Divergentie (JSD) om twee dingen te optimaliseren:
  1. Intra-subject consistentie: De attention kaarten voor één specifiek onderwerp moeten unimodaal en consistent zijn.
  2. Inter-subject scheiding: De gemiddelde attention kaarten van verschillende onderwerpen moeten zo ver mogelijk van elkaar verwijderd zijn (maximale divergentie).
• Dit resulteert in een probabilistische loss die verstrengeling straft zonder handmatige regels voor specifieke tokens.

C. Twee Implementaties
Het raamwerk biedt twee manieren om dit toe te passen:

1. Test-time Controller (Training-vrij):
   • Een "on-the-fly" controller die de sampling dynamiek tijdens de inferentie aanstuurt.
   • Het benadert de co-toestand (adjoint) lokaal en past de snelheid van het base model direct aan met een gradiënt-gebaseerde update.
   • Geen retraining nodig; werkt met bestaande checkpoints.
2. Fine-tuning via Adjoint Matching:
   • Een lichtgewicht fine-tuning methode die een controle-netwerk leert om de aanwijzingen (guidance) te amortiseren.
   • Het gebruikt "Adjoint Matching" om een efficiënt trainingsdoel te creëren zonder de volledige trajecten te hoeven differentiëren.
   • Het resultaat is een model dat sneller inferentie uitvoert (geen extra backpropagation tijdens generatie) en generaliseert naar nieuwe prompts.

3. Belangrijkste Bijdragen

1. FOCUS Loss: Een nieuwe, probabilistische attention-binding loss die verstrengeling meet op basis van divergentie tussen kansverdelingen, schaalbaar voor willekeurig aantal onderwerpen.
2. Unificatie via SOC: Een theoretisch raamwerk dat bestaande heuristieken (zoals Attend&Excite, CONFORM) verenigt onder één optimale controle-objectief, waardoor ze portabel worden naar moderne Flow Matching modellen.
3. Twee Praktische Instantiaties:
   • Een training-vrije test-time controller die werkt op commodity GPUs.
   • Een fine-tuning route (Adjoint Matching) die een geoptimaliseerd controller-netwerk leert met minimale extra parameters (<0.1% van de weights via LoRA).
4. Extensie naar Diffusion: De methode is niet beperkt tot Flow Matching; via een flow-diffusie correspondentie werkt het ook op klassieke Denoising Diffusion modellen (zoals SD 1.5).

4. Resultaten

De auteurs evalueren hun methode op Stable Diffusion 3.5 en FLUX.1, en testen transfer naar Stable Diffusion 1.5.

• Kwantitatieve Prestaties:
  • Zowel de test-time controller als de fine-tuned modellen verbeteren consequent de multi-subject uitlijning (Image-Text en Text-Text metrics) ten opzichte van de base modellen.
  • FOCUS behaalt de hoogste "Composite Scores" en presteert beter dan bestaande heuristieken zoals Attend&Excite, CONFORM en Divide&Bind.
  • Bij fine-tuning behoudt het model de stijl van het base model en diversiteit, terwijl de verstrengeling aanzienlijk afneemt.
• Menselijke Evaluatie:
  • In een pairwise preference studie (2.000 vergelijkingen) wint FOCUS consistent van de base modellen en andere heuristieken.
  • Menselijke beoordelaars geven de voorkeur aan afbeeldingen gegenereerd met FOCUS vanwege betere attributie en minder lekkage.
• Efficiëntie:
  • Test-time control voegt ongeveer 90-138% extra latentie toe (door backpropagation), maar werkt op standaard GPUs.
  • Fine-tuned modellen hebben geen extra inferentie- overhead ten opzichte van het base model.

5. Betekenis en Conclusie

FOCUS biedt een fundamentele doorbraak in het beheersen van complexe multi-entiteit scènes in generatieve AI. Door het probleem te herformuleren als een optimalisatieprobleem (Optimal Control) in plaats van een verzameling heuristieken, biedt het:

• Robuustheid: Werkt over verschillende architecturen heen (van U-Net tot Diffusion Transformers).
• Flexibiliteit: Biedt een enkele "knop" (de controlesterkte $\lambda$) om het compromis tussen creativiteit/stijl en strikte object-scheiding te regelen.
• Toekomstbestendigheid: De fine-tuning methode biedt een route om modellen permanent te verbeteren voor multi-subject taken zonder de noodzaak van zware retraining of complexe spatial constraints (zoals bounding boxes).

Kortom, FOCUS stelt tekst-naar-afbeelding modellen in staat om "wereldmodellen" te bouwen waarin objecten en hun eigenschappen correct en onafhankelijk worden gerespecteerd, wat essentieel is voor toepassingen zoals storytelling, wetenschappelijke visualisatie en complexe compositie.