Resumo Técnico: TallyTrain

Declaração do Problema

O aprendizado federado (FL) enfrenta dois gargalos de comunicação ortogonais que limitam a escalabilidade:

1. Tamanho do Modelo: Métodos de média de parâmetros (ex: FedAvg, DiLoCo) exigem a troca de pesos completos do modelo ou pseudo-gradientes, resultando em custos de largura de banda proporcionais ao número de parâmetros ($\Theta(|W|)$). Isso é impraticável para modelos de bilhões de parâmetros em dispositivos de borda (edge devices).
2. Contagem de Classes: Métodos no espaço de funções (ex: FedMD, FedDF) trocam predições de rótulos suaves (logits) sobre um conjunto de sondagem público. A largura de banda escala linearmente com o número de classes de saída ($\Theta(C \cdot |D_{pub}|)$). Para tarefas de vocabulário grande (ex: modelos de linguagem com $C \in [2.048, 50.000]$), a transmissão de vetores de rótulos suaves completos torna-se proibitiva.

Os paradigmas atuais tentam reduzir a comunicação diminuindo a frequência das trocas enquanto mantêm o tamanho da mensagem fixo. O TallyTrain aborda o eixo do tamanho ortogonal, mantendo a comunicação frequente, mas reduzindo drasticamente o tamanho da carga útil por mensagem.

Metodologia

O TallyTrain introduz um primitivo de comunicação baseado em votação de argmax sobre um conjunto de sondagem público compartilhado, em vez de trocar distribuições completas de rótulos suaves.

O Primitivo Central

Em vez de transmitir um vetor de $C$ dimensões de logits, cada par (peer) transmite apenas o índice de sua classe predita top-1 ($\text{arg max } f_n(x)$) para cada amostra no conjunto de sondagem público.

• Largura de Banda: A carga útil é reduzida de 4C bytes (floats de 32 bits) para $\lceil \log_2 C \rceil$ bits (alinhados em bytes: 1 byte para $C \le 256$, 2 bytes para $C \le 65.536$).
• Formação de Consenso: Os pares agregam esses rótulos rígidos (hard labels) em um histograma de votação empírica $\bar{H}(x)$, que serve como o alvo de consenso para a destilação.
• Filtragem de Ruído: Os autores argumentam que a votação majoritária atua como um filtro de ruído. Sob condições não-IID, pares subtreinados podem estar "confiantes, porém errados". A média de rótulos suaves amplifica esse ruído, enquanto a votação de rótulos rígidos o filtra, desde que a maioria dos pares esteja correta (satisfazendo uma condição de Condorcet onde a precisão individual $\bar{p} > 0,5$).

Variantes Operacionais

O protocolo suporta dois eixos operacionais ortogonais:

1. Destilação Pura no Espaço de Funções (Eixo A):
   • Sondas Rotuladas: Usa uma perda híbrida combinando Entropia Cruzada (CE) nos rótulos de verdade fundamental (ground-truth) e divergência KL para o histograma de consenso.
   • Sondas Não-Rotuladas: Usa divergência KL para o consenso com um esquema de decaimento linear ($\lambda_r$) para evitar o "colapso de classe de Condorcet" (deriva para um consenso incorreto autorreforçado) quando os rótulos públicos estão ausentes ou fora da distribuição (out-of-distribution).
2. Variante de Ponte de Largura de Banda (Eixo B):
   • Intercala o canal barato de rótulos rígidos com fusões esparsas de média de parâmetros (FedAvg) ocorrendo a cada $M$ rodadas.
   • Esta variante utiliza o canal de rótulos rígidos para estabilizar os pares entre as fusões, prevenindo a "deriva inter-fusão" observada no FedAvg padrão, enquanto a fusão periódica de parâmetros eleva o modelo em direção ao teto de precisão do espaço de parâmetros.

Fundamentos Teóricos

O artigo fornece três resultados teóricos:

