Resumen Técnico: Aprendizaje de Invarianzas Compartidas por Subconjuntos para la Generalización de Dominio mediante Mezcla de Expertos

1. Planteamiento del Problema

La Generalización de Dominio (DG) tiene como objetivo entrenar modelos en uno o más dominios de origen que se generalicen a dominios objetivo no vistos sin acceder a los datos del objetivo durante el entrenamiento. El enfoque predominante en DG se basa en el aprendizaje de representaciones invariantes de dominio, típicamente imponiendo que la distribución condicional $P(Z|Y)$ permanezca consistente en todos los dominios de origen. Los métodos como el alineamiento de características, el Riesgo de Invarianza de Invariante (IRM) y el emparejamiento de gradientes operan bajo el supuesto de que existe una estructura predictiva global compartida de forma única en todos los entornos.

Sin embargo, los autores argumentan que este supuesto suele ser demasiado restrictivo. En escenarios del mundo real, los dominios pueden compartir solo parcialmente sus estructuras predictivas; características específicas pueden ser estables en algunos dominios pero variar en otros. Imponer la invariancia global a través de un número creciente de dominios heterogéneos restringe progresivamente el espacio de representación factible. El análisis teórico y las observaciones empíricas (Figura 1) demuestran que, a medida que aumenta el número de dominios de origen, la información mutua entre la representación aprendida y la etiqueta ($I(Z; Y)$) disminuye monótonamente. En consecuencia, los factores predictivos valiosos que son estables solo dentro de subconjuntos específicos de dominios son descartados, lo que conduce a una degradación del rendimiento de la generalización cuando la heterogeneidad del dominio es alta.

2. Metodología: MESSI

Para abordar las limitaciones de la invariancia global, el artículo propone MESSI (Mixture-of-Experts for Subset-Shared Invariances). La hipótesis central es que la estructura invariante debe modelarse como compartida por subconjuntos en lugar de compartida globalmente. El método emplea una arquitectura de Mezcla de Expertos (MoE) para aprender una descomposición estructurada de las relaciones entre dominios.

2.1 Arquitectura

MESSI utiliza un codificador de columna vertebral compartido (por ejemplo, DeiT) seguido de un enrutador y múltiples redes de expertos:

1. Codificador Compartido: Produce una representación de características $u = b_\phi(x)$.
2. Enrutador: Computa una distribución suave sobre los expertos $\pi(x) = \text{softmax}(g_\theta(u))$, aproximando efectivamente una variable de subconjunto latente $S$.
3. Expertos: Cada experto $m$ produce una representación específica $z^{(m)} = h_m(u)$.
4. Agregación: La representación final es una suma ponderada $z(x) = \sum \pi_m(x) z^{(m)}$, lo que permite la composición adaptativa de componentes invariantes de subconjuntos.

2.2 Objetivos de Entrenamiento

El modelo se entrena para optimizar una función de pérdida compuesta (Ecuación 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}$$

• Alineamiento de Subconjuntos Condicionado al Enrutamiento ($\mathcal{L}_{OT}^{ssi}$): En lugar de alinear todos los pares de dominios, la invariancia se impone selectivamente. Para cada experto $m$, el método alinea las distribuciones de características condicionadas a la clase solo entre pares de dominios $(i, j)$ donde el experto tiene una alta responsabilidad de enrutamiento para una clase específica $c$. El peso de alineamiento $a_{ijc}^{(m)}$ se deriva de las probabilidades de enrutamiento. La discrepancia se mide utilizando Transporte Óptimo (OT) (o una variante de MMD), ponderada por estas masas de enrutamiento. Esto crea una "invarianza inducida por el enrutamiento" donde el modelo aprende dónde debe existir la invariancia.
• Especialización de Expertos ($\mathcal{L}_{div}$): Para evitar que los expertos colapsen en representaciones redundantes, un regularizador de diversidad minimiza la similitud de coseno entre las direcciones de las características de diferentes expertos.
• Enrutamiento Confiable y Balanceado ($\mathcal{L}_{sp}, \mathcal{L}_{bal}$):
  • $\mathcal{L}_{sp}$ fomenta el enrutamiento disperso a nivel de instancia (entropía baja), asegurando que cada muestra active un subconjunto específico de expertos.
  • $\mathcal{L}_{bal}$ asegura un uso equilibrado de los expertos a nivel de conjunto de datos, evitando que unos pocos expertos dominen el tráfico.

