Resumo Técnico: Aprendizado de Invariâncias Compartilhadas por Subconjuntos para Generalização de Domínio com Mistura de Especialistas

1. Definição do Problema

A Generalização de Domínio (DG) visa treinar modelos em um ou mais domínios de origem que se generalizem para domínios de destino não vistos sem acessar dados de destino durante o treinamento. A abordagem predominante em DG baseia-se no aprendizado de representações invariantes de domínio, tipicamente ao forçar a distribuição condicional $P(Z|Y)$ a permanecer consistente em todos os domínios de origem. Métodos como alinhamento de características, Invariant Risk Minimization (IRM) e correspondência de gradiente operam sob a premissa de que existe uma única estrutura preditiva globalmente compartilhada entre todos os ambientes.

No entanto, os autores argumentam que essa suposição é frequentemente restritiva demais. Em cenários do mundo real, os domínios podem compartilhar apenas parcialmente suas estruturas preditivas; características específicas podem ser estáveis em alguns domínios, mas variar em outros. Forçar a invariância global através de um número crescente de domínios heterogêneos restringe progressivamente o espaço de representação viável. Análises teóricas e observações empíricas (Figura 1) demonstram que, à medida que o número de domínios de origem cresce, a informação mútua entre a representação aprendida e o rótulo ($I(Z; Y)$) diminui monotonicamente. Consequentemente, fatores preditivos valiosos que são estáveis apenas dentro de subconjuntos específicos de domínios são descartados, levando a um desempenho de generalização degradado quando a heterogeneidade do domínio é alta.

2. Metodologia: MESSI

Para abordar as limitações da invariância global, o artigo propõe o MESSI (Mixture-of-Experts for Subset-Shared Invariances). A hipótese central é que a estrutura invariante deve ser modelada como compartilhada por subconjunto em vez de globalmente compartilhada. O método emprega uma arquitetura de Mistura de Especialistas (MoE) para aprender uma decomposição estruturada das relações entre domínios.

2.1 Arquitetura

O MESSI utiliza um codificador de base compartilhado (ex: DeiT) seguido por um roteador e múltiplas redes de especialistas:

1. Codificador Compartilhado: Produz uma representação de característica $u = b_\phi(x)$.
2. Roteador: Computa uma distribuição suave sobre os especialistas $\pi(x) = \text{softmax}(g_\theta(u))$, aproximando efetivamente uma variável latente de subconjunto $S$.
3. Especialistas: Cada especialista $m$ produz uma representação específica $z^{(m)} = h_m(u)$.
4. Agregação: A representação final é uma soma ponderada $z(x) = \sum \pi_m(x) z^{(m)}$, permitindo a composição adaptativa de componentes invariantes de subconjunto.

2.2 Objetivos de Treinamento

O modelo é treinado para otimizar uma função de perda composta (Equação 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}$$

• Alinhamento de Subconjunto Condicionado ao Roteamento ($\mathcal{L}_{OT}^{ssi}$): Em vez de alinhar todos os pares de domínios, a invariância é imposta seletivamente. Para cada especialista $m$, o método alinha as distribuições de características condicionadas à classe apenas entre pares de domínios $(i, j)$ onde o especialista possui alta responsabilidade de roteamento para uma classe específica $c$. O peso de alinhamento $a_{ijc}^{(m)}$ é derivado das probabilidades de roteamento. A discrepância é medida usando Transporte Ótimo (OT) (ou uma variante de MMD), ponderada por essas massas de roteamento. Isso cria uma "invariância induzida pelo roteamento", onde o modelo aprende onde a invariância deve ocorrer.
• Especialização de Especialistas ($\mathcal{L}_{div}$): Para evitar que os especialistas colapsem em representações redundantes, um regularizador de diversidade minimiza a similaridade de cosseno entre as direções de características de diferentes especialistas.
• Roteamento Confiável e Balanceado ($\mathcal{L}_{sp}, \mathcal{L}_{bal}$):
  • $\mathcal{L}_{sp}$ incentiva o roteamento esparso ao nível da instância (baixa entropia), garantindo que cada amostra ative um subconjunto específico de especialistas.
  • $\mathcal{L}_{bal}$ garante a utilização balanceada dos especialistas ao nível do conjunto de dados, evitando que poucos especialistas dominem o tráfego.

