Technische Samenvatting: Leren van Subset-Gedeelde Invarianties voor Domein-generalisatie met Mixture-of-Experts

1. Probleemstelling

Domein-generalisatie (DG) streeft ernaar modellen te trainen op één of meer bron-domeinen die generaliseren naar ongeziene doel-domeinen zonder toegang te hebben tot doel-data tijdens de training. De gangbare aanpak in DG vertrouwt op het leren van domein-invariante representaties, waarbij doorgaans wordt gepoogd de conditionele distributie $P(Z|Y)$ consistent te houden over alle bron-domeinen. Methoden zoals feature alignment, Invariant Risk Minimization (IRM) en gradient matching werken onder de aanname dat er één globaal gedeelde voorspellende structuur bestaat over alle omgevingen.

De auteurs stellen echter dat deze aanname vaak te restrictief is. In realistische scenario's kunnen domeinen slechts gedeeltelijk overlappende voorspellende structuren hebben; specifieke kenmerken kunnen stabiel zijn in sommige domeinen maar variëren in andere. Het afdwingen van globale invariantie over een toenemend aantal heterogene domeinen beperkt de bruikbare representatieruimte progressief. Theoretische analyse en empirische observaties (Figuur 1) laten zien dat naarmate het aantal bron-domeinen groeit, de wederzijdse informatie tussen de geleerde representatie en het label ($I(Z; Y)$) monotoon afneemt. Bijgevolg worden waardevolle voorspellende factoren die slechts stabiel zijn binnen specifieke subsets van domeinen weggegooid, wat leidt tot een verslechterde generalisatieprestatie wanneer de domein-heterogeniteit hoog is.

2. Methodologie: MESSI

Om de beperkingen van globale invariantie aan te pakken, stelt het artikel MESSI voor (Mixture-of-Experts voor Subset-Gedeelde Invariantes). De kernhypothese is dat de invariante structuur gemodelleerd moet worden als subset-gedeeld in plaats van globaal gedeeld. De methode maakt gebruik van een Mixture-of-Experts (MoE) architectuur om een gestructureerde decompositie van domeinrelaties te leren.

2.1 Architectuur

MESSI maakt gebruik van een gedeelde backbone encoder (bijv. DeiT) gevolgd door een router en meerdere expert-netwerken:

1. Gedeelde Encoder: Produceert een feature-representatie $u = b_\phi(x)$.
2. Router: Berekent een zachte distributie over experts $\pi(x) = \text{softmax}(g_\theta(u))$, wat effectief een latente subset-variabele $S$ benadert.
3. Experts: Elke expert $m$ produceert een specifieke representatie $z^{(m)} = h_m(u)$.
4. Aggregatie: De uiteindelijke representatie is een gewogen som $z(x) = \sum \pi_m(x) z^{(m)}$, wat zorgt voor de adaptieve compositie van subset-invariante componenten.

2.2 Trainingsdoelstellingen

Het model wordt getraind om een samengestelde verliesfunctie te optimaliseren (Vergelijking 5):
$$\mathcal{L} = \mathcal{L}_{cls} + \lambda_{ssi}\mathcal{L}_{OT}^{ssi} + \lambda_{sp}\mathcal{L}_{sp} + \lambda_{bal}\mathcal{L}_{bal} + \lambda_{div}\mathcal{L}_{div}$$

• Routing-geconditioneerde Subset Alignment ($\mathcal{L}_{OT}^{ssi}$): In plaats van alle domeinparen uit te lijnen, wordt invariantie selectief afgedwongen. Voor elke expert $m$ lijnt de methode de klasse-conditionele feature-distributies alleen uit tussen domeinparen $(i, j)$ waar de expert een hoge routing-verantwoordelijkheid heeft voor een specifieke klasse $c$. Het alignment-gewicht $a_{ijc}^{(m)}$ wordt afgeleid van de routing-kansen. De discrepantie wordt gemeten met behulp van Optimal Transport (OT) (of een MMD-variant), gewogen door deze routing-massa's. Dit creëert een "routing-geïnduceerde invariantie" waarbij het model leert waar invariantie moet gelden.
• Expert Specialisatie ($\mathcal{L}_{div}$): Om te voorkomen dat experts vervallen in redundante representaties, minimaliseert een diversiteits-regularisator de cosinus-gelijkenis tussen de feature-richtingen van verschillende experts.
• Zelfverzekerde en Gebalanceerde Routing ($\mathcal{L}_{sp}, \mathcal{L}_{bal}$):
  • $\mathcal{L}_{sp}$ stimuleert sparse routing op instantie-niveau (lage entropie), waardoor elke sample een specifieke subset van experts activeert.
  • $\mathcal{L}_{bal}$ zorgt voor een gebalanceerd expert-gebruik op dataset-niveau, om te voorkomen dat enkele experts het verkeer domineren.

