DeepFusion: Accelerating MoE Training via Federated Knowledge Distillation from Heterogeneous Edge Devices

Resumo Técnico: DEEPFUSION

1. O Problema

Os Modelos de Linguagem de Grande Escala (LLMs) baseados em Mistura de Especialistas (MoE), como Qwen-MoE e DeepSeek-MoE, estão revolucionando a IA generativa. No entanto, eles enfrentam dois desafios principais:

• Escassez de Dados: Há uma projeção de que os dados públicos de texto de alta qualidade para treinamento se esgotarão entre 2026 e 2032. É necessário acessar dados privados e diversificados de dispositivos de borda.
• Limitações de Recursos na Borda: O Aprendizado Federado (FL) tradicional exige que os dispositivos de borda hospedem uma cópia local do modelo global para treinamento. Para modelos MoE massivos, isso é impraticável para dispositivos com recursos limitados (como smartphones antigos ou dispositivos IoT embarcados), pois mesmo as versões "podadas" (pruned) dos modelos MoE locais ainda são muito pesadas para o hardware disponível.
• Heterogeneidade e Desalinhamento de Visão: Dispositivos de borda possuem arquiteturas e capacidades de hardware diversas. Tentar alinhar diretamente modelos locais (professores) com o modelo global MoE (aluno) falha devido a uma mismatch de perspectiva (view-mismatch). As arquiteturas diferentes resultam em espaços latentes e vieses indutivos distintos, tornando a distilação de conhecimento tradicional (baseada apenas em logits ou alinhamento de características em espaço compartilhado) ineficaz.

2. Metodologia: DEEPFUSION

O DEEPFUSION é um novo framework de treinamento federado de MoE baseado em Distilação de Conhecimento Federada. Diferente das abordagens anteriores, ele não exige que os dispositivos de borda executem modelos MoE. Em vez disso, cada dispositivo treina seu próprio LLM leve e personalizado.

O processo ocorre em três fases principais no servidor central:

• Fase I: Agrupamento de Conhecimento Local (Local Knowledge Clustering)

  • Os dispositivos enviam seus LLMs treinados e embeddings de baixo rank (que representam a semântica dos dados locais sem revelar os dados brutos) para o servidor.
  • O servidor agrupa esses modelos em K domínios de conhecimento locais (usando K-Means) que correspondem aos K especialistas do modelo MoE global.
  • Para cada domínio, um modelo proxy de conhecimento local é gerado pela média dos pesos dos modelos agrupados. Isso reduz a sobrecarga computacional, permitindo que o servidor distile apenas K modelos proxy em vez de N modelos individuais.

• Fase II: Distilação de Conhecimento entre Arquiteturas (Cross-architecture Knowledge Distillation)

  • Este é o núcleo da inovação. O objetivo é transferir conhecimento dos modelos proxy (arquitetura compacta) para os especialistas do MoE global (arquitetura complexa).
  • Módulo VAA (View-Aligned Attention): Para resolver o problema de "mismatch de visão", o DEEPFUSION introduz o módulo VAA. Ele permite que o modelo aluno (base do MoE) integre representações de características de múltiplos estágios do modelo professor.
  • O VAA utiliza um mecanismo de atenção multi-head para alinhar a perspectiva preditiva do aluno com a do professor, mesmo que residam em espaços latentes diferentes. Isso permite que o aluno aprenda não apenas os logits finais, mas as representações de características intermediárias e o raciocínio subjacente.
  • A função de perda combina perda de entropia cruzada (hard labels), perda de alinhamento de características (MSE) e perda de divergência KL (logits).

• Fase III: Ajuste Fino do Modelo MoE Global

  • Os K modelos base do MoE (resultantes da distilação) são mesclados para formar o modelo MoE global.
  • Os parâmetros das camadas FFN (Feed-Forward Networks) de cada modelo base são copiados diretamente para os respectivos especialistas do MoE global.
  • As camadas de Embedding, Self-Attention e Saída são inicializadas pela média dos parâmetros dos modelos base.
  • O modelo global é então ajustado (fine-tuned) apenas nas camadas de Gate e nas camadas globais, mantendo os especialistas (FFN) congelados, o que garante eficiência computacional.

3. Contribuições Chave

1. Primeiro Framework Federado Escalável para MoE: É a primeira solução que permite a participação de dispositivos de borda heterogêneos e com recursos limitados no treinamento de MoE, sem exigir que eles hospedem modelos MoE locais.
2. Módulo VAA (View-Aligned Attention): Uma inovação técnica que resolve o problema de desalinhamento de perspectivas entre arquiteturas heterogêneas, permitindo uma transferência de conhecimento eficaz através de diferentes espaços latentes.
3. Eficiência de Comunicação e Escalabilidade: Adota um design de FL de "um único tiro" (one-shot), onde os modelos são enviados apenas uma vez após o treinamento local, reduzindo drasticamente a sobrecarga de comunicação. O agrupamento de modelos proxy reduz a carga computacional no servidor.

4. Resultados Experimentais

O framework foi avaliado em dois cenários do mundo real usando modelos industriais (Qwen-MoE e DeepSeek-MoE) e conjuntos de dados médicos e financeiros:

• Desempenho: O DEEPFUSION atingiu um desempenho comparável ao treinamento centralizado (usando DeepSpeed como limite superior) e superou as linhas de base federadas existentes (FedJETS, FedKMT, OFA-KD).
• Perplexidade: Melhorou a perplexidade dos tokens em até 5,28% em tarefas de perguntas e respostas abertas (open-ended QA), demonstrando maior capacidade de raciocínio.
• Custo de Comunicação: Reduziu os custos de comunicação em até 71% em comparação com métodos federados MoE tradicionais, devido à eliminação de múltiplas rodadas de troca de modelos grandes.
• Uso de Memória: Reduziu a pegada de memória nos dispositivos de borda em 3,3x a 9,3x em comparação com métodos que exigem modelos MoE locais, permitindo a participação de dispositivos com hardware limitado (ex: NVIDIA Jetson Nano, iPhones).

5. Significado e Impacto

O DEEPFUSION representa um avanço significativo na viabilidade do treinamento de LLMs de ponta na borda. Ao superar as barreiras de hardware e heterogeneidade, ele permite:

• Privacidade Preservada: Acesso a vastos conjuntos de dados privados (médicos, financeiros) sem centralização.
• Democratização da IA: Dispositivos de baixo custo podem contribuir para modelos globais massivos, enriquecendo o conhecimento de especialistas do MoE com dados diversificados.
• Sustentabilidade de Dados: Oferece uma solução prática para a iminente escassez de dados públicos de treinamento, explorando o potencial dos dados gerados localmente em dispositivos IoT e móveis.

Em resumo, o DEEPFUSION transforma a limitação da heterogeneidade de dispositivos em uma vantagem, agregando perspectivas de conhecimento diversas para criar um modelo MoE global mais robusto e capaz.