🤖 machine learning

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

DeepFusion is a novel federated learning framework that accelerates Mixture-of-Experts (MoE) training by enabling resource-constrained edge devices to train custom local LLMs and distill their heterogeneous knowledge into a global MoE model via a new View-Aligned Attention module, thereby overcoming architecture mismatches and significantly reducing communication costs while maintaining high performance.

Original authors: Songyuan Li, Jia Hu, Ahmed M. Abdelmoniem, Geyong Min, Haojun Huang, Jiwei Huang

Published 2026-02-17

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Songyuan Li, Jia Hu, Ahmed M. Abdelmoniem, Geyong Min, Haojun Huang, Jiwei Huang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you want to build the world's smartest encyclopedia, but you don't have enough books in your own library. You know that millions of people have valuable knowledge in their heads, but they can't share it because their data is private, and their phones or smart devices are too weak to run a massive encyclopedia on their own.

This is the problem DEEPFUSION solves.

Here is a simple breakdown of how it works, using some everyday analogies.

The Problem: The "Too Big to Fit" Dilemma

Think of modern AI models (like the ones powering chatbots) as giant, complex factories.

The MoE (Mixture of Experts) Factory: This is a super-smart factory with hundreds of specialized workers (experts). To train this factory, you need to show it everything—medical records, financial reports, legal documents, etc.
The Edge Devices: These are your smartphones, smart fridges, and cars. They are like small, cozy home workshops.
The Conflict: You want the home workshops to help train the giant factory. But the factory is so huge that it won't even fit inside a home workshop. If you try to force the factory into a small workshop, the workshop collapses (runs out of memory).

Traditional methods tried to shrink the factory just enough to fit in the workshop, but even the "shrunken" version was too heavy for many devices.

The Solution: DEEPFUSION

DEEPFUSION changes the rules. Instead of forcing the small workshops to build the giant factory, it lets them build their own specialized tools and then teaches the giant factory how to use them.

Here is the step-by-step process:

1. The "Specialized Tool" Phase (Local Training)

Every edge device (phone, car, etc.) is different. Some are powerful; some are weak.

What happens: Instead of forcing everyone to build the same thing, each device builds a custom tool that fits its own size and needs.
- A powerful phone might build a slightly larger tool.
- A weak smart speaker builds a tiny, efficient tool.
The Benefit: They train these tools using their own private data (like your personal medical history or your local bank transactions). No data leaves the device. The device just learns to be an expert in its specific area.

2. The "Grouping" Phase (Clustering)

Once the devices are done training, they send their tools to a central "Master Builder" (the Server).

The Problem: The Master Builder has thousands of tools. Some are for medicine, some for finance, some for cooking. It's a mess.
The Fix: The Master Builder sorts these tools into buckets based on what they know. All the "medical tools" go in one bucket, all the "finance tools" go in another.
The "Proxy": Inside each bucket, the Master Builder mixes the tools together to create one Super-Tool that represents the best knowledge of that group.

3. The "Perspective Shift" (The Secret Sauce: VAA)

This is the most clever part.

The Mismatch: The "Super-Tools" (teachers) think in a simple, compact way because they were built for small devices. The "Giant Factory" (student) thinks in a complex, massive way. If you try to teach the Giant Factory using the Super-Tool's instructions directly, the Factory gets confused. It's like trying to teach a PhD student using a kindergarten textbook; the concepts are right, but the way they are explained is too simple.
The Fix (View-Aligned Attention): The researchers invented a special translator called VAA.
- Imagine the Super-Tool is looking at a puzzle through a magnifying glass (focusing on small details).
- The Giant Factory is looking at the same puzzle from high above (seeing the whole picture).
- The VAA translator takes the "magnifying glass view" and teaches the Giant Factory how to see those small details as if it were looking through the glass too. It aligns their "perspectives" so the Factory can actually learn from the Tool.

4. The "Grand Fusion" (Global Model)

Once the Giant Factory has learned from all the Super-Tools, it becomes the ultimate expert.

It now has a "Medical Expert" who knows everything about diseases (learned from thousands of phones).
It has a "Finance Expert" who knows everything about markets (learned from thousands of banking apps).
All this knowledge is combined into one massive, privacy-safe model.

Why is this a Big Deal?

Privacy First: Your data never leaves your device. The device just learns and sends "knowledge" (the weights of the model), not your actual photos or messages.
Inclusivity: You don't need a supercomputer to participate. Even a weak, old device can contribute because it only has to run a small, custom tool.
Efficiency: It saves a massive amount of internet traffic. Instead of sending huge model updates back and forth many times, the devices train once, send their tool up, and the Master Builder does the heavy lifting.
Better Results: Because it gathers knowledge from many different types of devices (heterogeneous), the final model is smarter and more versatile than if it were trained on just one type of computer.

In short: DEEPFUSION is like organizing a global potluck where everyone brings a dish they made in their own kitchen (using their own ingredients). Instead of trying to fit everyone's kitchen into one giant room, you let them cook at home, bring the dish to a central hall, and then a master chef combines all the recipes to create the world's greatest menu.

1. Problem Statement

The paper addresses the critical bottleneck in training Mixture-of-Experts (MoE) based Large Language Models (LLMs), such as Qwen-MoE and DeepSeek-MoE. These models require vast, diverse, and high-quality training data to activate all expert subnetworks effectively. However, public text data is becoming scarce, and privacy concerns prevent the centralization of private data from edge devices.

