Technische Samenvatting: UNIGP – Het temmen van de Diffusion Transformer voor Uniforme Generatie en Perceptie met Behoud van Priors

Probleemstelling

Recente vooruitgang in diffusiemodellen heeft opmerkelijk succes behaald in zowel controleerbare beeldgeneratie (bijv. Text-to-Image gestuurd door diepte of randen) als taken voor dichte voorspelling (bijv. monoculaire diepte- en oppervlaktenormaal-schatting). De huidige benaderingen behandelen deze domeinen echter doorgaans als afzonderlijke entiteiten. Controleerbare generatiemethoden (zoals ControlNet) zijn beperkt tot generatietaken en vereisen vaak enorme datasets voor precieze controle, terwijl methoden voor dichte voorspelling (zoals Marigold of regressiegebaseerde modellen) vaak zijn gefinetuned vanuit generatieve backbones, wat leidt tot "catastrofale vergetelheid" (catastrophic forgetting) van de veelzijdige generatieve priors van het model.

Huidige pogingen tot unificatie (bijv. JointNet, UniCon) slagen er vaak niet in om de complementaire relatie tussen generatie en perceptie volledig uit te buiten. Ze veroorzaken ofwel enorme computationele overhead door volledige backbones te dupliceren, of ze brengen de generatieve capaciteiten in gevaar door de backbone direct te finetunen voor perceptie. De kernuitdaging is om deze heterogene taken te verenigen binnen een enkel framework dat de vooraf getrainde generatieve priors behoudt, terwijl het nauwkeurige dichte voorspelling en controleerbare generatie mogelijk maakt, zonder complexe taakspecifieke architecturen of verliezen (losses) te vereisen.

Methodologie: UNIGP Framework

De auteurs stellen UNIGP voor, een framework gebouwd op de Multi-Modal Diffusion Transformer (MMDiT) architectuur. De methode bestaat uit twee primaire innovaties: een gedisentangleerde architecturale tak en een uniforme trainingsstrategie.

1. Gedisentangleerde Uniforme Generatie en Perceptie (DUGP)

Volgend aan het principe van Occam's scheermes, vermijdt UNIGP het dupliceren van de volledige MMDiT-backbone. In plaats daarvan introduceren zij de DUGP-tak, die alleen de beeldtak ($\mathcal{F}_i$) van de vooraf getrainde MMDiT kopieert.

• Architectuur: De DUGP-tak wordt geïnitialiseerd door de parameters van de beeldtak te kopiëren. Deze is verdeeld in een Condition Patcher, Stacked Control Layers en Stacked Perception Layers.
• Behoud van Priors (Gradiëntisolatie): Een cruciale ontwerpkeuze is het "omzeilen" (bypass) van de additieve verbinding in de perceptielagen. In de controlelagen wordt de output van de DUGP-tak toegevoegd aan de output van de backbone via een met nul geïnitialiseerde lineaire laag ($\mathcal{L}_{zero}$), waardoor de conditie de generatie kan sturen. Echter, in de perceptielagen wordt dit additieve pad onderbroken. Dit zorgt ervoor dat gradiënten van de dichte voorspellingsverlies (dense prediction loss) uitsluitend door de DUGP-tak stromen, waardoor voorkomt dat de vooraf getrainde gewichten van de backbone worden aangepast. Dit kwantelt de generatieve priors effectief af, terwijl de DUGP-tak in staat wordt gesteld perceptie te leren als een niet-invasieve "lezer".
• Multi-Conditie Afhandeling: Voor scenario's die meerdere condities vereisen, wordt de DUGP condition patcher voor elke conditie gedupliceerd, waarbij interacties worden afgehandeld via sequentie-concatenatie om de efficiëntie tijdens inferentie te behouden.

2. Uniforme Dataset en Trainingsstrategie

Om de kloof tussen generatie en perceptie te overbruggen, gebruikt UNIGP een uniforme trainingsstrategie op een heterogene dataset die gefilterde generatiedata combineert met synthetische perceptiedata.

• Binaire Verliesweging: In plaats van complexe verliesbalancering gebruikt het model een eenvoudige binaire routeringsmechanisme. Voor elke sample in een gemengde batch worden de verliesgewichten ($\lambda_g, \lambda_p$) ingesteld op 1 of 0 op basis van de databron. Als de sample uit de generatiedataset komt, is de generatieverlies actief; als de sample uit de perceptiedataset komt, is de perceptieverlies actief.
• Doelstelling: Het netwerk leert simultaan het snelheidveld (velocity field) voor RGB, diepte en normaal-latents te voorspellen, waarbij geoptimaliseerd wordt voor zowel generatieve kwaliteit als geometrische nauwkeurigheid binnen één enkele forward pass.

