Subspace Optimization for Efficient Federated Learning under Heterogeneous Data

1. Formulação do Problema

O artigo aborda o "triplo desafio" que enfrenta o Aprendizado Federado (FL) na era de modelos em grande escala (centenas de milhões a bilhões de parâmetros):

1. Cálculo: As etapas de treinamento local são computacionalmente pesadas.
2. Memória: Armazenar estados completos do modelo e variáveis auxiliares em dispositivos de borda com recursos limitados é proibitivo.
3. Comunicação: Transmitir atualizações de dimensão completa através de largura de banda limitada é ineficiente.

Esses desafios são exacerbados pela heterogeneidade de dados (Não-IID). Em cenários Não-IID, os objetivos locais divergem do objetivo global, causando "deriva do cliente" (client drift).

• Soluções Existentes e Limitações:
  • Correção de Heterogeneidade (ex: SCAFFOLD): Utiliza variáveis de controle para corrigir a deriva, mas exige o armazenamento e a comunicação de estados auxiliares de dimensão completa, anulando os ganhos de eficiência para modelos grandes.
  • Compressão (ex: Quantização): Reduz os bits transmitidos, mas frequentemente exige buffers de erro de dimensão completa, falhando em reduzir a pegada de memória.
  • Métodos de Subespaço/Baixa Riqueza: Reduzem a dimensionalidade para eficiência, mas tipicamente carecem de mecanismos principiais para corrigir a deriva do cliente. Combinações ingênuas frequentemente descartam informações históricas de correção quando o subespaço muda, levando à instabilidade.

A Lacuna: Não existe nenhum algoritmo existente que realize correção nativa de heterogeneidade dentro de um subespaço de baixa dimensão enquanto preserva a robustez de métodos de dimensão completa como o SCAFFOLD.

2. Metodologia: Subespaço-SCAFFOLD (SSF)

Os autores propõem o SSF, um algoritmo que integra a correção de heterogeneidade diretamente em um subespaço de otimização de baixa dimensão.

Mecanismos Principais

1. Projeção de Subespaço:

   • Em cada rodada $t$, uma matriz de projeção ortogonal aleatória compartilhada $P_t \in \mathbb{R}^{r \times d}$ (onde $r \ll d$) é gerada.
   • A otimização (atualizações locais e agregação) ocorre estritamente no subespaço de dimensão $r$.
   • Os gradientes são projetados: $g_{proj} = P_t g$.

2. Variáveis de Controle Projetadas (Correção de Heterogeneidade):

   • O SSF imita o mecanismo do SCAFFOLD, mas opera sobre quantidades projetadas.
   • Os clientes mantêm variáveis de controle locais $c_i$ e o servidor mantém uma variável de controle global $c$.
   • A direção de atualização local é corrigida usando a diferença entre as variáveis de controle locais e globais projetadas: $P_t g - c_{i, proj} + c_{proj}$.

3. Reabastecimento com Preservação de Resíduo (Inovação Chave):

   • Diferentemente dos métodos de subespaço padrão que descartam informações fora do subespaço ativo, o SSF mantém um estado de controle de dimensão completa ($c_i$ e $c$).
   • Regra de Atualização: Quando o subespaço muda, as variáveis de controle são atualizadas como:
     $$c^{t+1} = (I - P_t^\top P_t)c^t + P_t^\top P_t (\text{nova informação de gradiente})$$
   • Isso significa que o complemento ortogonal (resíduo) da variável de controle é preservado inalterado, enquanto apenas o componente do subespaço ativo é atualizado.
   • Reabastecimento (Backfill): Após a atualização do subespaço, o modelo de espaço completo é reconstruído combinando as novas coordenadas projetadas com o resíduo preservado: $x^{t+1} = P_t^\top x^{t+1}_{proj} + x^t_{res}$.

Fluxo do Algoritmo

