Technische Samenvatting: Equivariante Latente Uitlijning via Flow Matching onder Groepssymmetrieën

Probleemstelling

Het artikel behandelt een kritieke beperking in Equivariant Representation Learning (ERL) voor geometrie-bewuste generatieve modellen en novel view synthesis (NVS). Hoewel ERL ernaar streeft om latente ruimtes te construeren waarin groepshandelingen (bijv. rotaties in $SO(2)$ of $SO(3)$) analytisch en voorspelbaar werken, identificeren de auteurs een fenomeen dat zij latente misalignement (latente uitlijningsfout) noemen.

In de praktijk, zelfs wanneer een encoder $\Phi$ wordt getraind met equivariantie-doelstellingen, wijkt de analytisch getransformeerde latente representatie $\rho(g)\Phi(x)$ (waarbij $\rho(g)$ de bekende groeprepresentatie is) vaak af van de empirisch gecodeerde latente van de getransformeerde afbeelding $\Phi(g \circ x)$. Deze discrepantie ontstaat door:

1. Imperfecte Optimalisatie: De encoder slaagt er niet in om een perfect equivariante mapping te leren vanwege de complexiteit van het gezamenlijk comprimeren van objectinhoud en symmetriestructuren.
2. Aliasing: Standaardarchitecturen (CNN's, Transformers) introduceren aliasing via discrete bemonstering en niet-bandgelimiteerde filters, wat zorgt voor persistente equivariantie-fouten.
3. Multi-modale Doelstellingen: Imperfecte encoders kunnen verschillende inputs naar hetzelfde latente punt mappen, wat een multi-modale doelwitdistributie creëert wanneer groepshandelingen worden toegepast, die deterministische punt-naar-punt mappings niet kunnen oplossen.

Deze misalignement accumuleert onder transformaties, wat de geometrische consistentie en visuele getrouwheid van gesynthetiseerde nieuwe gezichten (novel views) verslechtert.

Methodologie: Residual Latent Flow (RLF)

Om deze misalignement te corrigeren zonder de waardevolle groepstheoretische priors weg te gooien, stellen de auteurs Residual Latent Flow (RLF) voor, een framework gebaseerd op Flow Matching.

Kernconcept

In plaats van de analytisch getransformeerde latente $\rho(g)\Phi(x)$ als de uiteindelijke grondwaarheid te beschouwen, behandelen de auteurs deze als een eerste-orde benadering. Ze leren een continue flow (een tijd-afhankelijke diffeomorfisme) die deze benaderde latente naar de empirisch correcte doel-latente $\Phi(g \circ x)$ transporteert.

Probleemformulering

De taak wordt geformuleerd als een conditioneel distributie-transportprobleem:

• Bron ($z_0$): De analytisch getransformeerde latente, $z_0 = \rho(g)\Phi(x)$.
• Doel ($z_1$): De ware gecodeerde latente van de getransformeerde afbeelding, $z_1 = \Phi(g \circ x)$.
• Restrictie: In tegenstelling tot standaard flow matching dat onafhankelijke distributies transporteert, moet RLF $z_0$ transporteren naar de specifieke gekoppelde $z_1$ afgeleid van hetzelfde object $(x, g)$. Dit vereist het modelleren van de conditionele gezamenlijke distributie $\pi_{0,1}(z_0, z_1 | x, g)$ in plaats van onafhankelijke marginalen.

Trainingsprocedure

De methode volgt een driefasige trainingspipeline:

1. Baseline Autoencoder Training: Een encoder $\Phi$ en decoder $\Psi$ worden getraind met een standaard ERL-verlies (Vergelijking 2) om zowel de equivariantie- als de reconstructiefouten te minimaliseren.
2. RLF Training: De encoder wordt bevroren. Een flow-model wordt getraind om het snelheidveld $v_\theta$ te leren dat $z_0$ naar $z_1$ transporteert. De doelstelling minimaliseert het verschil tussen de geleerde snelheid en het conditionele snelheidveld (lineaire interpolatie tussen $z_0$ en $z_1$):
   $$\mathcal{L}_{RLF}(\theta) = \mathbb{E}_{x, g, t, z_t} [\|v_\theta(z_t, t) - (z_1 - z_0)\|_2^2]$$
   Dit stelt het model in staat om de multi-modale aard van de doelwitdistributie aan te pakken die voortkomt uit imperfecte encoders.
3. Decoder Fine-tuning: De decoder $\Psi$ wordt gefinetuned op de door de flow gecorrigeerde latentes om zich aan te passen aan de verschoven latente distributie, wat zorgt voor hoogwaardige reconstructie.

De architectuur van het flow-model varieert per groep: Transformer-gebaseerd voor $SO(2)$ en U-Net-gebaseerd voor $SO(3)$ (waarbij gebruik wordt gemaakt van de vierkante rasterstructuur van Wigner D-matrix blokken).

Belangrijkste Bijdragen

1. Identificatie van Latente Misalignement: Het artikel definieert en analyseert formeel de discrepantie tussen analytisch geroteerde latentes en empirisch gecodeerde latentes in ERL, en toont aan hoe dit de geometrische consistentie ondermijnt.
2. Residual Latent Flow Framework: De auteurs stellen een nieuw mechanisme voor de latente correctie voor met behulp van flow matching. In tegenstelling tot standaard toepassingen die ruis naar data transporteren, transporteert RLF groepstheoretische benaderingen naar empirische doelwitten, waardoor de groepsstructuur behouden blijft terwijl imperfecties in de encoder worden gecorrigeerd.
3. Empirische Validatie: De methode wordt gevalideerd over meerdere datasets met $SO(2)$ en $SO(3)$ symmetrieën, waarbij een significante reductie in latente fout en verbeteringen in de kwaliteit van novel view synthesis worden aangetoond.

Experimentele Resultaten

De auteurs evalueerden RLF op zes datasets, inclusief synthetische (ABO-Material, ModelNet10-SO(3), ComplexBRDFs) en real-image (RotatedMNIST, SmallNORB) benchmarks.

• Novel View Synthesis (NVS): RLF presteerde consequent beter dan de baselines (inclusief NFT, Spatial-VAE en LGA) in zowel in-distribution als out-of-distribution (OOD) settings.
• Metrieken:
  • Latente Fout: Significante reductie in de $L_2$-afstand tussen voorspelde en grondwaarheid latentes.
  • Hoekfout: Substantiële afname van de afwijking van voorspelde rotatiehoeken ten opzichte van de grondwaarheid, met name voor $SO(3)$-taken.
  • Visuele Getrouwheid: Verbeteringen in PSNR en LPIPS, met kwalitatieve resultaten die een betere preservatie van de achtergrond en fijne details tijdens rotatie laten zien vergeleken met baselines, die vaak corruptie of artefacten vertoonden.
• Generalisatie: De methode toonde robuustheid aan bij variërende rotatiehoeken en onbekende objecten, wat valideert dat de flow-correctie de equivariantie-relatie effectief generaliseert.

Betekenis en Claims

Het artikel claimt dat Residual Latent Flow een praktische oplossing biedt voor het "latente misalignement"-probleem dat bestaande equivariante modellen teistert. Door de groepshandeling als een eerste-orde benadering te behandelen en een residuele correctie te leren, kan de methode:

• Consistentie Herstellen: Het dwingt af dat de latente ruimte de onderliggende groepsymmetrie trouwer volgt dan standaard ERL-training alleen.
• Getrouwheid Verbeteren: Het verbetert direct de kwaliteit van novel view synthesis, vooral in uitdagende scenario's met rotaties buiten het vlak (out-of-plane) waar occlusie en geometrisch redeneren vereist zijn.
• Structuur Preserveren: Het behoudt de interpreteerbaarheid en theoretische garanties van groepen (bijv. Wigner D-matrices) terwijl het corrigeert voor de praktische beperkingen van neurale netwerkarchitecturen.

De auteurs merken op dat hoewel de methode de uitlijning aanzienlijk verbetert, de uiteindelijke beeldkwaliteit begrensd blijft door de capaciteit van de decoder. Zij identificeren het uitbreiden van het framework naar meer algemene transformaties (bijv. $SE(3)$, gearticuleerde beweging) en het ontwikkelen van stabiele end-to-end trainingsstrategieën als belangrijke toekomstige richtingen.