• Contração no Espaço de Funções: Sob suposições padrão de suavidade, os pares contraem em direção ao acordo sobre o conjunto de sondagem público.
• Limite de Condorcet: Se a precisão individual do par exceder 50%, a votação majoritária converge para a classe verdadeira com alta probabilidade.
• Redução de Variância: Para pares com margens top-1 suficientes, a variância do gradiente de destilação de rótulo rígido é limitada e inferior à dos gradientes de rótulo suave, pois o argmax trunca as caudas de alta entropia de modelos subtreinados.

Principais Contribuições

1. Primitivo de Comunicação de Rótulo Rígido: Introdução do TallyTrain, que valida o histograma de votação de predições argmax como uma poderosa distribuição de consenso para destilação.
2. Desempenho com Largura de Banda Reduzida: Demonstração de que o consenso de rótulo rígido iguala ou supera a precisão da destilação de rótulo suave, reduzindo a largura de banda por sonda em fatores de 40 a 4.096 (dependendo de $C$).
3. Regimes Operacionais Duplos:
   • Um modo puramente de espaço de função adequado para arquiteturas heterogêneas.
   • Um modo de "ponte de largura de banda" (TallyTrain+faM) que supera as linhas de base padrão de média de parâmetros (FedAvg, FedProx, FedDF) através da fronteira de Pareto de largura de banda vs. precisão.
4. Análise Teórica: Formalização das propriedades de contração e dos benefícios de redução de variância da votação de rótulo rígido.

Resultados Experimentais

Experimentos foram conduzidos em CIFAR-10, CIFAR-100 (não-IID) e WikiText-2 (modelagem de linguagem com $C=2048$).

• Precisão vs. Largura de Banda:
  • No CIFAR-100 (não-IID), o TallyTrain alcançou 33,16% de precisão de cauda com 400x menos largura de banda que o FedMD (rótulo suave), superando o FedMD em 1,35 pontos percentuais.
  • No CIFAR-10, a variante de ponte TallyTrain+fa200 alcançou 71,92% de precisão, superando significativamente o FedAvg (54,63%) e o FedDF (52,34%) com custos de largura de banda similares ou menores.
  • No WikiText-2, o TallyTrain+fa200 igualou a precisão do FedAvg-fa200 com apenas ~6% de largura de banda extra, enquanto o TallyTrain puro igualou a precisão do FedMD com 4.096x menos largura de banda.
• Filtragem de Ruído: No CIFAR-100, o consenso de rótulo suave na verdade degradou o desempenho (contribuindo com -0,22 pp) em comparação com uma linha de base sem consenso, enquanto a votação de rótulo rígido melhorou o desempenho (+1,13 pp), confirmando a hipótese de filtragem de ruído.
• Estabilidade: A variante de ponte demonstrou o menor desvio padrão entre pares ($\sigma = 0,38$ no CIFAR-10, $\sigma = 0,03$ no WikiText-2), indicando alta reprodutibilidade e estabilidade em comparação ao FedAvg e FedDF.

Significância e Alegações

O artigo afirma que o TallyTrain aborda os gargalos de engenharia do aprendizado federado de duas maneiras:

1. Escalabilidade para Grandes Vocabulários: Ao colapsar o eixo de contagem de classes para $\lceil \log_2 C \rceil$ bits, o TallyTrain permite a destilação federada para tarefas de vocabulário grande (como modelos de linguagem) onde métodos de rótulo suave são atualmente inviáveis devido às restrições de largura de banda.
2. Dominância de Pareto: A variante de ponte de largura de banda (TallyTrain+faM) cria um novo ponto operacional que domina os métodos padrão de média de parâmetros. Ele aproveita a "folga do eixo de frequência" do FedAvg de baixa frequência (onde os pares derivam entre as fusões) preenchendo a lacuna com um canal de rótulo rígido barato e estabilizador.

Os autores enfatizam que o método não amortiza o computacional de treinamento (cada par ainda executa o SGD local completo), mas otimiza estritamente o recurso de comunicação. A abordagem é particularmente eficaz em configurações não-IID, onde a média de rótulo suave tende a amplificar os erros de pares que estão confiantes, porém errados, e subtreinados, enquanto a votação majoritária os filtra.