Class-Grouped Normalized Momentum and Faster Hyperparameter Exploration to Tackle Class Imbalance in Federated Learning

Resumo Técnico: FedCGNM e FedHOO para Desequilíbrio de Classes em Aprendizado Federado

Declaração do Problema
O desequilíbrio de classes permanece um desafio crítico no Aprendizado Federado (FL), particularmente quando a distribuição global de rótulos é de cauda longa (long-tailed). Nesses cenários, as classes minoritárias sofrem com um desempenho preditivo insatisfatório devido à sub-representação nos dados de treinamento. Técnicas centralizadas padrão, como reponderação de perda (loss reweighting), amostragem avançada ou aumento generativo, são difíceis de implementar em FL. Restrições de privacidade impedem a troca de dados, e geradores sintéticos frequentemente falham em produzir amostras realistas para regimes esparsos ou tarefas específicas de domínio (por exemplo, detecção de defeitos em semicondutores). Além disso, abordagens existentes de nível de otimização, como a Normalização por Classe (Per-Class Normalization - PCN), enfrentam limitações em configurações multi-classe: elas são suscetíveis ao ruído direcional de classes minoritárias e introduzem desalinhamentos de escala ao somar muitos vetores unitários, o que desestabiliza a convergência conforme o número de classes aumenta.

Metodologia
O artigo propõe uma abordagem de duas frentes para abordar esses problemas: um otimizador do lado do cliente chamado FedCGNM (Federated Class-Grouped Normalized Momentum) e um sintonizador de hiperparâmetros auxiliar chamado FedHOO.

1. FedCGNM (Otimizador):

   • Agrupamento de Classes: Em vez de normalizar gradientes para cada classe individual ($C$), o FedCGNM particiona as classes em um pequeno número de grupos ($H \ll C$, tipicamente maioria vs. minoria). A regra de agrupamento é derivada pela minimização da variância intra-grupo das proporções de classes, uma estratégia que se mostrou um proxy confiável para a busca exaustiva de limiares.
   • Momento Normalizado por Grupo: Para cada grupo, o algoritmo mantém um vetor de momento. Durante o treinamento local, os gradientes estocásticos são computados por grupo, e o momento é atualizado. Crucialmente, o vetor de momento de cada grupo é normalizado para comprimento unitário (com uma constante de estabilização $\delta$) antes da soma.
   • Direção de Atualização: A direção de atualização final é a soma desses momentas de grupo normalizados. Este design equaliza a magnitude do gradiente entre os grupos de maioria e minoria, ao mesmo tempo em que mitiga o ruído direcional inerente aos gradientes de classes raras. Ao operar sobre um punhado de grupos em vez de classes individuais, ele evita o desalinhamento de escala e a instabilidade associados ao PCN.
   • Alinhamento do Cliente: Ao forçar cada atualização de grupo a ter uma norma unitária, as direções de atualização agregadas entre clientes heterogêneos tornam-se mais alinhadas, abordando uma causa primária de degradação de desempenho em FL.

2. FedHOO (Sintonizador de Hiperparâmetros):

   • Contexto: Selecionar taxas de reamostragem apropriadas não é trivial em FL devido ao espaço de busca combinatório entre clientes e ao risco de desestabilizar a otimização com cronogramas flutuantes.
   • Mecanismo: O FedHOO é um algoritmo baseado em X-armed bandit (XAB) projetado para federações de pequena escala (onde $K$ é pequeno). Ele explora a natureza linear da agregação de FL. Em cada rodada de comunicação, o servidor instrui cada cliente a treinar com duas taxas de reamostragem candidatas.
   • Eficiência: Ao combinar linearmente as atualizações retornadas pelos clientes, o servidor pode inferir o desempenho de validação de todas as $2^K$ combinações possíveis de taxas locais sem a necessidade de $2^K$ rodadas de treinamento separadas. Isso permite a identificação rápida de taxas de amostragem eficazes no início do treinamento, mantendo a privacidade (os clientes retornam apenas métricas de validação escalares).

Análise Teórica
O artigo fornece uma análise de convergência para o FedCGNM que considera explicitamente taxas de reamostragem variantes no tempo. Sob suposições padrão de suavidade (smoothness), gradientes limitados e ausência de viés (unbiasedness), os autores provam que o FedCGNM atinge uma taxa de ponto estacionário padrão de $O(T^{-1/2})$. A análise destaca que, para a convergência ocorrer, a variação cumulativa das taxas de amostragem entre as rodadas ($V_T$) deve satisfazer condições específicas, implicando que reamostragens agressivas, porém instáveis, podem prejudicar a convergência.

Principais Contribuições

1. Otimizador FedCGNM: O primeiro otimizador do lado do cliente que agrupa classes e aplica momento de norma unitária por grupo. Ele equilibra a influência de maioria/minoria, reduz o ruído e mitiga o desalinhamento de escala.
2. Agrupamento Consciente da Variância: Uma estratégia fundamentada e leve que enquadra o particionamento de classes como a minimização da variância intra-grupo, produzindo divisões ótimas comparáveis à busca exaustiva.
3. Análise de Convergência: Um novo framework teórico que combina FL, normalização de grupo e taxas de amostragem dinâmicas, demonstrando o impacto de cronogramas de taxas na convergência.
4. FedHOO: O primeiro algoritmo de otimização para determinar taxas de amostragem locais combinatórias em FL usando um framework de X-armed bandit que aproveita o paralelismo federado para uma exploração eficiente.

Resultos Experimentais
Os métodos foram avaliados em quatro benchmarks públicos de cauda longa (CIFAR-10-LT, CIFAR-100-LT, Adult Income, UNSW-NB15) e um conjunto de dados proprietário de defeitos de chips semicondutores (CDD).

• Desempenho: O FedCGNM superou consistentemente os baselines (incluindo FedAvg, FedProx, CLIMB, FedGraB e variantes de PCN) em todas as configurações. Em regimes de poucos clientes ($K=5$), a combinação do FedCGNM com o FedHOO produziu os melhores resultados, com melhorias de até 29% sobre o FedAvg com reamostragem.
• Aplicação Industrial: No conjunto de dados CDD (taxa de defeito de 1,7%), o método proposto alcançou uma melhoria de 16% sobre o melhor baseline, demonstrando eficácia em ambientes industriais reais e sensíveis à privacidade.
• Robustez: A abordagem mostrou resiliência em regimes de muitos clientes com participação parcial e sob diferentes graus de heterogeneidade de dados (Non-IID).
• Ablação: Experimentos confirmaram que dois grupos ($H=2$) oferecem o melhor equilíbrio entre viés e variância, e que fatores de momento moderados ($\beta \approx 0.5$) produzem o desempenho ideal.

Significância e Alegações
O artigo afirma que o FedCGNM e o FedHOO formam um framework prático para mitigar o desequilíbrio de classes global em ambientes federados que preservam a privacidade, sem depender de troca de dados ou geração sintética. A significância reside em abordar o desequilíbrio de classes no nível de otimização, o que é mais amplamente aplicável do que truques de nível de dados restritos pela privacidade ou características do dataset. Os autores enfatizam que sua abordagem é particularmente eficaz em federações de pequena escala (comuns em implementações industriais), onde o ajuste conjunto de hiperparâmetros é viável e crítico, enquanto a estratégia de agrupamento garante estabilidade e alinhamento em cenários multi-classe onde técnicas de normalização anteriores falham. O trabalho avança o campo ao fornecer tanto um otimizador robusto quanto um método eficientemente fundamentado para ajustar estratégias de reamostragem na presença de desequilíbrio de classes.