FedRep

FedRep explicitly separates the neural network into two distinct components:

• Global Representation (Body): The lower, foundational layers of the model (e.g., convolutional feature extractors, transformer encoders) are learned collaboratively across all clients. This shared representation captures universal, transferable patterns from the collective data.
• Local Head (Classifier/Regressor): The final layers of the model are unique to each client and are never shared or aggregated. These heads are trained exclusively on the client's local data to perform the specific prediction task, allowing for fine-grained personalization to the client's unique data distribution.

Training in FedRep follows a strict, two-phase alternating schedule within each communication round:

1. Local Representation Learning Phase: Clients freeze their local heads and perform multiple epochs of training only on the global representation parameters using their local data. The computed updates for the representation are sent to the server for aggregation (e.g., via FedAvg).
2. Local Head Personalization Phase: Clients freeze the newly updated global representation and perform multiple epochs of training only on their unique local head parameters. These head updates are kept entirely private and are never communicated. This alternation ensures the shared representation evolves to be useful for all clients, while each head specializes for its local task.

A core challenge in personalized federated learning is representation collapse, where clients learn overly specialized features that are useless for others. FedRep's architecture and training schedule directly combat this:

• The global representation is trained using data from all participating clients, forcing it to extract features that are generally useful.
• The alternating phases prevent the head from overly biasing the representation during its learning phase. By separating the optimization objectives, the representation is regularized towards a common, transferable feature space, improving generalization for new clients or tasks.

Unlike FedAvg, which produces a single global model, FedRep is designed with a formal personalization objective. Each client i aims to find optimal parameters (θ, w_i) that minimize: F_i(θ, w_i) = E_{(x,y)~D_i}[L(f(x; θ, w_i), y)] Where:

• θ is the shared global representation.
• w_i is the client-specific local head.
• D_i is the client's local data distribution.
• L is the loss function. The algorithm's goal is to find a shared θ* that serves as a good foundation for all clients, while allowing each to find their own optimal w_i*. This formulation is particularly effective under statistical heterogeneity (non-IID data).

FedRep is often compared to FedPer (Federated Learning with Personalization Layers). Key distinctions:

• Training Order: FedPer trains the entire model (base + head) locally each round, then aggregates only the base layers. FedRep uses alternating phases for representation and head.
• Representation Focus: FedRep's explicit alternating schedule is theorized to learn a more robust and general-purpose global representation compared to FedPer, as it dedicates optimization steps solely to the representation without head interference.
• Communication: Both algorithms communicate only base/representation updates. FedRep may require careful tuning of the phase alternation frequency relative to communication rounds.

FedRep excels in scenarios requiring personalized models on decentralized, heterogeneous data:

• Healthcare Diagnostics: A shared representation learns general medical features from all hospitals, while each hospital's local head specializes in its patient population or diagnostic codes.
• Next-Word Prediction on Smartphones: A global representation learns general language structure, while local heads adapt to individual user writing style and frequently used vocabulary.
• Industrial IoT Predictive Maintenance: A shared representation learns vibration/thermal patterns, while local heads on each machine specialize to its unique wear characteristics and operating environment.
• Retail Recommendation: A global representation learns general product embeddings, while local heads at each storefront personalize recommendations based on local customer purchase history.

FedRep is ideal for training diagnostic models across multiple hospitals without sharing sensitive patient data. Hospitals collaboratively learn a shared feature extractor to identify common patterns in medical imagery (e.g., X-rays, MRIs). Each hospital then trains a local prediction head fine-tuned to its specific patient demographics, imaging equipment, and diagnostic protocols. This ensures high global feature quality while preserving site-specific accuracy and absolute data privacy.

Mobile keyboard applications use FedRep to improve language models without uploading personal typing data. All devices learn a universal representation of language structure and common phrases. Simultaneously, each device trains a personalized head that adapts to the user's unique vocabulary, slang, and typing style. This delivers a responsive, personalized typing experience while keeping all personal messages and habits entirely on-device.

In networks of industrial sensors or smart home devices, FedRep handles significant data heterogeneity. Sensors in different locations (e.g., a factory floor vs. a warehouse) learn a shared representation for common signal patterns like vibration or temperature anomalies. Each sensor cluster then personalizes its head to predict failures specific to its local machinery and environment. This enables predictive maintenance models that are both broadly knowledgeable and contextually precise.

