Sintesi Tecnica: Apprendimento di Invarianze Condivise tra Sottoinsiemi per la Generalizzazione del Dominio con Mixture-of-Experts

1. Definizione del Problema

La Generalizzazione del Dominio (DG) mira ad addestrare modelli su uno o più domini sorgente affinché generalizzino su domini target non visti, senza accedere ai dati target durante l'addestramento. L'approccio prevalente nella DG si basa sull'apprendimento di rappresentazioni invarianti rispetto al dominio, imponendo tipicamente che la distribuzione condizionata $P(Z|Y)$ rimanga coerente in tutti i domini sorgente. Metodi come l'allineamento delle caratteristiche, l'Invariant Risk Minimization (IRM) e il gradient matching operano sotto l'assunto che esista una singola struttura predittiva globalmente condivisa tra tutti gli ambienti.

Tuttavia, gli autori sostengono che questo assunto sia spesso troppo restrittivo. In scenari reali, i domini possono condividere solo parzialmente le loro strutture predittive; specifiche caratteristiche potrebbero essere stabili in alcuni domini ma variare in altri. Imporre l'invarianza globale attraverso un numero crescente di domini eterogenei restringe progressivamente lo spazio delle rappresentazioni fattibili. L'analisi teorica e le osservazioni empiriche (Figura 1) dimostrano che, all'aumentare del numero di domini sorgente, l'informazione mutua tra la rappresentazione appresa e l'etichetta ($I(Z; Y)$) diminuisce monotonicamente. Di conseguenza, fattori predittivi preziosi che sono stabili solo all'interno di specifici sottoinsiemi di domini vengono scartati, portando a una degradazione delle prestazioni di generalizzazione quando l'eterogeneità del dominio è elevata.

2. Metodologia: MESSI

Per affrontare i limiti dell'invarianza globale, il documento propone MESSI (Mixture-of-Experts for Subset-Shared Invariances). L'ipotesi centrale è che la struttura invariante debba essere modellata come condivisa tra sottoinsiemi piuttosto che globalmente condivisa. Il metodo utilizza un'architettura Mixture-of-Experts (MoE) per apprendere una decomposizione strutturata delle relazioni tra i domini.

2.1 Architettura

MESSI utilizza un encoder di base condiviso (ad es., DeiT) seguito da un router e da molteplici reti esperte:

1. Encoder Condiviso: Produce una rappresentazione delle caratteristiche $u = b_\phi(x)$.
2. Router: Calcola una distribuzione "soft" sugli esperti $\pi(x) = \text{softmax}(g_\theta(u))$, approssimando efficacemente una variabile latente di sottoinsieme $S$.
3. Esperti: Ogni esperto $m$ produce una specifica rappresentazione $z^{(m)} = h_m(u)$.
4. Aggregazione: La rappresentazione finale è una somma pesata $z(x) = \sum \pi_m(x) z^{(m)}$, che consente la composizione adattiva dei componenti invarianti del sottoinsieme.

2.2 Obiettivi di Addestramento

Il modello viene addestrato per ottimizzare una funzione di perdita composta (Equazione 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}$$

• Allineamento del Sottoinsieme Condizionato dal Routing ($\mathcal{L}_{OT}^{ssi}$): Invece di allineare tutte le coppie di domini, l'invarianza è imposta selettivamente. Per ogni esperto $m$, il metodo allinea le distribuzioni delle caratteristiche condizionate alla classe solo tra le coppie di domini $(i, j)$ in cui l'esperto ha un'alta responsabilità di routing per una specifica classe $c$. Il peso di allineamento $a_{ijc}^{(m)}$ è derivato dalle probabilità di routing. La discrepanza è misurata utilizzando il Trasporto Ottimo (OT) (o una variante MMD), pesato da queste masse di routing. Ciò crea un'"invarianza indotta dal routing" dove il modello apprende dove l'invarianza debba sussistere.
• Specializzazione dell'Esperto ($\mathcal{L}_{div}$): Per evitare che gli esperti collassino in rappresentazioni ridondanti, un regolarizzatore di diversità minimizza la similarità del coseno tra le direzioni delle caratteristiche di diversi esperti.
• Routing Convincente e Bilanciato ($\mathcal{L}_{sp}, \mathcal{L}_{bal}$):
  • $\mathcal{L}_{sp}$ incoraggia un routing sparso a livello di istanza (bassa entropia), assicurando che ogni campione attivi un sottoinsieme specifico di esperti.
  • $\mathcal{L}_{bal}$ garantisce un utilizzo bilanciato degli esperti a livello di dataset, evitando che pochi esperti dominino il traffico.