1. Decomposição: O servidor divide o modelo e as variáveis de controle em partes projetadas e residuais.
2. Etapas Locais: Os clientes reconstroem o modelo completo usando o modelo de subespaço transmitido e o resíduo do servidor. Eles calculam gradientes no espaço completo, projetam-nos e realizam atualizações corrigidas para heterogeneidade no subespaço.
3. Agregação: O servidor agrega os pontos finais projetados no subespaço.
4. Reabastecimento e Atualização: O servidor eleva o modelo de volta ao espaço completo e atualiza as variáveis de controle, preservando os componentes residuais.

3. Contribuições Principais

1. Design do Algoritmo (SSF):

   • O primeiro algoritmo de FL a realizar correção nativa de heterogeneidade em um subespaço de baixa dimensão.
   • Introduz uma atualização estilo reabastecimento que retém componentes residuais das variáveis de controle, garantindo que informações históricas de correção de deriva não sejam perdidas quando o subespaço gira ou muda.

2. Análise Teórica:

   • Prova que, sob suposições padrão de suavidade e variância limitada, o SSF alcança uma taxa de convergência não assintótica de $\tilde{O}(1/T + 1/\sqrt{NKT})$.
   • Essa taxa corresponde à aceleração linear do SCAFFOLD de dimensão completa, demonstrando que as restrições de subespaço não degradam o comportamento de convergência assintótica em relação à deriva do cliente.
   • A análise vale sem assumir independência entre o projetor aleatório e os gradientes estocásticos.

3. Validação Empírica:

   • Demonstrado tanto em um problema de brinquedo de regressão matricial controlado quanto em um benchmark de aprendizado profundo (CIFAR-100 com ResNet-110).
   • Mostrou que o SSF reduz significativamente os custos de memória, comunicação e cálculo, mantendo uma acurácia comparável ao SCAFFOLD de dimensão completa.

4. Resultados Experimentais

Problema de Brinquedo (Regressão Matricial)

• Robustez: Sob níveis variados de heterogeneidade (baixo, médio, alto), o SSF superou consistentemente o FedSub (um método de subespaço de base sem preservação de resíduo) e o FedAvg.
• Estabilidade: O FedSub exibiu instabilidade numérica e divergência (erros NaN) em dimensões de subespaço mais altas ($r=50$) devido à perda de informação ortogonal. O SSF permaneceu estável em todas as dimensões.
• Acurácia: Com uma razão de subespaço de $r/d = 0.2$, o SSF alcançou erros relativos muito próximos ao Full-SCAFFOLD, significativamente melhores que o Full-FedAvg.

Aprendizado Profundo (CIFAR-100 / ResNet-110)

• Classificação: Full-SCAFFOLD > SSF > FedAvg > FedSub.
• Desempenho: O SSF alcançou uma acurácia de teste final de 45,43%, superando o FedAvg (34,35%) e o FedSub (23,17%) por uma grande margem.
• Eficiência: O SSF reduziu com sucesso a sobrecarga de comunicação e memória, preservando a robustez do SCAFFOLD contra a heterogeneidade de dados.

5. Significado e Impacto

• Ponte da Lacuna: O SSF resolve a tensão fundamental entre robustez (lidar com dados Não-IID) e eficiência (lidar com modelos grandes). Ele prova que não são necessários estados auxiliares de dimensão completa para corrigir a deriva do cliente de forma eficaz.
• Escalabilidade: Ao reduzir a pegada de comunicação e memória de $O(d)$ para $O(r)$ (onde $r \ll d$) enquanto mantém garantias de convergência teórica, o SSF torna o Aprendizado Federado de modelos grandes viável em dispositivos de borda com recursos limitados.
• Insight Teórico: O mecanismo de "preservação de resíduo" fornece um novo paradigma para otimização em subespaço, mostrando que manter memória de dimensão completa para variáveis de controle (mesmo que apenas parcialmente acessada) é crucial para a estabilidade em ambientes de subespaço dinâmicos.

Em conclusão, o SSF oferece uma solução prática e teoricamente sólida para implantar Aprendizado Federado em grande escala e heterogêneo, alcançando o "melhor dos dois mundos": a eficiência dos métodos de subespaço e a robustez da correção por variáveis de controle.