Banks can collaborate against fraud without exposing customer transaction histories. Institutions train a global representation to recognize fundamental patterns of fraudulent activity. Each bank then personalizes a local classifier head tuned to its specific customer base, product types, and regional fraud tactics. This collective learning improves detection of novel, cross-institutional fraud schemes while maintaining competitive advantage and regulatory compliance for individual banks.

Streaming services or news platforms can use FedRep to refine recommendation algorithms. A global content encoder is learned collaboratively across all users to understand item features (genre, topics). On each user's device, a local preference head is trained exclusively on that user's implicit feedback (watch time, clicks). This creates highly tailored recommendations without centralizing sensitive viewing history, addressing core privacy concerns in digital advertising.

FedRep's architecture is adaptable to vertical federated learning, where different organizations hold different features for the same entities (e.g., a bank has financial data, a retailer has purchase data). The parties can collaboratively train the representation layers on aligned sample IDs, learning a fused feature space. Each party then maintains private task-specific heads for their own prediction objectives (e.g., credit scoring, product recommendation), enabling collaboration without direct data pooling.

A subfield of federated learning focused on producing models tailored to individual clients' unique data distributions, rather than a single global model. Core approaches include:

• Local Fine-Tuning: Clients further train the received global model on their local data after aggregation.
• Multi-Task Learning: Framing the problem as learning related but distinct tasks for each client.
• Model Mixture: Combining a global model with a local model for predictions.
• Meta-Learning: Learning a model initialization that can be quickly adapted to any client. FedRep falls into the category of learning shared representations with local personalization layers.

A foundational personalized FL algorithm and direct precursor to FedRep. FedPer splits a neural network into:

• Base Layers (Body): Learned collaboratively across all clients. These layers extract general feature representations from the input data.
• Personalization Layers (Head): The final layers of the network (e.g., classifier). These layers are not shared with the server; each client trains and retains its own unique head. Key Difference from FedRep: In FedPer, clients train both the body and head locally in each round. FedRep explicitly decouples this process, learning the shared representation and personalized heads in separate, alternating phases for more stable representation learning.

The defining challenge FedRep is designed to address. This refers to the scenario where data across federated clients is not Independent and Identically Distributed (IID). Manifestations include:

• Label Distribution Skew: Some clients have examples of only a few classes (e.g., a user's phone only has photos of cats and dogs).
• Feature Distribution Skew: The same feature (e.g., handwriting style) varies significantly between clients.
• Quantity Skew: Vast differences in the amount of data per client. This heterogeneity causes client drift, where local models diverge, harming the performance of a single global model. Personalized FL methods like FedRep provide a solution.

The machine learning paradigm of discovering useful features or representations from raw data. In FedRep, this is the "body" of the model. The goal is to learn a shared, lower-dimensional embedding space that is:

• Informative: Captures the essential factors of variation in the data (e.g., object shape, texture).
• Disentangled: Separates these factors into distinct, independent dimensions.
• Transferable: Useful for a variety of downstream tasks (the personalized "heads"). By collaboratively learning this representation, FedRep allows clients to benefit from the broader data landscape while specializing the final decision layer to their local context.

A major convergence problem in standard federated averaging (FedAvg) under non-IID data. It occurs because each client's local stochastic gradient descent (SGD) step is a biased estimate of the global gradient, causing local models to diverge from the global optimum. FedRep's Mitigation: By fixing the shared representation (body) during the personalization phase and only updating local heads, FedRep explicitly prevents the representation from drifting due to local label skew. The representation is updated only in the separate, collaborative phase, where updates are averaged across many clients, reducing bias.

The core optimization strategy employed by FedRep. Instead of jointly optimizing all model parameters, the algorithm alternates between two sub-problems:

1. Representation (R) Learning Phase: The server's global body R is updated. Client heads h_k are fixed. Clients compute gradients for R based on their data and local head, and these gradients are aggregated by the server.
2. Head (h) Personalization Phase: The global body R is fixed and broadcast. Each client k updates only its own personal head h_k using its local data. These heads are never shared. This alternating approach provides theoretical convergence benefits and practical stability by decoupling the learning objectives.