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

1. Problématique et Contexte

Les modèles de langage de grande taille (LLM) basés sur l'architecture Mixture-of-Experts (MoE) (comme Qwen-MoE et DeepSeek-MoE) révolutionnent le traitement du langage naturel grâce à leur capacité à activer de manière sparse des sous-réseaux d'experts spécialisés. Cependant, ces modèles nécessitent des volumes massifs et diversifiés de données d'entraînement, ce qui pose deux défis majeurs :

• Pénurie de données publiques : Les données textuelles de haute qualité publiques sont en voie d'épuisement.
• Limites de l'apprentissage fédéré (FL) traditionnel : Bien que le FL permette d'utiliser des données privées d'appareils périphériques (edge devices) tout en préservant la vie privée, les méthodes actuelles (comme FedAvg) exigent que chaque appareil héberge une copie locale du modèle global. Pour les modèles MoE de grande taille, cela est impossible pour les appareils contraints en ressources (smartphones, capteurs IoT) en raison de la taille du modèle et de la complexité de l'architecture MoE (mécanismes de gating, multiples experts).
• Hétérogénéité des appareils : Les approches récentes tentant de résoudre ce problème supposent souvent que tous les appareils peuvent exécuter un "expert local" compact, ce qui reste trop lourd pour les dispositifs à faible puissance (CPU/GPU limités, mémoire restreinte).
• Problème de "View-Mismatch" (Inadéquation des perspectives) : Dans la distillation de connaissances fédérée, les modèles locaux (enseignants) et le modèle global MoE (élève) ont des architectures et des espaces latents différents. Les méthodes traditionnelles d'alignement des logits ou des caractéristiques échouent car les perspectives prédictives ne sont pas alignées.

2. Méthodologie : DEEPFUSION

DEEPFUSION est un cadre d'entraînement fédéré MoE scalable conçu pour permettre aux appareils périphériques hétérogènes et contraints de participer à l'entraînement via une distillation de connaissances fédérée. L'approche repose sur trois piliers principaux :

A. Conception "One-Shot" (En un seul coup)

Contrairement au FL traditionnel qui nécessite de multiples rounds de communication pour échanger des mises à jour de modèles, DEEPFUSION adopte une approche one-shot :

1. Chaque appareil périphérique entraîne indépendamment son propre LLM local (léger et adapté à ses besoins et contraintes matérielles) sur ses données privées.
2. Une fois l'entraînement local terminé, le modèle complet est envoyé une seule fois au serveur central.
3. Le serveur effectue la distillation des connaissances pour mettre à jour le modèle MoE global.

• Avantage : Réduction drastique des coûts de communication et de la charge computationnelle sur les appareils périphériques.

B. Clustering des Connaissances Locales

Pour gérer l'échelle et éviter une surcharge computationnelle sur le serveur lors de la distillation de milliers de modèles :

• Les modèles LLM locaux uploadés sont regroupés en K domaines de connaissances locaux (correspondant aux K experts du modèle MoE global) en utilisant un algorithme de clustering (K-Means) basé sur des embeddings de caractéristiques à faible rang extraits des données d'entraînement (pour préserver la confidentialité).
• Pour chaque domaine, un modèle proxy de connaissance locale est généré par moyennage des poids des modèles locaux du cluster.
• Ce modèle proxy sert d'enseignant pour distiller les connaissances vers un expert spécifique du modèle MoE global.

C. Distillation de Connaissances Cross-Architecture avec VAA

C'est l'innovation centrale pour résoudre le problème de l'inadéquation des perspectives (view-mismatch) entre les architectures légères (enseignants) et le modèle MoE complexe (élève) :

