FedPall: Federated Learning for Feature Drift

• FedPall is a federated learning framework designed for feature drift where label marginals are shared but class-conditional features differ across clients.
• It employs prototype-based adversarial alignment and collaborative contrastive learning to unify feature spaces while preserving class-discriminative structures.
• The framework uses mixed prototype features and a global classifier to enhance model generalization and privacy, achieving significant accuracy gains on heterogeneous image datasets.

Searching arXiv for the FedPall paper and closely related federated learning work on feature drift. FedPall is a federated learning framework for feature drift, a heterogeneous FL setting in which the label marginal is shared across clients while the class-conditional feature distributions differ, formally $P_i(x \mid y) \neq P_j(x \mid y)$ with $P(y)$ remaining the same across clients. In this regime, samples with the same label may occupy substantially different regions of feature space because of different devices, environments, or acquisition conditions. FedPall—short for Prototype-based Adversarial and Collaborative Learning—addresses this problem by combining prototype-based adversarial learning for feature-space unification, collaborative prototype contrastive learning for preserving class-discriminative structure, and mixed prototype features for training a stronger global classifier (Zhang et al., 7 Jul 2025).

1. Problem formulation and scope

FedPall is designed specifically for feature drift in federated learning. The paper characterizes this setting by the mismatch of class-conditional feature distributions across clients while keeping the label marginal shared. This differs from more generic non-IID formulations because the primary issue is not simply label imbalance or quantity skew, but the fact that identical classes are represented differently across client domains (Zhang et al., 7 Jul 2025).

The paper identifies two direct consequences. First, local feature extraction becomes unstable because each client observes only a restricted local view and may overfit to client-specific appearance patterns. Second, global aggregation becomes less meaningful because parameter averaging does not by itself reconcile the fact that the same class occupies different client-specific feature regions. The paper further notes that standard classification objectives such as cross-entropy do not explicitly align same-class features across clients. This leaves same-class features poorly aligned, decision boundaries ambiguous, and both personalized local models and naively aggregated global models less discriminative.

Within the paper’s framing, feature drift is prevalent in real-life data, particularly when data come from different sources, devices, or environments. This motivates an FL design in which the representation space itself becomes an object of coordination rather than assuming that aggregation alone will induce a useful shared feature geometry.

2. System architecture and training stages

Each local client model $F(\cdot)$ in FedPall contains a feature extractor $G(\cdot)$ and a classifier $H(\cdot)$. The server maintains a global prototype set $\mathcal{G}$, a global amplifier $A$, and a global classifier $C_g$. The amplifier is described as a 3-layer MLP used in the adversarial process; the client-side amplifier is used in the forward pass, while the server-side amplifier is trained on uploaded mixed features (Zhang et al., 7 Jul 2025).

The overall framework is organized into four stages:

1. Generating global prototypes
2. Training local models
3. Training global model
4. Decentralizing the global classifier

This decomposition is central to the method’s logic. Prototype generation supplies a shared class-level reference; local training adapts client encoders under alignment and contrastive constraints; global-model training uses uploaded mixed features rather than raw local features; and decentralization replaces local classifiers with the server-trained global classifier, after which clients may retrain that classifier locally for personalization.

A plausible implication is that FedPall separates the representation problem from the classification problem more explicitly than standard FL pipelines. The feature extractor is coordinated through prototypes and adversarial alignment, whereas the classifier is strengthened using server-visible mixed features that are intended to encode global class information without directly exposing raw features.

3. Prototype-based adversarial alignment

For client $n$, the prototype for class $k$ is defined as the mean feature vector

$P(y)$0

where $P(y)$1 denotes samples of class $P(y)$2 on client $P(y)$3, $P(y)$4 is the number of such samples, and $P(y)$5 is the output of the client feature extractor. The local prototype set is

$P(y)$6

with $P(y)$7 classes and feature dimension $P(y)$8. The server aggregates local prototypes into global prototypes by

$P(y)$9

and forms the global prototype set

$F(\cdot)$0

FedPall then uses an adversarial setup between client and server via the amplifier. The server-side amplifier attempts to amplify heterogeneous information, while the client-side feature encoder is trained to reduce heterogeneity. The KL term is

$F(\cdot)$1

