Mixture of Experts (MoE) for PFL

The expert sub-models are the core, specialized components of the MoE architecture. In PFL, these are neural network modules (e.g., feed-forward layers) that are trained collaboratively across all federated clients. Each expert learns to handle specific patterns or features within the heterogeneous global data distribution. The key innovation is that these experts are shared globally—their parameters are aggregated on the server via standard federated averaging—while the logic for using them remains local.

• Function: Capture diverse, reusable knowledge from the federated population.
• Property: Parameters are sparse-activated; only a subset is used for any given input.
• Example: In a vision task, one expert might specialize in texture patterns, while another specializes in shape contours.

The gating network is a small, client-specific neural network that acts as a router or selector. For each input, the gating network computes a set of weights (often a sparse, top-k selection) that determines which expert sub-models are activated and how their outputs are combined. This network is not aggregated federated; it remains entirely local to each client and is trained only on the client's private data.

• Function: Dynamically selects the most relevant experts for a client's local data distribution.
• Property: Enables personalization; two clients with different data will have gating networks that activate different expert combinations.
• Implementation: Often a simple network, like a linear layer followed by a softmax or top-k function, to keep computational overhead low.

Sparse activation is the operational principle that makes MoE architectures computationally feasible at scale. Instead of using all experts for every forward pass, the gating network selects only a small, fixed number (k) of top-scoring experts. This is known as top-k routing. The outputs of the selected experts are then combined via a weighted sum based on the gating weights.

• Benefit: Dramatically reduces computation versus a dense model of equivalent parameter count, as only a fraction of parameters are active per inference.
• PFL Advantage: Allows the global model to have a very large capacity (many experts) to capture population diversity, while each client's device only bears the cost of activating a few experts, making it suitable for edge device heterogeneity.

A critical challenge in MoE training is load balancing. Without intervention, the gating network can develop a preference for a few popular experts, leaving others under-trained—a 'rich-get-richer' problem. In federated PFL, this is compounded by non-IID client data.

Common techniques to ensure fair expert utilization include:

• Auxiliary Load Balancing Loss: An added loss term that penalizes uneven routing distributions.
• Random Routing: Introducing stochasticity, like noise, to the gating mechanism during training.
• Capacity Factor: Setting a limit on the number of tokens routed to each expert per batch.

Failure to balance load leads to expert collapse, degrading model performance and the benefits of specialization.

Training an MoE for PFL typically follows a two-phase federated process:

1. Collaborative Expert Training: The parameters of the expert sub-models are learned via federated averaging. Clients compute gradients for the experts activated by their local gating network and send these updates to the server for secure aggregation.
2. Local Gating Network Training: The client-specific gating network parameters are updated only locally using the client's private data. These parameters are never shared or aggregated, forming the core of the personalization.

This decoupling is key: it allows the learning of a rich, shared repository of knowledge (experts) while enabling fully decentralized, data-private personalization via the gating mechanism.

MoE for PFL is architecturally distinct from but complementary to other personalization strategies:

• vs. FedPer/FedRep: Unlike these methods which personalize only the final personalized model head, MoE personalizes the pathway through the model via the gating network, often allowing personalization within intermediate layers.
• vs. Local Fine-Tuning: MoE provides a structured, sparse alternative to fine-tuning all parameters of a global model, which can be computationally heavy and lead to client drift.
• vs. Clustered FL: MoE can be seen as a soft, dynamic clustering where the gating network assigns each input (or client) to a mixture of clusters (experts) in a data-dependent way, rather than hard-assigning each client to a single cluster model.
• Synergy: MoE can be combined with personalized federated distillation, where the global mixture of experts acts as a teacher for smaller local models.

The core benefit of MoE for PFL is enabling specialized expertise without forcing a single global model to fit all data distributions. Each expert sub-model can learn to excel on a specific type of data pattern (e.g., urban vs. rural sensor data, formal vs. informal text). The gating network acts as a router, allowing each client to access only the relevant experts for inference. This architecture prevents catastrophic forgetting, where learning new patterns erases old ones, as experts dedicated to niche distributions are preserved and refined over federated rounds.

MoE architectures achieve conditional computation. During inference on a client device, only the subset of experts activated by the gating network needs to be executed, not the entire model. This sparse activation can lead to:

• Lower inference latency on edge devices.
• Reduced memory footprint during local personalization, as clients may only fine-tune their selected experts.
• More efficient use of on-device compute resources, which is critical for federated edge learning. However, this efficiency depends heavily on the sparsity and accuracy of the gating mechanism.

MoE scales effectively to scenarios with extreme statistical heterogeneity (non-IID data). As the number of distinct client data distributions grows, the system can, in principle, scale by adding more experts rather than continually enlarging a monolithic model. This makes MoE suitable for large-scale, real-world PFL deployments where client populations are diverse. The gating network learns to map clients to experts based on the underlying data patterns, creating a form of automatic client clustering without explicit server-side clustering algorithms.

