Resumo Técnico: Alinhamento de Distribuição para Aprendizado Federado de Passo Único via Transporte Ótimo

1. Declaração do Problema

O artigo aborda os desafios inerentes ao Aprendizado Federado de Passo Único (One-Shot Federated Learning - OSFL), um regime de comunicação onde os clientes interagem com um servidor central exatamente uma vez. Enquanto o Aprendizado Federado (FL) padrão mitiga a heterogeneidade de dados por meio de otimização iterativa cliente-servidor, o OSFL carece dessa capacidade de refinamento.

O problema central identificado é o efeito cumulativo de dois tipos específicos de heterogeneidade estatística:

1. Desvio de Domínio (Domain Shift): Os clientes observam diferentes distribuições de entrada $P_k(x)$ devido a variadas condições de aquisição ou ambientes.
2. Desvio de Rótulo (Label Shift): Os clientes exibem diferentes marginais de rótulos $P_k(y)$, levando a proporções de classes variadas.

Na ausência de correção iterativa, esses desvios induzem representações de características desalinhadas e funções de decisão de classe-posterior inconsistentes entre os clientes. Os métodos de OSFL existentes (que dependem de destilação, síntese de dados ou ensembling) frequentemente assumem representações alinhadas ou abordam esses desvios separadamente. Consequentemente, as discrepâncias na geometria das características e os vieses condicionais de classe permanecem não corrigidos, degradando o desempenho do modelo global agregado.

2. Metodologia: SLOT-Align

Os autores propõem o SLOT-Align (Single-round, Learning-free Optimal Transport Alignment), um framework consciente da geometria projetado para harmonizar as distribuições de características dos clientes antes do treinamento de OSFL downstream. O método opera como uma etapa de pré-processamento e baseia-se em um codificador congelado compartilhado $f_\theta$ (pré-treinado em um conjunto de dados externo) para mapear todas as entradas locais em um espaço latente comum.

O fluxo de trabalho consiste em quatro etapas principais:

A. Extração de Estatísticas de Características Locais

Cada cliente $k$ calcula estatísticas compactas de primeira e segunda ordem das características $z = f_\theta(x)$ extraídas de seus dados locais:

• Média: $\mu_k = \mathbb{E}[f_\theta(x)]$
• Covariância: $\Sigma_k$ (regularizada via encolhimento de Ledoit–Wolf para garantir positividade definida).
  Estas estatísticas servem como um proxy Gaussiano $N(\mu_k, \Sigma_k)$ para a distribuição de características induzida, capturando as discrepâncias dominantes em localização e dispersão causadas pelos desvios de domínio e de rótulo.

B. Construção da Referência Global

O servidor agrega estas estatísticas locais para construir uma distribuição de referência global:

• Média Global: Uma média ponderada das médias dos clientes.
• Covariância Global: Calculada como o baricentro de Bures–Wasserstein das covariâncias dos clientes. Isso é resolvido via uma iteração de ponto fixo para encontrar o minimizador único da soma das distâncias de Bures ao quadrado, equilibrando efetivamente a variabilidade específica do cliente com a consistência global.

C. Mapeamento de Transporte Ótimo Geodésico

Cada cliente constrói um mapa de transporte para alinhar sua aproximação Gaussiana local à referência global.

• Mapa Ótimo: Sob custo quadrático, o mapa de transporte ótimo entre duas Gaussianas é afim e possui uma solução de forma fechada envolvendo as raízes quadradas das matrizes de covariância.
• Interpolação: Para evitar a correção excessiva (já que a verdadeira distribuição de características não é estritamente Gaussiana), o método introduz um parâmetro de interpolação $\tau \in [0, 1]$. A transformação final $T^{(\tau)}_k$ é uma combinação convexa do mapa identidade e do mapa de transporte ótimo completo. Isso cria uma interpolação de deslocamento de Wasserstein, movendo a distribuição do cliente ao longo de uma geodésica de velocidade constante em direção à referência global.