3. Belangrijkste Bijdragen

Het artikel schetst drie primaire bijdragen:

1. Identificatie van een Fundamentele Beperking: De auteurs tonen aan dat het afdwingen van invariantie over een toenemend aantal domeinen de voorspellende informatie vermindert. Zij introduceren subset-gedeelde invariantie, waarbij DG wordt geherformuleerd als een probleem van het leren van gestructureerde, subset-afhankelijke invarianties in plaats van een enkele globale invariante representatie.
2. Routing-gebaseerd MoE-framework: Zij stellen een methode voor die routing gebruikt om een data-gestuurde decompositie van domeinrelaties te induceren. Door experts alleen uit te lijnen met de domein-subsets die zij bedienen, dwingt de methode invariantie selectief af, wat de adaptieve compositie van invariante componenten mogelijk maakt zonder dat expliciete subset-labels vereist zijn.
3. Empirische Validatie: De aanpak vertoont verbeterde out-of-domain generalisatie op standaard benchmarks, met name onder heterogene domein-shifts en naarmate het aantal bron-domeinen toeneemt.

4. Experimentele Resultaten

Experimenten werden uitgevoerd op standaard DG-benchmarks (DomainBed: PACS, OfficeHome, TerraIncognita, DomainNet) en een gecontroleerde Rotated-Colored MNIST dataset ontworpen om domeingroei te testen.

• Standaard Benchmarks: MESSI (specifiek de DeiT-S/16 variant met OT) bereikte een state-of-the-art gemiddelde nauwkeurigheid (66,2%) op DomainBed, waarmee het sterke baselines verslaat, waaronder ERM, CORAL, Fishr, SAGM, en andere MoE-gebaseerde methoden zoals GMoE en OMoE.
• Robuustheid tegen Domeingroei: Op het Rotated-Colored MNIST fixed-budget domeingroei-protocol toonde MESSI superieure robuustheid. Terwijl globale alignment-methoden (zoals CORAL) leden onder prestatieverlies naarmate het aantal bron-domeinen toenam (door over-constrictie), behield MESSI een hoge nauwkeurigheid met een minimale daling (1,86% daling versus 2,58% voor de beste CORAL-variant).
• Ablatie-studies:
  • Het verwijderen van de subset-geconditioneerde alignment ($\mathcal{L}_{ssi}$) veroorzaakte de grootste prestatiedaling, wat bevestigt dat dit de kritieke component is.
  • Het verwijderen van diversiteit- of routing-regularisaties leidde tot expert-redundantie en diffuse routing, wat de prestaties verminderde.
  • Vergelijkingen met controles (Global-MoE, Random-Subset-MoE) toonden aan dat MESSI een onderscheidende alignment-structuur leert in plaats van simpelweg een globale doelstelling te versnipperen.
• Efficiëntie: Hoewel de OT-gebaseerde variant een trainings-overhead introduceert (ongeveer 12,39x de tijd van ERM), vermindert de MMD-variant dit tot ~4,24x. De kosten voor inferentie blijven vergelijkbaar met andere MoE-methoden.

5. Betekenis en Claims

Het artikel claimt dat Domein-generalisatie voorbij het paradigma van het afdwingen van een enkele globale invariantie moet gaan. In plaats daarvan betogen zij dat de voorspellende structuur inherent subset-afhankelijk is.

• Theoretisch Inzicht: Het werk biedt een theoretische rechtvaardiging (via wederzijdse informatie-analyse) voor waarom globale invariantie faalt wanneer de domein-heterogeniteit toeneemt: het dwingt het model om voorspellende factoren weg te gooien die niet universeel gedeeld worden.
• Praktische Implicatie: Door invariantie te modelleren via gedeeltelijk gedeelde structuren over domein-subsets, bereikt MESSI een grotere robuustheid. Het routing-mechanisme stelt het model in staat om deze subsets direct uit de data te ontdekken zonder dat expliciete subset-labels of vooraf gedefinieerde domein-groeperingen nodig zijn.
• Beperkingen: De auteurs erkennen bescheiden dat de methode afhankelijk is van het vermogen van de router om te specialiseren; als de bron-domeinen schaars zijn of de klasse-dekking ongebalanceerd is, kunnen de routing-signalen ruis bevatten. Daarnaast zijn de computationele kosten van pairwise alignment, met name met Optimal Transport, hoger dan die van standaard globale alignment-methoden.

Concluderend biedt MESSI een mechanisme om subset-niveau invarianties te bewaren en samen te stellen, wat suggereert dat de toekomst van DG ligt in gestructureerde, adaptieve decompositie in plaats van uniforme globale alignment.