Resumo Técnico: COSMOS

Declaração do Problema

A Aprendizagem Federada (AF) enfrenta desafios significativos em ambientes heterogêneos, onde os clientes diferem não apenas na distribuição de dados (heterogeneidade estatística), mas também na arquitetura do modelo (heterogeneidade arquitetural). Embora métodos de Aprendizagem Federada Personalizada (PFL) tenham sido desenvolvidos para adaptar modelos a clientes individuais, a maioria das abordagens existentes assume compatibilidade estrutural ou exige que os clientes divulguem arquiteturas proprietárias de modelos. Isso cria uma barreira prática, pois os clientes frequentemente desejam utilizar modelos arbitrários ou proprietários sem revelar parâmetros ou gradientes.

As abordagens atuais de AF agnósticas a modelos, como FedMD e COMET, dependem de comunicação no nível de saída (por exemplo, rótulos suaves ou duros) utilizando um conjunto de dados público não rotulado. No entanto, esses métodos sofrem de limitações específicas:

1. Coordenação Passiva do Servidor: Frequentemente tratam o servidor como um coordenador passivo, falhando em aproveitar os recursos computacionais do servidor para personalização ativa.
2. Agrupamento Heurístico: Métodos como o COMET dependem de agrupamento K-means, o que requer conhecimento prévio da quantidade de clusters ($K$) e carece de mecanismos formais para garantias de personalização quando a diversidade dos clientes é alta.
3. Falta de Garantias Teóricas: A literatura existente de PFL agnóstica a modelos geralmente carece de garantias teóricas rigorosas quanto à contração de risco, dependendo frequentemente de suposições simplificadas (por exemplo, modelos lineares, distribuições Gaussianas) ou fornecendo apenas resultados padrão de convergência para estacionariedade.

Metodologia: Framework COSMOS

Os autores propõem o COSMOS (Modelos de Servidor Baseados em Saída Agrupada), um framework de PFL agnóstico a modelos que habilita personalização no lado do servidor usando apenas comunicação de pseudo-rótulos. O framework opera sob a suposição de que os clientes possuem dados rotulados privados e que um conjunto de dados público não rotulado $U$ compartilhado (por exemplo, dados raspados da web ou sintéticos) está disponível.

O COSMOS procede em duas fases distintas:

1. Fase de Pré-treinamento e Agrupamento (PTC)

Esta fase estabelece o mapeamento inicial cliente-servidor e os modelos específicos de cluster:

• Pré-treinamento Local: Os clientes treinam seus modelos locais $f_i$ em seus dados privados $D_i$ e geram pseudo-rótulos no conjunto de dados público $U$.
• Agrupamento Ganancioso: O servidor agrupa os clientes com base na similaridade de seus pseudo-rótulos em $U$. Diferentemente do K-means, o COSMOS emprega um algoritmo de agrupamento ganancioso que não exige que $K$ seja especificado antecipadamente. Seleciona iterativamente um cliente com os vizinhos mais "próximos" (dentro de um limite de distância $B_0$) para formar um cluster até que todos os clientes sejam atribuídos. Este processo é endógeno ao procedimento de treinamento.
• Destilação no Lado do Servidor: Para cada cluster $G_k$, o servidor agrega os pseudo-rótulos de seus clientes membros para criar uma distribuição média de pseudo-rótulos. Em seguida, treina um modelo "professor" dedicado $h_k$ para aquele cluster usando esses rótulos agregados e um termo de regularização específico para garantir suavidade local e robustez.
• Destilação no Lado do Cliente: Os clientes recebem os pseudo-rótulos do modelo professor do cluster atribuído e ajustam finamente seus modelos locais usando esses rótulos.

2. Fase de Ajuste Fino Federado Iterativo (IFFT)

Após a fase PTC, o sistema entra em um loop iterativo:

• Ajuste Fino Local: Os clientes refinam seus modelos em seus dados rotulados privados.
• Destilação do Servidor: Os clientes enviam pseudo-rótulos atualizados ao servidor. O servidor atualiza os modelos professores específicos de cluster com base nos novos pseudo-rótulos agregados.
• Destilação do Cliente: Os clientes refinam ainda mais seus modelos locais usando os pseudo-rótulos atualizados do professor.

