Technische Samenvatting: Universele Few-Shot Ruimtelijke Controle voor Diffusiemodellen

Probleemstelling

Pre-trained tekst-naar-afbeelding (T2I) diffusiemodellen hebben opmerkelijk succes behaald bij het genereren van hoogwaardige afbeeldingen vanuit natuurlijke taalprompts. Echter, taalgebaseerde controle mist vaak de precisie die vereist is voor fijnmazige visuele manipulatie, zoals het afdwingen van specifieke objectposities, vormen of scène-indelingen. Om dit aan te pakken, maken bestaande methoden gebruik van ruimtelijke controle-adapters (bijv. ControlNet, Uni-ControlNet) die expliciete ruimtelijke signalen (randkaarten, diepte, pose) injecteren in bevroren diffusiemodellen.

Ondanks hun effectiviteit, kampen deze bestaande adapters met twee kritieke beperkingen:

1. Hoge Trainingskosten: Ze worden doorgaans onafhankelijk getraind voor elke specifieke ruimtelijke controletaak, wat aanzienlijke computationele middelen en duizenden gelabelde voorbeelden vereist (bijv. 150K afbeeldingen per taak).
2. Slechte Aanpasbaarheid: Het hergebruiken van vooraf getrainde multi-task adapters of training-vrije methoden faalt vaak bij het generaliseren naar nieuwe ruimtelijke condities die significant verschillen van de trainingsdistributie. Training-vrije methoden vermijden weliswaar de kosten voor data, maar vertrouwen vaak op langdurige latente optimalisatieprocessen die de generatietijd verhogen en een beperkte controle bieden.

Het artikel stelt de noodzaak vast voor een veelzijdig framework dat in staat is om diffusiemodellen te besturen met nieuwe ruimtelijke condities met behulp van slechts een klein aantal geannoteerde voorbeelden (few-shot learning) tijdens de testfase, zonder significante overhead tijdens de generatietijd.

Methodologie: Universele Few-Shot Controle (UFC)

De auteurs stellen Universal Few-so-Shot Control (UFC) voor, een universele controle-adapter ontworpen om bevroren T2I-diffusiemodellen te sturen met willekeurige nieuwe ruimtelijke condities met behulp van een support-set van beeld-conditie paren.

Kernmechanisme

UFC pakt de uitdaging van heterogene inputformaten aan (bijv. ijle keypoints versus dichte dieptekaarten) en ernstige dataschaarste door middel van twee primaire componenten:

1. Taak-agnostische Feature Interpolatie via Patch-wise Matching:
   In plaats van de specifieke condities direct te coderen in een vaste representatie, maakt UFC gebruik van de analogie tussen een query-conditie ($y^q_\tau$) en een support-set ($S_\tau = \{(x_i, y^i_\tau)\}$).

   • Encoders: De methode gebruikt een bevroren beeldencoder $f$ (die visuele kenmerken extraheert uit support-afbeeldingen) en een leerbare conditie-encoder $g_\tau$ (die zowel query- als support-condities codeert).
   • Matching: Controle-features voor een query-patch worden geconstrueerd door de visuele features van support-patches te interpoleren, gewogen door de gelijkenis tussen de query-conditie en de support-condities.
   • Formule: De controle-feature $I(y^{q,k}_\tau; S_\tau)$ voor de $k$-de query-patch wordt berekend als:
     $$I(y^{q,k}_\tau; S_\tau) = \sum_{i=1}^N \sum_{j=1}^M \sigma(g_\tau(y^{q,k}_\tau), g_\tau(y^{i,j}_\tau)) \cdot f(x^{i,j})$$
     waarbij $\sigma$ een patch-gewijze gelijkenisfunctie is. Dit zorgt ervoor dat de resulterende controle-features zich in een verenigde visuele feature-ruimte bevinden, ongeacht de specifieke taak $\tau$, wat robuuste generalisatie vergemakkelijkt.

2. Parameter-efficiënte Adaptatie:
   Om snel aan nieuwe taken aan te passen zonder te overfitten op de kleine support-set, gebruikt UFC een meta-learning strategie gecombineerd met parameter-efficiënte fine-tuning:

   • Meta-Training: Het model wordt getraind op een episodische meta-training dataset die meerdere taken bevat. In elk episode leert het model een query-conditie en een support-set te mappen naar controle-features.
   • Fine-Tuning: Voor een nieuwe taak worden alleen een kleine set taak-specifieke parameters ($\theta_\tau$) in de conditie-encoder (bijv. bias-termen of LoRA-parameters) en de matching/projectie-lagen gefinetuned. De gedeelde backbone-parameters blijven bevroren.

Architectuur-instantiatie

De auteurs instantiëren UFC binnen het ControlNet framework met behulp van Stable Diffusion (zowel UNet als DiT backbones).