which the paper expands as

$F(\cdot)$2

Here, $F(\cdot)$3 is the client amplifier on client $F(\cdot)$4, and the target is a uniform distribution. The intended effect is to encourage less client-specific, less heterogeneous features. In the paper’s interpretation, this adversarial component pushes client representations toward a more unified feature space while preserving privacy because raw features are not directly shared (Zhang et al., 7 Jul 2025).

The method’s reliance on prototypes is significant. Rather than aligning instance distributions directly across clients, FedPall performs class-level coordination through aggregated prototype anchors. This suggests a compromise between full feature sharing, which raises privacy concerns, and purely parameter-level aggregation, which does not directly address feature drift.

4. Collaborative prototype contrastive learning

The paper argues that adversarial alignment alone can blur class information. FedPall therefore adds prototype contrastive learning based on the server-supplied global prototypes. For each local feature $F(\cdot)$5 with class $F(\cdot)$6, the model minimizes distance to the corresponding global prototype $F(\cdot)$7 and maximizes distance to prototypes of other classes through an InfoNCE loss: $F(\cdot)$8 In this expression, $F(\cdot)$9, $G(\cdot)$0 is a temperature hyperparameter, and $G(\cdot)$1 is cosine similarity.

The full local objective is

$G(\cdot)$2

where $G(\cdot)$3 is cross-entropy loss and $G(\cdot)$4 weight the KL and InfoNCE terms.

This collaborative term uses shared global prototypes to restore class semantics that may be weakened by feature-space alignment. Operationally, clients pull same-class features together and push different-class features apart, but they do so relative to a server-aggregated global class structure rather than only local minibatch geometry. The ablation reported in the paper indicates that the combination of KL and InfoNCE is important; either term alone can be unstable under strong feature drift (Zhang et al., 7 Jul 2025).

A common misconception in this area is that alignment and discrimination are naturally compatible under a single objective. FedPall explicitly treats them as partially competing pressures: KL-based alignment reduces inter-client discrepancy, whereas InfoNCE restores class separation that alignment may otherwise erode. The method’s loss design is built around balancing those two effects.

5. Mixed prototype features and server-side classifier learning

A distinctive component of FedPall is its use of mixed prototype features for training the global classifier. For a local feature $G(\cdot)$5 of class $G(\cdot)$6 on client $G(\cdot)$7, the mixed feature is

$G(\cdot)$8

where $G(\cdot)$9 and $H(\cdot)$0 is the global prototype for class $H(\cdot)$1. This constructs a representation between the client-local feature and the global class anchor.

To reduce privacy leakage, the mixed feature is masked: $H(\cdot)$2 and the transmitted feature becomes

$H(\cdot)$3

Clients upload mixed features $H(\cdot)$4, labels $H(\cdot)$5, and client IDs $H(\cdot)$6. The server trains the global amplifier $H(\cdot)$7 using $H(\cdot)$8 and the global classifier $H(\cdot)$9 using $\mathcal{G}$0, both with cross-entropy empirical risk minimization (Zhang et al., 7 Jul 2025).

The paper’s rationale is that mixed features provide global-view, classification-relevant representations, allowing the server to learn better decision boundaries than it could from raw local models alone. This is not equivalent to ordinary centralized training because the server does not receive raw features; it receives prototype-guided, masked mixtures. This suggests an intermediate representation-sharing strategy in which the server is given enough information to shape a stronger classifier while the masking and mixing mechanisms seek to limit privacy leakage.

The global classifier is then decentralized back to clients, replacing the original local classifier. Clients may retrain it locally for personalization. In effect, the classifier becomes a server-coordinated component operating over locally adapted feature extractors.

6. Communication protocol, evaluation setting, and empirical findings

The round structure reported in the paper proceeds as follows: clients compute local class prototypes; clients send prototypes and label counts to the server; the server aggregates global prototypes and returns them; clients train local feature extractors using cross-entropy, KL adversarial alignment, and InfoNCE with global prototypes; clients construct mixed features; clients send masked mixed features with labels and client IDs; the server trains the global amplifier and classifier; the server deploys the global classifier back to clients; and clients may retrain that classifier locally to adapt it to personalized data (Zhang et al., 7 Jul 2025).

