Multi-Task Federated Learning

MTFL is built upon multi-task learning (MTL), a machine learning paradigm where multiple related tasks are learned simultaneously. The core hypothesis is that sharing representations between tasks improves generalization. In the federated context, each client's local data distribution defines a unique but related learning task. The challenge is to leverage the shared structure across all client tasks while preventing negative interference from highly heterogeneous data.

• Key Mechanism: Learning a shared representation or model component that captures common knowledge.
• Objective: Minimize the sum of all clients' empirical risks: ∑_i L_i(θ_shared, θ_i), where θ_i are client-specific parameters.

A standard MTFL architecture decomposes the model into shared global parameters and client-specific (task-specific) parameters. Only the shared parameters are aggregated on the server, while the personalized parameters remain local. This directly addresses statistical heterogeneity (non-IID data) by allowing the model to adapt to each client's unique data distribution.

• Common Pattern: θ_i = {θ_shared, θ_personal_i}. The shared layers (θ_shared) learn a common feature space, while the final layers (θ_personal_i) make client-specific predictions.
• Examples: FedPer keeps base layers global and classification heads local. FedRep learns a global representation encoder and local client-specific heads.

These methods add a regularization term to each client's local loss function to explicitly control the relationship between the local personalized model and the global model. This balances personalization with the benefits of collaboration.

• Mechanism: Local objective: L_local(w_i) + λ * R(w_i, w_global).
• Proximal Regularization (e.g., FedProx): Uses R(w_i, w_global) = ||w_i - w_global||² to limit client drift.
• Moreau Envelopes: A generalized formulation used in algorithms like Ditto, which provides personalized models with robustness guarantees by solving a bi-level optimization problem.

MTFL is naturally aligned with meta-learning, particularly Model-Agnostic Meta-Learning (MAML). The goal is to learn a global model initialization that is explicitly optimized for fast adaptation to new clients (tasks) with only a few gradient steps on local data.

• Process: The server maintains a meta-model. Each federated round involves:
  1. Clients receive the meta-model.
  2. Clients perform several steps of local adaptation.
  3. Clients compute gradients with respect to the pre-adaptation meta-model and send them back.
  4. Server aggregates these gradients to update the meta-model.
• Outcome: The resulting meta-initialization is highly adaptable, a framework known as Personalized Federated Meta-Learning (PFML).

These approaches handle heterogeneity by grouping clients or model components, acknowledging that not all clients should share the same global knowledge.

• Clustered Federated Learning: Clients are partitioned into clusters based on data distribution similarity (e.g., via gradient similarities). A separate global model is then learned for each cluster, providing coarse-grained personalization.
• Mixture of Experts (MoE): A global set of expert sub-networks is learned collaboratively. A client-specific gating network (learned locally) dynamically selects or combines these experts for inference. This allows personalization at the sub-model level without requiring a unique full model per client.

These are advanced architectural strategies for generating personalized models.

• Hypernetworks: A central neural network (the hypernetwork) learns to generate the weights of each client's personalized model based on a client identifier or context vector. The hypernetwork's parameters are trained federatedly.
• Model Interpolation: The final personalized model for client i is a weighted combination: w_i = α * w_global + (1-α) * w_local_i. The interpolation coefficient α can be tuned per client or learned. This provides a simple but effective knob to control the personalization-collaboration trade-off.

Personalized Federated Learning (PFL) is the overarching paradigm where a global model is trained collaboratively across clients, but the final deployed model is adapted to each client's unique data distribution. MTFL is a specific formulation of PFL.

• Core Goal: Produce client-specific models from a shared foundation.
• Contrast with MTFL: While PFL describes the objective, MTFL provides a specific mathematical framework (multi-task learning) to achieve it.

Non-IID (Non-Independent and Identically Distributed) data is the statistical heterogeneity across client datasets that makes standard federated averaging ineffective and motivates PFL/MTFL.

• Challenge: Client data distributions differ in feature space (covariate shift) or label space (label distribution skew).
• MTFL Solution: Treats each client's non-IID distribution as a distinct but related learning task, explicitly modeling the inter-client relationships.

Clustered Federated Learning is a related approach where clients are partitioned into groups (clusters) based on data distribution similarity, and a separate model is learned for each cluster.

• Mechanism: Uses client update similarity or data statistics to form clusters.
• Relation to MTFL: Can be seen as a form of MTFL where tasks are defined per cluster rather than per client, offering a coarser-grained personalization to reduce complexity.

A Mixture of Experts (MoE) architecture for PFL maintains a global set of expert sub-models, with a client-specific gating network that selects or combines experts for local inference.

• How it works: Experts are trained federated; the gating network is personalized.
• MTFL Connection: This implements a multi-task learning prior by assuming client tasks can be solved by a sparse combination of shared, reusable expert functions.

Hypothesis Transfer Learning in PFL involves using the global federated model as a source hypothesis, which is then transferred and adapted locally by each client.

• Common Technique: Local fine-tuning with regularization (e.g., L2 penalty) to stay close to the global model.
• MTFL Perspective: This is a simple yet effective MTFL strategy where the relationship between client tasks is enforced via regularization towards a shared hypothesis, preventing destructive forgetting of useful common knowledge.

Personalized Model Aggregation refers to server-side algorithms that compute a customized aggregate update for each client, rather than a single global average.

• Examples: Weighting updates based on client similarity or using attention mechanisms.
• MTFL Role: This is a core server-side mechanism in MTFL systems, enabling the server to leverage inter-task relationships to produce a better starting point for each client's local personalization phase.