pFedPM: Feature Fusion & Relation Networks

• The paper introduces pFedPM, a federated learning framework that combines low-dimensional feature fusion with relation networks to integrate local adaptation and global consensus.
• Feature fusion techniques, including prototype and embedding fusion with alignment, reduce communication burdens and enhance privacy compared to conventional gradient aggregation.
• Empirical studies show that pFedPM significantly boosts performance—with notable gains in F1 score and AUC—across heterogeneous, non-IID datasets.

Feature fusion and relation networks in personalized federated learning frameworks—referred to as pFedPM in several recent works—address the challenge of aggregating knowledge across heterogeneous clients while preserving personalization, privacy, and communication efficiency. Rooted in advances in federated feature fusion and nonlinear classification layers, these methods decouple local pre-training from global consensus and exploit representation-level collaboration to overcome the limitations of classical gradient or model aggregation approaches under data, model, and label distribution heterogeneity (Ma et al., 2023, Xing et al., 2024).

1. Federated Learning Preliminaries and Motivation

Classical federated learning (FL) frameworks optimize a global model $w$ by minimizing the average classification loss across all $m$ clients, each with its own data distribution $P_i$, i.e.,

$$f(w) = \frac{1}{m} \sum_{i=1}^m \mathbb{E}_{(x,y)\sim P_i}\left[\ell(w;x,y)\right].$$

This approach implicitly assumes that all clients observe all classes (labels), and that client data is at least partially IID. In practical deployments, however, clients may have strong statistical heterogeneity, architectural constraints, or access only to a subset of classes (label distribution skew). Under such circumstances, directly aggregating gradients or model parameters fails to produce models that generalize well to each user or device (Xing et al., 2024).

Personalized federated learning relaxes the global model constraint by learning a collection of client-specific models $\{w_i\}$, optimizing

$$\min\quad \sum_{i=1}^m\left(\frac{|D_i|}{N}\cdot \ell_i(w_i)\right),\qquad N = \sum_{i}|D_i|,$$

with the aim of balancing local adaptation and knowledge sharing.

2. Feature Fusion Mechanisms

The cornerstone of pFedPM is the replacement of high-dimensional gradient/model exchanges with low-dimensional feature or prototype fusion. Two principal approaches are found in the literature:

• Prototype-Based Feature Fusion (Xing et al., 2024): Each client maintains a feature extractor $f_{\phi_i}$ and computes class prototypes

$$C_i^{(j)} = \frac{1}{|D_{i,j}|}\sum_{(x,y)\in D_{i,j}} f_{\phi_i}(x).$$

The central server forms global prototypes per class via a weighted average across clients with class $j$:

$$\bar{C}^{(j)} = \frac{1}{|\mathcal{C}_j|}\sum_{i\in\mathcal{C}_j} \left(\frac{|D_{i,j}|}{N_j}\right)C_i^{(j)},$$

where $m$0 and $m$1 is the set of clients with class $m$2.

• Direct Embedding Fusion with Alignment ((Ma et al., 2023); Editor’s term: FF-align): After local model pre-training, clients send feature embeddings $m$3 for all or a subset of samples to the server, which learns linear alignment matrices $m$4 for each client to align heterogeneous embeddings:

$m$5

with $m$6 either unconstrained ("soft" alignment) or (approximately) doubly-stochastic ("hard" alignment via Sinkhorn–Knopp normalization).

This feature-level aggregation supports arbitrary model architectures, as only the feature dimension $m$7 must be shared across clients.

3. Personalization-Consensus Tradeoff

To jointly exploit local specificity and global knowledge, pFedPM fuses local and global feature/prototype representations with a tunable parameter $m$8:

$m$9

where $P_i$0 yields a fully personalized model and $P_i$1 yields a fully global (shared) model (Xing et al., 2024). The optimal value depends on the degree of label distribution skew, with higher $P_i$2 preferred as heterogeneity increases. This framework directly interpolates between global consensus and local adaptation.

Empirical studies demonstrate that convergence speed and accuracy are sensitive to $P_i$3, with suboptimal values impeding either personalization or knowledge transfer.

4. Relation Networks for Nonlinear Classification

pFedPM incorporates a relation network as an additional, flexible, non-linear classifier layer. Each client, after extracting a feature $P_i$4, forms class-pair inputs $P_i$5 and passes them through a relation network $P_i$6 to produce similarity scores:

$P_i$7

Softmax over $P_i$8 produces class probabilities. The relation network's parameters are optimized by minimizing the mean-squared error between $P_i$9 and the one-hot indicator $$f(w) = \frac{1}{m} \sum_{i=1}^m \mathbb{E}_{(x,y)\sim P_i}\left[\ell(w;x,y)\right].$$0 over the client’s dataset. This approach generalizes linear or fixed similarity measures and captures complex relationships between instance features and class prototypes (Xing et al., 2024).