3. Principais Contribuições

O artigo delineia três contribuições primordiais:

1. Identificação de uma Limitação Fundamental: Os autores demonstram que forçar a invariância através de um número crescente de domínios reduz a informação preditiva. Eles introduzem a invariância compartilhada por subconjunto, reformulando a DG como um problema de aprendizado de invariâncias estruturadas e dependentes de subconjuntos, em vez de uma única representação invariante global.
2. Estrutura MoE Baseada em Roteamento: Eles propõem um método que utiliza o roteamento para induzir uma decomposição de relações de domínio orientada pelos dados. Ao alinhar especialistas apenas com os subconjuntos de domínios que eles atendem, o método impõe a invariância de forma seletiva, permitindo a composição adaptativa de componentes invariantes sem exigir rótulos explícitos de subconjuntos.
3. Validação Empírica: A abordagem demonstra melhor generalização fora do domínio em benchmarks padrão, particularmente sob mudanças de domínio heterogêneas e conforme o número de domínios de origem aumenta.

4. Resultados Experimentais

Experimentos foram conduzidos em benchmarks padrão de DG (DomainBed: PACS, OfficeHome, TerraIncognita, DomainNet) e em um conjunto de dados controlado de Rotated-Colored MNIST projetado para testar o crescimento de domínios.

• Benchmarks Padrão: O MESSI (especificamente a variante DeiT-S/16 com OT) alcançou precisão média de estado da arte (66,2%) no DomainBed, superando fortes baselines, incluindo ERM, CORAL, Fishr, SAGM e outros métodos baseados em MoE como GMoE e OMoE.
• Robustez ao Crescimento de Domínios: No protocolo de crescimento de domínio de orçamento fixo do Rotated-Colored MNIST, o MESSI mostrou robustez superior. Enquanto métodos de alinhamento global (como o CORAL) sofreram quedas de desempenho à medida que o número de domínios de origem aumentava (devido à sobre-restrição), o MESSI manteve alta precisão com uma queda mínima (queda de 1,86% vs. 2,58% para a melhor variante do CORAL).
• Estudos de Ablação:
  • A remoção do alinhamento condicionado ao subconjunto ($\mathcal{L}_{ssi}$) causou a maior queda de desempenho, confirmando-o como o componente crítico.
  • A remoção de regularizadores de diversidade ou de roteamento levou à redundância de especialistas e ao roteamento difuso, reduzindo o desempenho.
  • Comparações com controles (Global-MoE, Random-Subset-MoE) mostraram que o MESSI aprende uma estrutura de alinhamento distinta, em vez de simplesmente esparsificar um objetivo global.
• Eficiência: Embora a variante baseada em OT introduza um overhead de treinamento (aprox. 12,39x o tempo do ERM), a variante MMD reduz isso para ~4,24x. Os custos de inferência permanecem comparáveis a outros métodos de MoE.

5. Significância e Alegações

O artigo afirma que a Generalização de Domínio deve ir além do paradigma de forçar uma invariância global única. Em vez disso, argumenta que a estrutura preditiva é inerentemente dependente de subconjuntos.

• Insight Teórico: O trabalho fornece uma justificativa teórica (via análise de informação mútua) para o porquê da invariância global falhar conforme a heterogeneidade do domínio aumenta: ela força o modelo a descartar fatores preditivos que não são universalmente compartilhados.
• Implicação Prática: Ao modelar a invariância através de estruturas parcialmente compartilhadas entre subconjuntos de domínios, o MESSI alcança maior robustez. O mecanismo de roteamento permite que o modelo descubra esses subconjuntos diretamente dos dados sem exigir rótulos de subconjuntos explícitos ou agrupamentos de domínios predefinidos.
• Limitações: Os autores reconhecem modestamente que o método depende da capacidade do roteador de se especializar; se os domínios de origem forem poucos ou a cobertura de classes for desequilibrada, os sinais de roteamento podem ser ruidosos. Adicionalmente, o custo computacional do alinhamento por pares, particularmente com Transporte Ótimo, é superior aos métodos de alinhamento global padrão.

Em conclusão, o MESSI oferece um mecanismo para preservar e compor invariâncias de nível de subconjunto, sugerindo que o futuro da DG reside na decomposição estruturada e adaptativa, em vez do alinhamento global uniforme.