The evaluation uses three feature-drifted benchmarks: Digits (MNIST, SCHN, USPS, SynthDigits, MNIST-M), Office-10 (Amazon, Caltech, DSLR, WebCam), and PACS (Art Painting, Cartoon, Photo, Sketch). Digits uses 10% of the data for feature transformation experiments, while Office-10 and PACS use all data; all datasets are split 8:2 into train/test. The model and training setup specifies ResNet-50 as the local backbone, a 3-layer MLP classifier, a 3-layer MLP amplifier, SGD with learning rate 0.01, local epochs of 5 for Digits and 10 for Office-10/PACS, and 100 global training epochs. The loss weights are $\mathcal{G}$1 for Digits and $\mathcal{G}$2 for the other datasets.

FedPall is compared against SingleSet, FedAvg, FedProx, PerFedAvg, FedRep, FedBN, MOON, FedProto, ADCOL, and RUCR. The paper reports the best overall average top-1 accuracy on all three datasets: 67.45 on Office-10, 88.74 on Digits, and 60.56 on PACS. The next best methods are ADCOL with 64.51 on Office-10 and FedBN with 87.61 on Digits and 59.48 on PACS. The strongest improvement is reported on Office-10, where FedPall exceeds ADCOL by about 3 points on average; on Digits and PACS, the margin over FedBN is about 1.1 points in each case (Zhang et al., 7 Jul 2025).

These results are interpreted in the paper as evidence that FedPall is particularly effective when feature drift is severe and when train-test distribution mismatch is pronounced. Because Office-10 and PACS are described as real-world image datasets with strong feature drift, the larger gain on Office-10 is consistent with the framework’s emphasis on feature-space unification plus class-aware prototype coordination.

7. Ablations, privacy considerations, and limitations

The paper studies four loss configurations: CE only, CE + KL, CE + InfoNCE, and CE + KL + InfoNCE. The full loss performs best overall. On Office-10, CE + InfoNCE alone can even underperform CE-only, which the paper interprets as evidence that emphasizing class separation without alignment can amplify drift. CE + KL helps align clients but may weaken class discrimination in some cases. The three-term combination balances inter-client alignment, intra-client class separation, and classification accuracy (Zhang et al., 7 Jul 2025).

For hyperparameter sensitivity on Office-10, the method is reported to be not highly sensitive for $\mathcal{G}$3 roughly in $\mathcal{G}$4, with $\mathcal{G}$5 having less effect in that range. The best accuracy reported in the heatmap is 69.12 at $\mathcal{G}$6. The paper also ablates the global classifier. Removing it degrades performance, especially on more drifted datasets: on Office-10, the global classifier yields substantial gains, including 49.33% on Caltech, which is 15.55% higher than the baseline without the global classifier; on PACS, it improves all sub-datasets, with gains up to 3%.

The privacy leakage evaluation uses mutual information measured by DEMINE on Office-10, where lower is better. The reported values are 2.96 for Gaussian noise, 3.06 for prototype mixture only, 3.04 for Bernoulli masking only, and 2.10 for FedPall with prototype mixture plus masking. The corresponding average accuracy with FedPall is also the best among these settings at 67.55 (Zhang et al., 7 Jul 2025). This supports the paper’s claim that the combination of prototype mixing and Bernoulli masking is more effective than either mechanism alone in reducing leakage while preserving utility.

The framework’s stated limitations are explicit. It is specialized for image recognition and has not yet been extended to NLP or time-series analysis. The paper notes that such extensions would require task-specific feature and prototype definitions, redesigned losses, and possible architectural changes such as RNN-based temporal models. This suggests that FedPall should be understood as a representation-learning framework for feature-drifted visual FL rather than a modality-agnostic federated recipe.

In summary, FedPall addresses feature drift by attacking three coupled failure modes: inter-client feature-space discrepancy, loss of class-discriminative structure during alignment, and weak global decision boundaries under naïve aggregation. Its combination of adversarial KL alignment, prototype-based InfoNCE collaboration, and mixed-feature global classifier training is presented as the mechanism by which it unifies feature spaces, preserves class information, and improves classification under feature-drifted federated learning (Zhang et al., 7 Jul 2025).