Crucialmente, os modelos do servidor são treinados usando um termo de regularização derivado da teoria de aprendizado semi-supervisionado (SSL), que penaliza inconsistências de previsão sob transformações de dados (por exemplo, rotações, traduções) para garantir robustez.

Contribuições Teóricas

O artigo fornece a primeira análise teórica ponta a ponta da contração de risco de personalização para PFL agnóstica a modelos. Ao alavancar ferramentas da teoria de SSL (especificamente conectividade baseada em expansão e erro de pseudo-rótulo limitado), os autores derivam condições suficientes para contração exponencial de risco.

Componentes teóricos chave incluem:

• Suposições: A análise assume pesos de importância limitados entre as distribuições pública e privada, objetivos localmente bem condicionados e uma "margem de confiança" onde os clientes estão confiantes em suas previsões nos dados públicos. Também depende de uma propriedade de expansão da distribuição de dados, garantindo que os pseudo-rótulos carreguem informações não triviais sobre a verdade fundamental.
• Contração de Risco: Os autores provam que, sob essas condições, o risco de personalização para modelos do servidor e modelos do cliente contrai exponencialmente com o número de iterações. Especificamente, o limite de risco diminui a uma taxa determinada pela qualidade do cluster e pela cobertura do conjunto de dados público, eventualmente caindo abaixo do risco do treinamento local isolado.

Resultados Experimentais

Os autores avaliaram o COSMOS em quatro benchmarks de classificação de imagens: EMNIST-balanced, CIFAR-10, CIFAR-100 e Tiny ImageNet. Os experimentos utilizaram uma mistura de arquiteturas de clientes (AlexNet, ResNet, MobileNet, SqueezeNet) e um modelo de servidor VGG16.

• Desempenho: O COSMOS superou consistentemente todas as linhas de base agnósticas a modelos (FedMD, FedCT, COMET) em todos os benchmarks. A lacuna de desempenho aumentou à medida que a complexidade da tarefa crescia (por exemplo, de EMNIST para Tiny ImageNet).
• Heterogeneidade: O COSMOS permaneceu robusto a variados graus de heterogeneidade de dados (controlados pelo parâmetro Dirichlet $\alpha$) e diferentes misturas de arquiteturas de clientes. Mantive uma vantagem de desempenho significativa mesmo quando os clientes utilizavam modelos mais fracos.
• Eficiência de Comunicação: Ao comunicar apenas pseudo-rótulos (duros ou suaves) em vez de parâmetros do modelo ou gradientes, o COSMOS reduziu a largura de banda de comunicação em 1 a 2 ordens de grandeza em comparação com métodos de compartilhamento de parâmetros.
• Ambientes Homogêneos: Mesmo em configurações homogêneas onde todos os clientes compartilham a mesma arquitetura, o COSMOS permaneceu competitivo com métodos de PFL de última geração, demonstrando que seu design agnóstico a modelos não incorre em penalidade de desempenho quando a homogeneidade está presente.

Significado e Alegações

O artigo alega que o COSMOS aborda um gargalo crítico na AF ao habilitar personalização de alto desempenho sem exigir conhecimento estrutural dos modelos dos clientes. Seu significado reside em três áreas principais:

1. Implantação Prática: Remove as barreiras estruturais que limitam a PFL em ecossistemas diversos e do mundo real, onde os clientes utilizam arquiteturas proprietárias ou arbitrárias.
2. Avanço Teórico: Fornece as primeiras garantias gerais de contração de risco para PFL agnóstica a modelos, indo além dos resultados padrão de convergência para provar redução exponencial de risco sob suposições realistas.
3. Escalabilidade: Demonstra que o aprendizado personalizado no lado do servidor usando pseudo-rótulos é um paradigma viável e escalável para ambientes altamente heterogêneos, oferecendo ganhos substanciais de eficiência de comunicação.

Os autores concluem que, embora o trabalho atual se concentre em uma estrutura cliente-servidor, o framework estabelece as bases para futuros sistemas de aprendizado descentralizado, aprimorando ainda mais a relevância prática da aprendizagem federada.