While Federated Learning (FL) offers a solution by leveraging distributed edge data, traditional FL approaches face three major hurdles when applied to MoE training:

Resource Constraints: Existing federated MoE methods (e.g., FedJETS) require edge devices to host local "expert" models that are pruned versions of the global MoE. These models still retain the complex MoE backbone and gating mechanisms, making them too large and computationally heavy for resource-constrained devices (e.g., embedded IoT, low-end smartphones).
Communication Overhead: Transmitting large model updates (even pruned ones) across many devices creates prohibitive communication costs.
View-Mismatch (Heterogeneity): Traditional Federated Knowledge Distillation (FKD) assumes teacher and student models share similar architectures. In this scenario, the "teachers" are heterogeneous, lightweight on-device LLMs (e.g., TinyLlama, GPT-2), while the "student" is a massive global MoE model. Their differing architectures and latent spaces lead to a view-mismatch, rendering standard logit or feature alignment ineffective.

2. Methodology: DEEPFUSION Framework

DEEPFUSION is a scalable, one-shot federated training framework designed to fuse heterogeneous on-device knowledge into a global MoE model without requiring devices to run the global architecture. The system operates in three main phases:

A. One-Shot Federated Learning Design

Unlike traditional FL which involves multiple rounds of communication, DEEPFUSION employs a one-shot approach:

Local Training: Each edge device independently configures and trains a lightweight LLM (tailored to its hardware and local data) to optimize its specific application.
Single Upload: Once local training is sufficient, the fully trained on-device LLM is uploaded once to the central server. This drastically reduces communication costs compared to iterative FL.

B. Local Knowledge Clustering

To handle scalability and the computational burden of distilling from hundreds of individual devices:

Clustering: The server clusters uploaded on-device LLMs into $K$ local knowledge domains (matching the $K$ experts in the global MoE) based on semantic similarity of their training data (using low-rank feature embeddings).
Proxy Models: For each cluster, a Local-Knowledge Proxy Model is generated by averaging the weights of the clustered devices. This proxy acts as a robust "teacher" representing that specific domain.

C. Cross-Architecture Knowledge Distillation (The Core Innovation)

The paper introduces a novel View-Aligned Attention (VAA) module to solve the view-mismatch problem between the compact Proxy Models (Teachers) and the Global MoE Base Model (Student).

The Problem: Teachers (compact LLMs) build global semantic features in later layers, while Students (large LLMs) learn them early. Direct feature alignment fails.
The Solution (VAA):
1. Feature Extraction: Features from multiple representation stages (blocks) of the Student model are extracted.
2. Attention Mechanism: A multi-head self-attention module processes these features, allowing the Student to dynamically blend representations from different stages.
3. Alignment: This process constructs a predictive perspective for the Student that aligns with the Teacher's perspective.
4. Loss Function: The system minimizes a combined loss: Cross-Entropy (hard labels), Feature Matching (MSE between aligned features), and KL-Divergence (logits).

D. Global MoE Model Tuning

Merging: The $K$ distilled base models are merged into the global MoE. The FFN layers of the base models become the specialized Experts in the MoE.
Fine-Tuning: The global MoE's shared layers (Embedding, Self-Attention, Gating, Output) are fine-tuned on public benchmark data, while the specialized Experts are frozen. This ensures rapid convergence with minimal computational overhead.

3. Key Contributions

First Scalable Federated MoE Framework: DEEPFUSION is the first framework to enable ubiquitous, resource-constrained edge devices to participate in MoE training via knowledge distillation, eliminating the need for devices to host MoE architectures.
View-Aligned Attention (VAA): A novel module that resolves the architectural heterogeneity and view-mismatch problem, enabling effective knowledge transfer from compact, heterogeneous LLMs to a massive global MoE.
One-Shot Efficiency: By utilizing a one-shot upload strategy and local clustering, the framework significantly reduces communication overhead and server-side computational complexity.
Privacy Preservation: All raw data remains on-device; only model weights and lightweight semantic embeddings are transmitted.

4. Experimental Results

The framework was evaluated using industry-level models (Qwen-MoE and DeepSeek-MoE) on real-world Medical (multi-choice QA) and Financial (open-ended QA) datasets.

Performance: DEEPFUSION achieves performance comparable to centralized training (DeepSpeed) and outperforms state-of-the-art federated baselines (FedJETS, FedKMT, OFA-KD).
- Perplexity: Improved token perplexity by up to 5.28% on challenging open-ended QA tasks.
- Accuracy: Consistently achieved the highest or second-highest token accuracy across varying system scales (16 to 128 devices).
Communication Efficiency: Reduced communication costs by up to 71% compared to FedJETS, as it avoids the transmission of large, pruned MoE structures.
Memory Footprint: DEEPFUSION reduces on-device memory requirements by 3.3x to 9.3x compared to FedJETS, enabling participation from low-end devices (e.g., NVIDIA Jetson Nano) that cannot run MoE models.

5. Significance

DEEPFUSION represents a paradigm shift in federated learning for large models. It solves the "data hunger" of MoE models by unlocking the vast, diverse, and private data residing on heterogeneous edge devices without compromising device usability. By introducing the VAA module, it overcomes the fundamental theoretical barrier of distilling knowledge across vastly different model architectures. This approach makes it feasible to train next-generation, knowledge-rich MoE models in a privacy-preserving manner, addressing both the shortage of public training data and the hardware limitations of the edge ecosystem.