D. Integração

Os clientes aplicam a transformação $z' = T^{(\tau)}_k(z)$ às suas características locais. Estas características alinhadas são então alimentadas em qualquer algoritmo de OSFL existente (ex: FedAvg, FedCGS, FedPFT) sem modificar os procedimentos de otimização downstream.

3. Principais Contribuições

1. Formalização: O artigo formaliza o OSFL como um problema de alinhamento de distribuição sob desvio conjunto de domínio e rótulo, destacando como esses desvios limitam a eficácia das abordagens de passo único.
2. Proposta de Algoritmo: Apresenta o SLOT-Align, um framework baseado em transporte ótimo, livre de aprendizado (learning-free), que harmoniza as distribuições de características dos clientes através de uma única troca de estatísticas compactas. Não requer dados auxiliares, geradores treináveis ou otimização iterativa.
3. Validação Empírica: Experimentos extensos demonstram que o SLOT-Align melhora consistentemente a precisão e a robustez em múltiplos benchmarks (Office-Home, Digits, DomainNet), diversas backbones pré-treinadas (ViT, ResNet) e baselines de estado da arte de OSFL.

4. Resultados Experimentais

Os autores avaliaram o SLOT-Align sob variados graus de desvio de rótulo (controlado pelo parâmetro de concentração de Dirichlet $\alpha$) e diversidade de domínio.

• Ganhos de Desempenho: Sob desvio conjunto de domínio e rótulo ($\alpha=0.1$), o SLOT-Align melhorou consistentemente a acurácia média macro para todos os baselines testados (O-FedAvg, FedCGS, FedPFT). Por exemplo, no dataset Office-Home, ele melhorou o O-FedAvg em +7,86% e o FedPFT em +6,55%.
• Robustez: O método permaneceu eficaz sob desvio de rótulo mais forte ($\alpha=0.05$ e $\alpha=0.01$), mostrando ganhos mesmo em regimes de heterogeneidade severa.
• Ablação sobre a Força de Alinhamento: O parâmetro de interpolação $\tau$ mostrou-se crítico. Valores intermediários (ex: $\tau=0,4$) produziram o melhor equilíbrio, enquanto $\tau=1,0$ (transporte total) por vezes degradou o desempenho, sugerindo que o alinhamento parcial é preferível à correção total.
• Desvio Apenas de Domínio: Em configurações com rótulos equilibrados (apenas desvio de domínio), o SLOT-Align proporcionou ganhos marginais ou desempenho neutro, confirmando que o método visa especificamente a interação complexa entre desvios de domínio e de rótulo.
• Eficiência: O custo de comunicação é mínimo ($O(m + m^2)$ por cliente para média e covariância), e o overhead computacional é insignificante comparado à extração de características e ao treinamento downstream.

5. Significância e Alegações

O artigo afirma que o SLOT-Align oferece uma solução modular e computacionalmente eficiente para um gargalo crítico no OSFL: a incapacidade de corrigir o desalinhamento distributivo em uma única rodada de comunicação.

• Modularidade: Ele não substitui os pipelines de OSFL existentes, mas atua como uma camada de pré-processamento compatível, tornando-o aplicável a uma ampla gama de métodos atuais e futuros que dependem de codificadores congelados.
• Consciência Geométrica: Ao aproveitar a geometria do espaço 2-Wasserstein e os baricentros de Bures–Wasserstein, ele corrige explicitamente as discrepâncias tanto na localização (média) quanto na forma (covariância) das distribuições de características.
• Praticidade: O método evita os riscos de privacidade e o overhead computacional associados à síntese de dados ou destilação iterativa, baseando-se apenas em resumos estatísticos compactos.

Os autores concluem que, embora o OSFL seja limitado por orçamentos estritos de comunicação, o SLOT-Align mitiga efetivamente a degradação de desempenho causada por distribuições de dados heterogêneas, particularmente quando os desvios de domínio e de rótulo interagem para distorcer as estatísticas de características.