5. Consensus Graph Construction and GCN Fusion

In the federated feature fusion framework of (Ma et al., 2023), finer-grained consensus is achieved by learning both (i) linear alignment transformations across client embeddings and (ii) a client-wise consensus graph:

• Each client is treated as a node; the server learns the adjacency matrix $$f(w) = \frac{1}{m} \sum_{i=1}^m \mathbb{E}_{(x,y)\sim P_i}\left[\ell(w;x,y)\right].$$1 parameterized via Bernoulli variables $$f(w) = \frac{1}{m} \sum_{i=1}^m \mathbb{E}_{(x,y)\sim P_i}\left[\ell(w;x,y)\right].$$2.
• To enable gradient-based optimization through discrete sampling, the ICDF reparameterization is employed: for reference CDF $$f(w) = \frac{1}{m} \sum_{i=1}^m \mathbb{E}_{(x,y)\sim P_i}\left[\ell(w;x,y)\right].$$3, draw $$f(w) = \frac{1}{m} \sum_{i=1}^m \mathbb{E}_{(x,y)\sim P_i}\left[\ell(w;x,y)\right].$$4, set $$f(w) = \frac{1}{m} \sum_{i=1}^m \mathbb{E}_{(x,y)\sim P_i}\left[\ell(w;x,y)\right].$$5, and define $$f(w) = \frac{1}{m} \sum_{i=1}^m \mathbb{E}_{(x,y)\sim P_i}\left[\ell(w;x,y)\right].$$6; final $$f(w) = \frac{1}{m} \sum_{i=1}^m \mathbb{E}_{(x,y)\sim P_i}\left[\ell(w;x,y)\right].$$7 is normalized as in Kipf & Welling's GCN.
• Two layers of graph convolution are performed over the aligned feature matrix $$f(w) = \frac{1}{m} \sum_{i=1}^m \mathbb{E}_{(x,y)\sim P_i}\left[\ell(w;x,y)\right].$$8 to yield predictions.

Optionally, sparsity and orthogonality regularizers on $$f(w) = \frac{1}{m} \sum_{i=1}^m \mathbb{E}_{(x,y)\sim P_i}\left[\ell(w;x,y)\right].$$9 and $\{w_i\}$0 respectively are included to control graph complexity and enforce structure.

6. Communication Protocol and Privacy Considerations

Both prototype fusion and embedding fusion approaches achieve significant communication efficiency by sharing only low-dimensional feature representations per class or per sample, as opposed to entire model weights or gradients. For a typical classification task, this reduces server-client communication per round by up to 10–50$\{w_i\}$1 relative to standard FL (Xing et al., 2024).

Furthermore, these frameworks limit privacy risks: by averaging features across samples or classes, the risk of reconstructing individual instances ("deep leakage") is reduced compared to gradient-based updates. Only aggregated statistics, not raw data nor gradients, are communicated to the server.

The overall optimization proceeds by alternating local feature/statistics extraction, aggregation at the server, fusion/prototype mixing, and local as well as relation-network updates. This process is compatible with heterogeneous local model architectures provided the feature dimension is consistent.

7. Empirical Performance and Applications

Empirical evaluations on datasets subject to data and label heterogeneity (METR-LA, PEMS-Bay, MNIST, FEMNIST, CIFAR-10) validate the advantages of feature fusion with relation networks:

• On traffic and power grid tasks, feature fusion with alignment and learned consensus graphs yields up to double the F$\{w_i\}$2 score of horizontal FL or majority voting. Compared to a VFL "concatenate-and-train" baseline, adding alignment and graph learning improves F$\{w_i\}$3 by 5--10 points and AUC by 2–4 points (Ma et al., 2023).
• On image benchmarks under label skew, pFedPM reaches 98.50$\{w_i\}$40.42% (MNIST), 96.93$\{w_i\}$52.24% (FEMNIST), and 65.42$\{w_i\}$64.25% (CIFAR-10), outperforming FedAvg, FedProx, FeSEM, FedPer, FedRep, and L2GD in both accuracy and communication cost (Xing et al., 2024).
• Ablation studies indicate that soft alignment, consensus graph learning, and relation network layers substantively contribute to performance gains.

These approaches are especially practical in distributed sensing, non-IID classification, and settings where device or client model heterogeneity precludes classical federated synchronization, as they can be realized in a single federated communication round and scale to complex data modalities.

---

References

• "Federated Learning of Models Pre-Trained on Different Features with Consensus Graphs" (Ma et al., 2023)
• "Personalized federated learning based on feature fusion" (Xing et al., 2024)