---
title: Federated Attentive Message Passing (FedAMP)
url: https://www.emergentmind.com/topics/federated-attentive-message-passing-fedamp
type: topic
---

# Federated Attentive Message Passing (FedAMP)

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 [2007.03797].

## 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 $P_i$, while still exploiting cross-client similarities. This leads to the following aggregate optimization target:
$$
G(W) = \sum_{i=1}^{K} F_i(w_i) + \lambda \sum_{1\leq i<j\leq K} A(\|w_i - w_j\|^2),
$$
where $W = [w_1,\ldots,w_K] \in \mathbb{R}^{d \times K}$, $\lambda > 0$ balances personalization/collaboration, and $A: \mathbb{R}_+ \rightarrow \mathbb{R}$ is a concave, increasing penalty that induces attention. An example is $A(t)=1-\exp(-t/\sigma)$.

## 2. FedAMP Algorithmic Structure

FedAMP implements an alternating incremental-proximal optimization on $G(W)$. Each communication round $k$ proceeds as follows:

- **Message-Passing / Attention Step (Server-Side):**
  For each client $i$, compute the attentive aggregate:
  $$
  u_i^k = \sum_{j=1}^K \xi_{i,j} \cdot w_j^{k-1},
  $$
  where $\xi_{i,j} = \alpha_k\cdot A'(\|w_i^{k-1}-w_j^{k-1}\|^2)$ for $j \neq i$, and $\xi_{i,i} = 1 - \sum_{j\neq i} \xi_{i,j}$. The update can also be regarded as a perturbed gradient step:
  $$
  u_i^k = w_i^{k-1} - \alpha_k \sum_{j\neq i} A'(\|w_i^{k-1}-w_j^{k-1}\|^2) \cdot 2(w_i^{k-1} - w_j^{k-1}).
  $$
- **Local Proximal Update (Client-Side):**
  Each client $i$ solves:
  $$
  w_i^k = \arg\min_{w} F_i(w) + \frac{\lambda}{2\alpha_k}\|w - u_i^k\|^2.
  $$
  In practice this is implemented via a small number of local SGD or Adam steps.

The algorithm iterates these two steps for $k = 1,2,\ldots,K$. Pseudocode matching the above logic is presented in the original work.

## 3. Attention Mechanism and Similarity Adaptation

The attention kernel $A'(\cdot)$ serves as a nonincreasing, nonnegative similarity function:
- Small $\|w_i - w_j\|$ yields large $A'$, encouraging strong pairwise collaboration.
- Large $\|w_i - w_j\|$ yields small $A'$, limiting influence across dissimilar clients.

A widely used instantiation is the RBF kernel, $A(t) = 1 - \exp(-t/\sigma)$, so $A'(t) = \sigma^{-1} \exp(-t/\sigma)$. Consequently, $\xi_{i,j} \propto \exp(-\|w_i-w_j\|^2/\sigma)$.

The attention coefficients $\xi_{i,j}$ 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 $G(W)$, under bounded-gradient assumptions:

- **Convex Case:** If each $F_i$ and $A$ are convex, and $\alpha_k = \lambda/\sqrt{K}$,
  $$
  \min_{0\leq k\leq K} G(W^k) - G^* \leq O(1/\sqrt{K}).
  $$
  Diminishing $\alpha_k$ ensuring $\sum \alpha_k = \infty, \sum \alpha_k^2 < \infty$ yields $G(W^k)\to G^*$.

- **Smooth, Nonconvex Case:** If $F_i$ and $A$ are $L$-smooth, and $\alpha_k = \lambda/\sqrt{K}$,
  $$
  \min_{0\leq k\leq K} \|\nabla G(W^k)\|^2 \leq O(1/\sqrt{K}).
  $$
  With diminishing $\alpha_k$ as above, any limit point of $\{W^k\}$ 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 ($d$ large, as in DNNs), Euclidean distances become less meaningful. The heuristic variant "HeurFedAMP" alters the computation of $\xi_{i,j}$:

- Set self-attention $\xi_{i,i}$ to $\tau < 1$ (e.g., $\tau = 1/(n_i+1)$).
- For $j\neq i$,
  $$
  \xi_{i,j} = (1-\tau) \frac{\exp(\sigma\cdot \cos(w_i, w_j))}{\sum_{h\neq i}\exp(\sigma\cdot \cos(w_i, w_h))},
  $$
  where $\cos(\cdot,\cdot)$ is cosine similarity and $\sigma>0$ a temperature parameter.

This maintains $\sum_j \xi_{i,j}=1$ 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:** $\lambda$ balances personalization/collaboration; initial $\alpha_k$ should be $O(\lambda)$ then decay $O(1/\sqrt{k})$; attention kernel $\sigma$ must be tuned; self-attention $\tau$ in HeurFedAMP set to $1/(|\text{cluster}|+1)$.
- **Robustness:** Proximal step only requires available $u_i$—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 [2007.03797].

Source: https://www.emergentmind.com/topics/federated-attentive-message-passing-fedamp