• Module VAA (View-Aligned Attention) : Ce module permet au modèle élève (MoE base) d'intégrer les représentations de caractéristiques multi-étapes du modèle enseignant.
• Mécanisme :
  1. Les caractéristiques du modèle élève à différents stades de représentation sont fragmentées et projetées.
  2. Un mécanisme d'attention multi-têtes fusionne ces caractéristiques pour créer une nouvelle représentation ($F^{S'}$) qui aligne la perspective prédictive de l'élève sur celle de l'enseignant.
  3. Cela permet un alignement efficace des caractéristiques intermédiaires, au-delà de la simple correspondance des logits finaux.
• Fusion Globale : Les modèles de base MoE distillés sont fusionnés pour former le modèle MoE global. Les couches FFN (Feed-Forward Networks) des experts sont figées (héritées des modèles de base), tandis que les couches d'embedding, d'attention et de sortie sont affinées pour capturer les relations globales.

3. Contributions Clés

1. Premier cadre fédéré MoE scalable : DEEPFUSION est la première solution permettant d'exploiter des données d'entraînement diversifiées provenant d'appareils périphériques hétérogènes et contraints pour l'entraînement de modèles MoE, tout en préservant la vie privée.
2. Résolution du problème de "View-Mismatch" : Introduction du module VAA qui aligne les perspectives prédictives entre des architectures hétérogènes (LLM locaux vs MoE global), rendant la distillation de connaissances efficace là où les méthodes traditionnelles échouent.
3. Efficacité et Scalabilité : Le cadre permet une participation large d'appareils variés (du smartphone aux capteurs IoT) sans imposer de contraintes matérielles lourdes, tout en réduisant considérablement les coûts de communication grâce à l'approche one-shot.

4. Résultats Expérimentaux

Les expériences ont été menées sur des modèles industriels (Qwen-MoE et DeepSeek-MoE) avec des jeux de données réels (médical et financier).

• Performance : DEEPFUSION atteint des performances proches de l'entraînement centralisé (DeepSpeed) et surpasse les méthodes de base fédérées (FedJETS, FedKMT, OFA-KD).
  • Réduction de la perplexité des tokens : Jusqu'à 5,28 % d'amélioration sur des tâches de questions-réponses ouvertes (Open-ended QA), un domaine où le raisonnement complexe est crucial.
  • Précision : Performances supérieures en termes de précision des tokens sur des tâches de QCM médicaux.
• Efficacité des Communications : Réduction des coûts de communication jusqu'à 71 % par rapport aux méthodes fédérées MoE existantes, grâce à l'absence de multiples rounds d'échange et à l'utilisation de modèles locaux légers.
• Empreinte Mémoire : DEEPFUSION permet l'exécution sur des appareils à ressources limitées (ex: NVIDIA Jetson Nano, iPhone 16) avec une empreinte mémoire 3,3 à 9,3 fois inférieure à celle requise par les méthodes concurrentes (comme FedJETS) qui nécessitent le déploiement d'experts MoE locaux.

5. Signification et Impact

Ce travail marque une avancée significative dans le domaine de l'IA générative et de l'apprentissage fédéré :

• Démocratisation de l'entraînement MoE : Il rend possible l'entraînement de modèles MoE massifs en utilisant la puissance de calcul distribuée de milliards d'appareils IoT et mobiles, sans nécessiter de centraliser les données sensibles.
• Adaptabilité : En permettant à chaque appareil d'utiliser son propre modèle adapté à ses capacités, le système s'adapte à la réalité hétérogène du déploiement de l'IA en périphérie.
• Qualité des Données : Il offre une solution viable à la pénurie de données publiques de haute qualité en exploitant efficacement les données privées et diversifiées du monde réel (médical, financier).
• Innovation Architecturale : Le module VAA ouvre la voie à de nouvelles méthodes de transfert de connaissances entre architectures profondément différentes, un défi majeur pour l'avenir de l'IA fédérée.

En résumé, DEEPFUSION propose une solution pragmatique et performante pour combiner la puissance des modèles MoE avec la richesse des données distribuées sur les appareils périphériques, tout en surmontant les obstacles techniques liés à l'hétérogénéité et à la confidentialité.