• De conditie-encoder $g_\tau$ en de beeld-encoder $f$ worden geïnitialiseerd als kopieën van de pre-trained UNet encoder.
• Patch-wise matching wordt op meerdere niveaus van de encoder uitgevoerd met behulp van multi-head cross-attention blokken.
• De resulterende controle-features worden geprojecteerd en toegevoegd aan de features van het diffusiemodel via zero-geïnitialiseerde lineaire lagen, vergelijkbaar met standaard ControlNet-injectie.

Belangrijkste Bijdragen

1. Eerste Few-Shot Ruimtelijke Controle Framework: Het artikel stelt de eerste methode voor die specifiek is ontworpen voor few-shot ruimtelijke controle in tekst-naar-afbeelding diffusiemodellen, waardoor het mogelijk is afbeeldingen te genereren met nieuwe condities met minimale geannoteerde data (bijv. 30 voorbeelden).
2. Universeel Adapter Design: Door gebruik te maken van patch-wise matching om taak-agnostische visuele features te interpoleren, verenigt UFC diverse ruimtelijke modaliteiten (randen, diepte, pose, segmentatie) in één enkel controlemechanisme, waarmee het probleem van domein-discrepantie van eerdere methoden wordt overwonnen.
3. Data- en Computerefficiëntie: UFC bereikt concurrerende prestaties met volledig gesuperviseerde baselines met gebruik van slechts 0,1% van de volledige trainingsdata (150 voorbeelden) en aanzienlijk minder parameters voor adaptatie.
4. Architectuur-agnostisch: De methode wordt aangetoond effectief te zijn op zowel UNet als DiT (Diffusion Transformer) backbones, wat de toepasbaarheid op moderne diffusie-architecturen garandeert.

Experimentele Resultaten

De auteurs evalueerden UFC op zes nieuwe ruimtelijke controletaksen: Canny edges, HED edges, Depth, Surface Normal, Human Pose en Densepose. De evaluatie omvatte meta-training op subsets van deze taken en testen op uit de training gehouden taken met few-shot fine-tuning (30 shots).

• Controleerbaarheid: UFC presteerde consequent beter dan bestaande few-shot en training-vrije baselines (bijv. Prompt Diffusion, Ctrl-X, Uni-ControlNet) over alle taken. Zo behaalde UFC op de Canny-taak een SSIM van 0,3239 vergeleken met 0,2120 voor Prompt Diffusion (few-shot) en 0,2901 voor Ctrl-X (training-vrij).
• Beeldkwaliteit: UFC behield een hoge beeldkwaliteit (FID-scores vergelijkbaar met volledig gesuperviseerde baselines) terwijl het nauwkeurige ruimtelijke controle bereikte.
• Few-Shot Prestaties: Met slechts 30 support-voorbeelden bereikte UFC fijnmazige controle die consistent is met de ruimtelijke condities. Wanneer het werd gefinetuned met 150 voorbeelden (0,1% van de volledige data), was de prestatie van UFC op de Normal, Depth en Canny taken vergelijkbaar met volledig gesuperviseerde baselines.
• Ablatie-studies: Het verwijderen van het matching-mechanisme of de fine-tuning stap resulteerde in significante prestatieverliezen, wat de noodzaak bevestigt van zowel de analogie-gebaseerde interpolatie als de parameter-adaptatie voor generalisatie.
• Generalisatie naar 3D: De methode werd ook getest op nieuwe 3D-condities (meshes, wireframes, point clouds), wat de capaciteit demonstreerde om te generaliseren naar ongeziene ruimtelijke modaliteiten buiten de meta-training set.

Betekenis en Claims

Het artikel claimt dat UFC een belangrijke stap is naar de praktische bruikbaarheid en flexibiliteit van ruimtelijke controle in generatieve AI. Door het mogelijk te maken om pre-trained diffusiemodellen met nieuwe ruimtelijke condities te gebruiken met slechts een handvol voorbeelden, elimineert UFC de noodzaak voor kostbare, taak-specifieke trainingspipelines.

De auteurs benadrukken dat hun aanpak de kloof overbrugt tussen de rigiditeit van volledig gesuperviseerde controle-adapters en de beperkte aanpasbaarheid van training-vrije methoden. Zij positioneren UFC als een oplossing die "veelzijdige en data-efficiënte" controle mogelijk maakt, waardoor het haalbaar wordt om ruimtelijke controle te implementeren voor opkomende of niche-taken zonder dat enorme datasets vereist zijn.

Het artikel concludeert met een bescheiden erkenning van beperkingen: het framework is momenteel ontworpen voor ruimtelijke controle-generatie in plaats van uiterlijk-behoudende taken (zoals stijltransfer of inpainting), en vereist nog steeds een kleine hoeveelheid fine-tuning per nieuwe taak, in tegen tegenstelling tot Large Language Models die adapteren via in-context learning alleen. De auteurs beschouwen de ontwikkeling van vergelijkbare in-context capaciteiten voor ruimtelijke controle echter als een veelbelovende richting voor toekomstig onderzoek.