Resumo Técnico: Evidência Mecanística para Representações Preservadas, mas Desalinhadas em FedAvg Non-IID

1. Declaração do Problema

O Federated Averaging (FedAvg) é o procedimento de otimização padrão para Aprendizado Federado (FL), porém sofre uma degradação significativa de desempenho sob dados de clientes não-IID (heterogeneidade estatística), particularmente em desvio de rótulos (label skew). Explicações existentes atribuem essa falha ao client drift (deriva do cliente), inconsistência de objetivos e divergência de pesos, caracterizando o processo de otimização, mas deixando o efeito interno da agregação amplamente não observado.

Uma questão crítica permanece aberta: a degradação na acurácia global reflete a apagamento das representações aprendidas pelos clientes, ou reflete um desalinhamento onde as estruturas internas são preservadas, mas falham em se conectar efetivamente com a via de predição final? Este artigo investiga essa questão usando ferramentas de Interpretabilidade Mecanística (MI) para determinar se o FedAvg non-IID destrói o conhecimento aprendido ou apenas desconecta esse conhecimento do classificador.

2. Metodologia

Os autores analisam modelos de visão treinados em clientes esparsos (SimpleCNN e ResNet10) no CIFAR-10 e Fashion-MNIST. Para garantir a interpretabilidade, os clientes treinam sob um cronograma de esparsidade baseado em magnitude (aneculando para 99,9% de esparsidade), com controles de modelos densos confirmando que os achados não são artefatos de esparsidade. O estudo emprega quatro diagnósticos mecanísticos:

2.1. Descoberta de Circuitos

Usando portões binários diferenciáveis (differentiable binary gating), os autores identificam sub-redes mínimas (circuitos) suficientes para prever classes específicas. Eles medem:

• Estabilidade intra-cliente: Sobreposição de circuitos ao longo das rodadas de comunicação.
• Consistência inter-cliente: Similaridade de circuitos entre diferentes clientes.
• Preservação local-para-global: A Interseção sobre União (IoU) entre o circuito de um cliente e o circuito correspondente no modelo global agregado.
• Intervenções: Suficiência (acurácia quando apenas o circuito está ativo) e Necessidade (queda de acurácia quando o circuito é ablado).

2.2. Probing Linear

Um classificador linear é treinado nas ativações da penúltima camada congelada do modelo global. Isso testa se a informação discriminativa de classe permanece linearmente acessível no backbone, mesmo que o head nativo do modelo falhe em utilizá-la.

2.3. Finetuning Apenas do Head

O backbone é congelado, e apenas o head de classificação final é ajustado (finetuned) em um pequeno subconjunto IID balanceado (200 amostras). Isso isola se a perda de acurácia é localizada na camada de leitura (readout layer) ou no backbone de características.

2.4. Autoencoders Esparsos Universais (USAE)

Um USAE é treinado conjuntamente em ativações de modelos globais IID e non-IID para aprender um dicionário de características compartilhado. Isso testa se os dois modelos organizam suas características em uma base comum e permite a reconstrução cruzada para avaliar a compatibilidade de características.

3. Resultados Principais

3.1. Estabilidade e Divergência de Circuitos

• Estabilidade Local: A dinâmica de treinamento local é estável, independentemente da heterogeneidade dos dados. Os clientes convergem para circuitos altamente estáveis e específicos de classe (IoU $\approx$ 0,90), mesmo sob extrema disparidade ($\alpha=0,05$).
• Divergência Inter-Cliente: À medida que a heterogeneidade aumenta, os clientes aprendem circuitos diferentes para a mesma classe. Sob extrema disparidade, a sobreposição de circuitos inter-cliente cai para IoU $\approx$ 0,45.
• Preservação Parcial: O modelo global não faz a média dos circuitos resultando em ruído. A preservação local-para-global (IoU $\approx$ 0,55) permanece maior do que a sobreposição inter-cliente, indicando que o modelo global herda estrutura parcial dos clientes em vez de convergir para uma solução única.

3.2. Circuitos Funcionais Persistem Apesar de Acurácia Zero

Em configurações extremas de non-IID, o modelo global pode atingir acurácia próxima de zero para classes específicas (ex: 0,0% para "Gato" ou "Veado"). No entanto:

• Suficiência: Circuitos extraídos para essas classes alcançam quase perfeita suficiência (100% de acurácia) quando isolados.
• Necessidade: Ablar esses circuitos colapsa a acurácia para quase zero.
• Conclusão: As vias computacionais para essas classes existem e são causalmente responsáveis pela predição, mas o modelo falha em rotear as entradas através delas durante a inferência padrão.

3.3. Separabilidade Linear e Desalinhamento do Head

• Probes Lineares: Embora a acurácia global caia significativamente sob o desvio (ex: de 68,3% para 43,1%), a acurácia do probe linear permanece alta e estável (caindo apenas para 64,7%). A lacuna entre o probe e a acurácia global aumenta monotonicamente com a heterogeneidade (+21,7 pontos em $\alpha=0,05$).
• Finetuning do Head: O ajuste fino apenas do head de classificação em um pequeno conjunto balanceado recupera uma acurácia substancial (+10,2 pontos em $\alpha=0,05$), embora não restaure totalmente o desempenho IID. Isso confirma que o backbone retém características úteis, mas o head padrão está desalinhado.

3.4. Base de Características Compartilhada (USAE)

• Sobreposição de Conceitos: 98,7% dos conceitos do dicionário de 2.048 do USAE ativam de forma não trivial em ambos os modelos IID e non-IID, indicando uma base de características amplamente compartilhada.
• Transferência Cruzada de Modelos: Quando as ativações non-IID são decodificadas no espaço de características IID, o head IID performa bem. Inversamente, o head non-IID falha mesmo quando alimentado com ativações IID. Isso confirma que as características estão presentes e são compatíveis, mas o mecanismo de leitura (readout) non-IID é o ponto de falha.

4. Significância e Alegações

O artigo argumenta que a degradação do FedAvg non-IID não é totalmente explicada pelo apagamento representacional. Em vez disso, o principal mecanismo é o desalinhamento:

1. Preservação: As estruturas internas aprendidas pelos clientes (circuitos) e as bases de características são amplamente preservadas no modelo global agregado.
2. Desalinhamento: A via de predição final (o head de classificação) torna-se desalinhada com essas representações preservadas. O modelo "conhece" as classes, mas não consegue rotear as entradas para os circuitos corretos ou utilizar as características de forma eficaz.

Implicações:

• Reenquadramento do Problema: A questão não é necessariamente que os clientes "esquecem" ou que os pesos divergem para bacias incompatíveis que destroem o conhecimento, mas que o processo de agregação interrompe a relação estrutural entre o backbone e o head.
• Direções Algorítmicas:
  • Realinhamento de Head Pós-Agregação: O ajuste fino simples do head em um pequeno conjunto balanceado pode recuperar um desempenho significativo sem alterar o protocolo federado.
  • Agregação Consciente de Circuitos: Métodos futuros poderiam alinhar neurônios que implementam circuitos funcionais similares entre clientes antes da média (ex: via métodos de simetria de permutação) para reduzir o deslocamento estrutural na origem.

Limitações:
Os autores observam que, embora o realinhamento do head recupere um desempenho significativo, uma lacuna residual em relação ao desempenho IID permanece, sugerindo que a geometria do backbone também degrada levemente. Os achados são atualmente limitados a arquiteturas CNN/ResNet em conjuntos de dados de imagem, e a análise opera ao nível de canal, em vez de granularidade de sub-canal.