Belangrijkste Bijdragen

1. UNIGP Framework: Een uniform Diffusion Transformer-model dat controleerbare generatie en dichte voorspelling integreert. Het maakt gebruik van een DUGP-tak die alleen de beeldtak van MMDiT kopieert, waarbij het principe van Occam's scheermes wordt gevolgd om de computationele overhead te minimaliseren en de veelzijdige priors van de backbone te behouden.
2. Uniforme Trainingsstrategie: Een nieuwe dataset en trainingsaanpak die heterogene taken gezamenlijk leert met behulp van een eenvoudig binair verlieswegingsschema, waardoor de noodzaak voor taakspecifieke architecturale ontwerpen of complexe verliesbalancering wordt geëlimineerd.
3. Aangetoonde Synergie: Het artikel levert bewijs dat gezamenlijke training complementaire voordelen biedt: generatieve priors verrijken de perceptuele details (het vastleggen van fijne texturen), terwijl perceptieel leren zorgt voor strikte ruimtelijke uitlijning in generatie.

Experimentele Resultaten

Uit uitgebreide experimenten blijkt dat UNIGP eerdere uniforme benaderingen overtreft en op gelijke hoogte presteert met gespecialiseerde "expert"-methoden.

• Controleerbare Generatie: UNIGP presteert beter dan single-control (ControlNet) en multi-control (UniControlNet) modellen op het gebied van FID en CLIP-scores over diverse condities (diepte, normaal, schets, pose, Canny). Opvallend genoeg bereikt het de laagste RMSE bij diepte-naar-beeld en normaal-naar-beeld taken, wat wijst op een superieure structurele uitlijning.
• Dichte Voorspelling: In zero-shot diepte- en oppervlaktenormaal-schatting bereikt UNIGP prestaties die vergelijkbaar zijn met state-of-the-art regressiemodellen (bijv. DepthAnything V2, DSINE), ondanks dat het getraind is op aanzienlijk minder afbeeldingen (59K vs. 62,6M voor DepthAnything). Het presteert significant beter dan andere generatieve perceptiemethoden (Marigold, GeoWizard) en uniforme modellen (JointNet, UniCon).
• Gezamenlijke Generatie: UNIGP ondersteunt de gelijktijdige generatie van RGB, diepte- en oppervlaktenormalenkaarten vanuit tekstprompts, waarbij het baselines zoals LDM3D en JointNet overtreft op zowel beeldkwaliteit (FID) als geometrische consistentie (GD/GN metrieken).
• Ablatie-studies:
  • Synergie: Het verwijderen van perceptietraining verslechtert de structurele uitlijning van gegenereerde afbeeldingen (hogere RMSE), terwijl het verwijderen van generatietraining resulteert in te gladde dieptekunsten die fijne details missen.
  • Ontwerpkeuzes: Het vervangen van DUGP door een volledige backbone duplicatie (JointNet-stijl) of volledige fine-tuning (Marigold-stijl) leidt tot aanzienlijk minder goede prestaties, wat bevestigt dat de gedisentangleerde gradiëntstroom de sleutel tot succes is.
  • Laagbalans: Een gelijke toewijzing van gestapelde controle- en perceptielagen (12:12) levert de optimale balans tussen generatiecontrole en perceptienauwkeurigheid op.

Betekenis en Claims

Het artikel claimt dat UNIGP erin slaagt de gescheiden velden van diffusie-gebaseerde synthese en dichte geometrische voorspelling te overbruggen. De primaire betekenis ligt in het aantonen dat generatieve priors en perceptuele restricties niet wederzijds uitsluitend zijn, maar elkaar juist versterken.

• Behoud van Priors: Door de gradiëntstroom te isoleren, behoudt UNIGP de hoogwaardige generatieve capaciteiten van de vooraf getrainde MMDiT-backbone, die vaak verloren gaan bij perceptie-specifieke fine-tuning.
• Efficiëntie: De methode bereikt state-of-the-art resultaten met een relatief kleine toename in parameters (van 2B naar 3B) vergeleken met het dupliceren van de volledige backbone, wat het een computationeel efficiënte uniforme oplossing maakt.
• Veelzijdigheid: Het framework ondersteunt een breed scala aan taken binnen één enkel model, inclusief text-to-image, diepte/normaal schatting, en any-condition generatie, wat een meer veelzijdig alternatief biedt voor het onderhouden van meerdere gespecialiseerde modellen.

De auteurs erkennen dat hoewel UNIGP extra computationele overhead introduceert vergeleken met de basisversie (wat de inferentietijd verhoogt van ~5s naar ~9,5s), deze afweging gerechtvaardigd is door de mogelijkheid om meerdere gespecialiseerde architecturen te vervangen door één enkel, zeer veelzijdig uniform model.