3. Contribuciones Clave

El artículo describe tres contribuciones principales:

1. Identificación de una Limitación Fundamental: Los autores demuestran que imponer la invariancia a través de un número creciente de dominios reduce la información predictiva. Introducen la invarianza compartida por subconjuntos, replanteando la DG como un problema de aprendizaje de invarianzas estructuradas y dependientes de subconjuntos, en lugar de una única representación invariante global.
2. Marco de MoE Basado en Enrutamiento: Proponen un método que utiliza el enrutamiento para inducir una descomposición de las relaciones de dominio basada en los datos. Al alinear los expertos solo con los subconjuntos de dominios a los que sirven, el método impone la invariancia de forma selectiva, permitiendo la composición adaptativa de componentes invariantes sin requerir etiquetas de subconjunto explícitas.
3. Validación Empírica: El enfoque demuestra una mejora en la generalización fuera del dominio en bancos de pruebas estándar, particularmente bajo cambios de dominio heterogéneos y a medida que aumenta el número de dominios de origen.

4. Resultados Experimentales

Los experimentos se realizaron en bancos de pruebas de DG estándar (DomainBed: PACS, OfficeHome, TerraIncognita, DomainNet) y en un conjunto de datos controlado de Rotated-Colored MNIST diseñado para probar el crecimiento de dominios.

• Bancos de Pruebas Estándar: MESSI (específicamente la variante DeiT-S/16 con OT) alcanzó una precisión promedio de estado del arte (66.2%) en DomainBed, superando a bases sólidas incluyendo ERM, CORAL, Fishr, SAGM y otros métodos de MoE como GMoE y OMoE.
• Robustez al Crecimiento de Dominios: En el protocolo de crecimiento de dominios de presupuesto fijo de Rotated-Colored MNIST, MESSI mostró una robustez superior. Mientras que los métodos de alineamiento global (como CORAL) sufrieron caídas de rendimiento a medida que aumentaba el número de dominios de origen (debido a la sobre-restricción), MESSI mantuvo una alta precisión con una caída mínima (caída del 1.86% frente al 2.58% de la mejor variante de CORAL).
• Estudios de Ablación:
  • Eliminar el alineamiento condicionado al subconjunto ($\mathcal{L}_{ssi}$) causó la mayor caída de rendimiento, confirmándolo como el componente crítico.
  • Eliminar la diversidad o los regularizadores de enrutamiento llevó a la redundancia de expertos y a un enrutamiento difuso, reduciendo el rendimiento.
  • Las comparaciones con controles (Global-MoE, Random-Subset-MoE) mostraron que MESSI aprende una estructura de alineamiento distinta en lugar de simplemente dispersar un objetivo global.
• Eficiencia: Aunque la variante basada en OT introduce una sobrecarga de entrenamiento (aprox. 12.39x el tiempo de ERM), la variante MMD reduce esto a ~4.24x. Los costos de inferencia siguen siendo comparables a otros métodos de MoE.

5. Significado y Reivindicaciones

El artículo afirma que la Generalización de Dominio debe ir más allá del paradigma de imponer una única invariancia global. En su lugar, argumenta que la estructura predictiva es inherentemente dependiente de subconjuntos.

• Perspectiva Teórica: El trabajo proporciona una justificación teórica (vía análisis de información mutua) de por qué la invariancia global falla a medida que aumenta la heterogeneidad del dominio: obliga al modelo a descartar factores predictivos que no son universalmente compartidos.
• Implicación Práctica: Al modelar la invariancia a través de estructuras compartidas parcialmente entre subconjuntos de dominios, MESSI logra una mayor robustez. El mecanismo de enrutamiento permite al modelo descubrir estos subconjuntos directamente de los datos sin requerir etiquetas de subconjunto explícitas o agrupaciones de dominios predefinidas.
• Limitaciones: Los autores reconocen modestamente que el método depende de la capacidad del enrutador para especializarse; si los dominios de origen son pocos o la cobertura de clases es desequilibrada, las señales de enrutamiento pueden ser ruidosas. Además, el costo computacional del alineamiento por pares, particularmente con el Transporte Óptimo, es mayor que el de los métodos estándar de alineamiento global.

En conclusión, MESSI ofrece un mecanismo para preservar y componer invarianzas a nivel de subconjunto, sugiriendo que el futuro de la DG reside en la descomposición estructurada y adaptativa en lugar del alineamiento global uniforme.