Federated Learning with Energy-Based Structured Probabilistic Inference

Resumo Técnico: Aprendizado Federado com Inferência Probabilística Estruturada Baseada em Energia

Definição do Problema
O Aprendizado Federado (FL) normalmente depende do Médio Federado (FedAvg), onde o servidor agrega as atualizações dos modelos dos clientes usando pesos proporcionais às contagens de amostras locais. Embora simples, essa heurística torna-se subótima quando os clientes possuem distribuições de dados heterogêneas (não-IID). Em tais configurações, regras de ponderação fixas tratam todas as atualizações escalonadas como igualmente confiáveis, levando ao desvio do cliente (client drift), convergência lenta ou treinamento instável. Soluções existentes, como termos proximais (FedProx), variáveis de controle (SCAFFOLD) ou regras de agregação robustas (ex: média aparada, mediana geométrica), abordam modos de falha específicos, mas geralmente dependem de heurísticas fixas. Elas falham ao não modelar explicitamente tanto a confiabilidade individual do cliente quanto as interações entre os clientes dentro de um framework de inferência unificado.

Metodologia
Os autores propõem reformular o problema de agregação do lado do servidor como uma tarefa de inferência probabilística estruturada baseada em energia usando Campos Aleatórios Condicionais (CRFs). A ideia central é inferir "rótulos de confiabilidade" discretos para cada cliente com base em suas atualizações de modelo locais, em vez de depender apenas das contagens de amostras.

1. Formulação de Grafo:

   • Nós: Clientes participantes em uma rodada de comunicação.
   • Observações ($x$): Vetores de atualização de modelo local (deltas) e contagens de amostras.
   • Variáveis Latentes ($y$): Rótulos de confiabilidade discretos (ex: $\{0, 0.25, 1\}$), onde valores mais altos indicam maior confiança na atualização.

2. Função de Energia:
   O CRF define uma função de energia $E(y|x)$ composta por dois potenciais:

   • Potenciais Unários ($\psi_u$): Medem a confiabilidade individual do cliente. O servidor computa uma atualização de referência robusta (mediana coordenada a coordenada de todos os deltas dos clientes). Ele então deriva três estatísticas normalizadas para cada cliente: concordância de cosseno com a referência, distância relativa da referência e um escore de anomalia de norma de atualização robusta. Elas são combinadas via uma função sigmoide para gerar um escore de confiabilidade, que inicializa o potencial unário.
   • Potenciais Binários ($\psi_p$): Incentivam a consistência entre clientes com atualizações similares. A afinidade entre os clientes $i$ e $j$ é definida pela similaridade de cosseno de suas atualizações, controlada pelo mínimo de seus escores de confiabilidade individuais. Clientes de alta afinidade são penalizados se receberem rótulos de confiabilidade divergentes.

3. Inferência e Agregação:

   • Como a inferência de Máxima Verossimilhança (MAP) exata é intratável para um grafo totalmente conectado, os autores empregam uma aproximação de campo médio (mean-field). Isso atualiza iterativamente a distribuição a posteriori dos rótulos de confiabilidade para cada cliente com base nos escores unários e nos custos esperados dos vizinhos.
   • O peso de agregação final para um cliente é o produto de seu peso base do FedAvg (contagem de amostras) e sua confiabilidade esperada inferida. Este peso refinado é usado para computar o novo modelo global.

Principais Contribuições

• Formulação: O artigo formula a agregação de servidor em FL como um problema de inferência estruturada baseada em energia sobre variáveis de confiabilidade de clientes.
• Instanciação: Implementa um CRF de nível de cliente totalmente conectado onde os potenciais unários são derivados de estatísticas de atualização robustas (baseadas em mediana) e os potenciais binários são derivados de afinidade de atualização.
• Algoritmo: Propõe um algoritmo de agregação de campo médio compatível com loops de treinamento padrão de FedAvg, não exigindo mudanças no treinamento do lado do cliente.
• Avaliação: O método é avaliado contra baselines padrão de FL (FedAvg, FedProx, FedNova) e métodos de agregação robusta (Krum, média aparada, mediana geométrica, RFA).

Resultados Experimentais
Os autores avaliaram a abordagem em MNIST, CIFAR-10 e CIFAR-100 sob condições severas de não-IID (desvio de rótulo de Dirichlet com $\alpha=0.003$ e $\alpha=0.1$).

• Desempenho: A abordagem guiada por CRF (CRF-FedAvg e CRF-FedProx) superou consistentemente os baselines de agregação fixa e robusta em todos os datasets.
• Significância dos Ganhos: Enquanto as melhorias no MNIST foram modestas (já que a maioria dos métodos alcançou alta acurácia), os ganhos foram mais pronunciados no CIFAR-10 e CIFAR-100. O CRF-FedProx alcançou a maior acurácia no CIFAR-10, e o CRF-FedAvg alcançou a maior no CIFAR-100.
• Convergência: As curvas de perda de validação indicaram que a agregação guiada por CRF convergiu de forma mais rápida e estável do que os baselines robustos, particularmente nas rodadas de comunicação finais. Isso sugere que o método preserva efetivamente direções de atualização úteis enquanto suprime as não confiáveis, evitando o platô conservador visto em métodos como a mediana geométrica.

Significância e Alegações
O artigo alega que tratar a agregação como um problema de inferência estruturada oferece uma vantagem distinta sobre heurísticas fixas. Ao modelar explicitamente tanto a confiabilidade da atualização individual quanto a compatibilidade cliente-cliente, o servidor pode atribuir pesos que são adaptativos à geometria específica das atualizações da rodada atual.

• Os autores enfatizam que esta abordagem é complementar às técnicas de otimização do lado do cliente (como FedProx); enquanto o FedProx restringe o desvio local, a agregação guiada por CRF refina como as atualizações são combinadas pós-treinamento.
• O método é apresentado como um plug-in do lado do servidor que não altera o treinamento ou os protocolos de comunicação do cliente.
• Os autores observam limitações, incluindo o uso de potenciais desenhados manualmente (em vez de aprendidos), hiperparâmetros fixos entre experimentos e o custo computacional do grafo totalmente conectado ($O(m^2)$), que pode exigir esparsificação para federações de escala muito grande. Eles também declaram que a avaliação atual foca em heterogeneidade estatística, e não em robustez adversária (Bizantina).