Federated Attentive Message Passing (FedAMP)

• FedAMP is a federated learning strategy that uses pairwise adaptive attention to personalize client models on non-IID data.
• It alternates between attentive message passing on the server and local proximal updates on client devices to optimize personalized parameters.
• Empirical evaluations on FMNIST, EMNIST, and CIFAR100 show improved accuracy over traditional federated methods with robust convergence guarantees.

Federated Attentive Message Passing (FedAMP) is a federated learning methodology designed to enable client models to collaborate via adaptive, pairwise attention mechanisms, with the specific goal of improving performance on non-IID data distributions. By leveraging attention-inducing communication between models, FedAMP personalizes learned parameters for each client while maximizing the benefits of inter-client similarity. The approach was introduced as a solution to the persistent challenge of non-IID data in cross-silo federated learning, offering provable convergence, practical robustness, and demonstrably superior empirical results when compared to established methods (Huang et al., 2020).

1. Formal Problem Framework and Objective

Let $K$ denote the number of clients, each indexed by $i=1,\ldots,K$. Each client $i$ holds:

• A private dataset $D_i$ sampled from distribution $P_i$ (non-IID over $i$).
• Local model parameters $w_i \in \mathbb{R}^d$.
• A loss function $$F_i(w_i) := |D_i|^{-1} \sum_{(x,y)\in D_i} \ell(w_i;x,y)$$.

The global objective is to learn personalized parameters $\{w_i\}$ such that each $w_i$ is near-optimal for its own distribution $i=1,\ldots,K$0, while still exploiting cross-client similarities. This leads to the following aggregate optimization target:

$i=1,\ldots,K$1

where $i=1,\ldots,K$2, $i=1,\ldots,K$3 balances personalization/collaboration, and $i=1,\ldots,K$4 is a concave, increasing penalty that induces attention. An example is $i=1,\ldots,K$5.

2. FedAMP Algorithmic Structure

FedAMP implements an alternating incremental-proximal optimization on $i=1,\ldots,K$6. Each communication round $i=1,\ldots,K$7 proceeds as follows:

• Message-Passing / Attention Step (Server-Side):

For each client $i=1,\ldots,K$8, compute the attentive aggregate:

$i=1,\ldots,K$9

where $i$0 for $i$1, and $i$2. The update can also be regarded as a perturbed gradient step:

$i$3

• Local Proximal Update (Client-Side):

Each client $i$4 solves:

$i$5

In practice this is implemented via a small number of local SGD or Adam steps.

The algorithm iterates these two steps for $i$6. Pseudocode matching the above logic is presented in the original work.

3. Attention Mechanism and Similarity Adaptation

The attention kernel $i$7 serves as a nonincreasing, nonnegative similarity function:

• Small $i$8 yields large $i$9, encouraging strong pairwise collaboration.
• Large $D_i$0 yields small $D_i$1, limiting influence across dissimilar clients.

A widely used instantiation is the RBF kernel, $D_i$2, so $D_i$3. Consequently, $D_i$4.

The attention coefficients $D_i$5 thus implement a form of adaptive, pairwise, non-linear communication, automatically amplifying within-cluster collaboration on non-IID data.

4. Theoretical Convergence Analysis

FedAMP offers convergence guarantees for both convex and nonconvex formulations of the objective $D_i$6, under bounded-gradient assumptions:

• Convex Case: If each $D_i$7 and $D_i$8 are convex, and $D_i$9,

$P_i$0

Diminishing $P_i$1 ensuring $P_i$2 yields $P_i$3.

• Smooth, Nonconvex Case: If $P_i$4 and $P_i$5 are $P_i$6-smooth, and $P_i$7,

$P_i$8

With diminishing $P_i$9 as above, any limit point of $i$0 is stationary.

The two-step update is interpretable as a proximal-gradient procedure, and analysis leverages established incremental/proximal methods.

5. Heuristic Extension for Deep Neural Models

For high-dimensional parameterizations ($i$1 large, as in DNNs), Euclidean distances become less meaningful. The heuristic variant "HeurFedAMP" alters the computation of $i$2:

• Set self-attention $i$3 to $i$4 (e.g., $i$5).
• For $i$6,

$i$7

where $i$8 is cosine similarity and $i$9 a temperature parameter.

This maintains $w_i \in \mathbb{R}^d$0 while biasing attention based on angular rather than Euclidean closeness, empirically improving performance on DNNs.

6. Empirical Evaluation and Results

FedAMP and its heuristic extension are evaluated on MNIST, FMNIST, EMNIST, and CIFAR100 datasets with client partitions covering IID, pathological non-IID (each client only 2 labels), and practical non-IID (clients in 3 clusters with unbalanced samples).

Mean testing accuracy (BMTA) under the practical non-IID scenario (mean over clients):

Dataset	FedAvg	FedProx	APFL	FedAMP	HeurFedAMP
FMNIST	79.5%	78.7%	84.1%	91.0%	91.4%
EMNIST	N/A	N/A	N/A	81.2%	81.5%
CIFAR100	35.2%	37.3%	N/A	N/A	53.3%

Pairwise attention heatmaps (EMNIST, clients 0–61) reveal that attention coefficients form clear blocks, aligning with ground-truth clusters—FedAMP automatically learns and exploits such latent structure.

7. Practical Guidelines and Implications

Key operational insights include:

• Data regime sensitivity: On IID data, FedAMP reduces to global averaging (like FedAvg); on clustered non-IID data, it amplifies within-cluster collaboration.
• Hyperparameters: $w_i \in \mathbb{R}^d$1 balances personalization/collaboration; initial $w_i \in \mathbb{R}^d$2 should be $w_i \in \mathbb{R}^d$3 then decay $w_i \in \mathbb{R}^d$4; attention kernel $w_i \in \mathbb{R}^d$5 must be tuned; self-attention $w_i \in \mathbb{R}^d$6 in HeurFedAMP set to $w_i \in \mathbb{R}^d$7.
• Robustness: Proximal step only requires available $w_i \in \mathbb{R}^d$8—drops are naturally handled; attention down-weights corrupted or noisy clients, conferring resilience to label noise.

FedAMP constitutes a principled, provably convergent, and empirically validated framework for federated learning with adaptive, pairwise, non-linear collaboration, with particular effectiveness on non-IID problems and high-dimensional models (Huang et al., 2020).