A primary challenge is the federated training of the gating network. The gating function must learn to assign inputs to experts based on global knowledge, but it is trained on decentralized, non-IID data. Solutions include:

• Training the gating network centrally on a proxy dataset (compromising privacy).
• Using federated training with techniques like FedAvg, which risks the gating network becoming biased toward dominant client groups.
• Communication overhead increases as both expert parameters and gating network updates must be transmitted, though gating networks are typically small.

A critical systems challenge is load imbalance. Without careful design, the gating network may consistently route most clients to a few popular experts, leaving others underutilized ("expert starvation"). This wastes capacity and can degrade model performance. Mitigation strategies require auxiliary load-balancing losses during training (e.g., a loss term that penalizes uneven routing). However, enforcing this balance in a federated setting, without access to centralized batch statistics, adds significant algorithmic and engineering complexity.

Deploying a federated MoE system introduces substantial orchestration complexity beyond standard PFL. The server must manage:

• Multiple expert models with potentially different update frequencies.
• The versioning and synchronization of experts and gating networks.
• Robust aggregation for experts that may only receive updates from a small, shifting subset of clients.
• Failure resilience for scenarios where an expert receives no updates in a round. This complexity demands sophisticated federated learning orchestrators and complicates debugging and performance evaluation.

A client-specific neural network that dynamically selects or combines expert sub-models for local inference. In MoE for PFL, this network is personalized and remains on the client device.

• Function: Takes local data or context as input and outputs a set of weights or a sparse selection of relevant experts.
• Key Property: It is not aggregated on the server; its parameters are learned locally to reflect individual client data distributions.
• Example: A mobile device uses its gating network to activate only 2 out of 10 global experts for a text prediction task, based on the user's unique writing style.

A specialized neural network module within the global MoE architecture, each trained to handle a distinct region of the input data space or a specific type of task. In PFL, experts are learned collaboratively across all clients.

• Training: Experts are global parameters updated via federated aggregation (e.g., FedAvg) of client gradients.
• Sparsity: During inference, a client's gating network activates only a small subset of all experts, making the system computationally efficient.
• Analogy: Think of experts as a panel of specialists (e.g., cardiologist, dermatologist). The gating network is the triage nurse that routes a patient (input data) to the correct specialist.

The mechanism by which, for a given input, only a small, fixed number of experts in a MoE system are activated for computation. This is a core efficiency feature that makes large MoE models feasible.

• Process: The gating network produces a sparse weight vector, where only the top-k weights are non-zero. Only the corresponding expert networks are executed.
• Benefit for PFL: Reduces on-device inference cost, which is critical for edge and mobile deployments. A client only runs the 1-2 experts it needs.
• Contrast with Dense Models: In a standard dense model (e.g., a Transformer), all parameters are used for every input, leading to higher fixed computational cost.

A related PFL paradigm that frames each client's objective as a separate but related task. The goal is to learn a set of models that share knowledge while performing well on individual tasks.

• Relation to MoE: MoE can be viewed as an architectural instantiation of multi-task learning, where each expert implicitly specializes in a sub-task or data modality.
• Key Difference: Multi-task FL often learns a shared feature extractor with task-specific heads. MoE uses a shared set of expert modules with a client-specific gating function to combine them.
• Objective: Both aim to leverage shared structure across clients while accommodating heterogeneity, but through different parameterization strategies.

A PFL-adjacent approach where clients are partitioned into clusters based on data distribution similarity. A separate global model is then learned for each cluster.

• Comparison to MoE: Clustering is a hard assignment of clients to a single model. MoE provides a soft, dynamic assignment of inputs to multiple experts via the gating network.
• MoE as Soft Clustering: An MoE architecture can be seen as performing soft clustering in the function space, where experts represent cluster centers and the gating network computes membership probabilities.
• Use Case: Clustered FL is effective when client distributions fall into a few clear groups (e.g., geographic regions). MoE is more flexible for continuous or overlapping distribution shifts.

A server-side PFL technique where the aggregation rule is customized per client, rather than using a single global average (FedAvg). The server computes a weighted aggregate of client updates tailored for each recipient.

• Connection to MoE: In MoE for PFL, expert aggregation on the server is global, but expert utilization is personalized via the local gating network. Personalized aggregation is an alternative pathway to personalization.
• Mechanism: Weights can be based on client similarity, data quality, or past update history. For example, a client might receive an aggregate model biased towards updates from statistically similar peers.
• Contrast: MoE personalizes via architecture (gating + experts). Personalized aggregation personalizes via the optimization and communication protocol.