3. Contributi Chiave

Il documento delinea tre contributi primari:

1. Identificazione di un Limite Fondamentale: Gli autori dimostrano che imporre l'invarianza attraverso un numero crescente di domini riduce l'informazione predittiva. Introducono l'invarianza condivisa tra sottoinsiemi, riformulando la DG come un problema di apprendimento di invarianze strutturate e dipendenti dai sottoinsiemi, piuttosto che di una singola rappresentazione invariante globale.
2. Framework MoE Basato sul Routing: Propongono un metodo che utilizza il routing per indurre una decomposizione dei dati delle relazioni tra i domini. Allineando gli esperti solo con i sottoinsiemi di domini che servono, il metodo impone l'invarianza in modo selettivo, consentendo la composizione adattiva di componenti invarianti senza richiedere etichette esplicite dei sottoinsiemi.
3. Validazione Empirica: L'approccio dimostra una migliore generalizzazione fuori dal dominio su benchmark standard, in particolare sotto shift di dominio eterogenei e all'aumentare del numero di domini sorgente.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su benchmark standard di DG (DomainBed: PACS, OfficeHome, TerraIncognita, DomainNet) e su un dataset controllato Rotated-Colored MNIST progettato per testare la crescita dei domini.

• Benchmark Standard: MESSI (specificamente la variante DeiT-S/16 con OT) ha raggiunto un'accuratezza media allo stato dell'arte (66,2%) su DomainBed, superando i forti baseline inclusi ERM, CORAL, Fishr, SAGM e altri metodi MoE come GMoE e OMoE.
• Robustezza alla Crescita dei Domini: Nel protocollo di crescita dei domini a budget fisso Rotated-Colored MNIST, MESSI ha mostrato una robustezza superiore. Mentre i metodi di allineamento globale (come CORAL) hanno subito cali di prestazioni all'aumentare del numero di domini sorgente (a causa del sovra-vincolo), MESSI ha mantenuto un'alta accuratezza con un calo minimo (calo del 1,86% rispetto al 2,58% della migliore variante CORAL).
• Studi di Ablazione:
  • Rimuovere l'allineamento condizionato dal sottoinsieme ($\mathcal{L}_{ssi}$) ha causato il calo di prestazioni più significativo, confermandolo come componente critico.
  • Rimuovere i regolarizzatori di diversità o di routing ha portato a ridondanza degli esperti e a un routing diffuso, riducendo le prestazioni.
  • I confronti con i controlli (Global-MoE, Random-Subset-MoE) hanno mostrato che MESSI apprende una struttura di allineamento distinta piuttosto che limitarsi a rendere sparso un obiettivo globale.
• Efficienza: Sebbene la variante basata su OT introduca un overhead di addestramento (circa 12,39x il tempo di ERM), la variante MMD riduce questo valore a circa 4,24x. I costi di inferenza rimangono comparabili ad altri metodi MoE.

5. Significato e Rivendicazioni

Il documento sostiene che la Generalizzazione del Dominio debba andare oltre il paradigma di imporre un'unica invarianza globale. Al contrario, argomenta che la struttura predittiva è intrinsecamente dipendente dal sottoinsieme.

• Approfondimento Teorico: Il lavoro fornisce una giustificazione teorica (tramite l'analisi dell'informazione mutua) del perché l'invarianza globale fallisca all'aumentare dell'eterogeneità dei domini: essa forza il modello a scartare fattori predittivi che non sono universalmente condivisi.
• Implicazione Pratica: Modellando l'invarianza attraverso strutture parzialmente condivise tra i sottoinsiemi di domini, MESSI ottiene una maggiore robustezza. Il meccanismo di routing permette al modello di scoprire questi sottoinsiemi direttamente dai dati senza richiedere etichette di sottoinsieme esplicite o raggruppamenti di domini predefiniti.
• Limitazioni: Gli autori riconoscono con modestia che il metodo dipende dalla capacità del router di specializzarsi; se i domini sorgente sono pochi o la copertura delle classi è sbilanciata, i segnali di routing potrebbero essere rumorosi. Inoltre, il costo computazionale dell'allineamento a coppie, in particolare con il Trasporto Ottimo, è superiore rispetto ai metodi standard di allineamento globale.

In conclusione, MESSI offre un meccanismo per preservare e comporre le invarianti a livello di sottoinsieme, suggerendo che il futuro della DG risieda in una decomposizione strutturata e adattiva piuttosto che in un allineamento globale uniforme.