Layerwise Cosine Aggregation

• Layerwise Cosine Aggregation is a method that applies cosine similarity at individual neural network layers to robustly aggregate client updates in federated learning.
• It uses median-norm clipping and unit-norm normalization to control scale and ensure directional alignment, addressing the challenges of high-dimensional and heterogeneous data.
• Empirical evaluations show that LCA improves performance under both IID and non-IID conditions and resists Byzantine attacks, making it a practical solution for real-world distributed learning.

Layerwise Cosine Aggregation (LCA) refers to a robust parameter aggregation methodology for federated learning (FL) that incorporates cosine-based similarity measures at the granularity of individual neural network layers. Its primary motivation is to mitigate the shortcomings of classical Euclidean- or norm-based aggregation in high-dimensional or heterogeneous-data regimes, where classical approaches become vulnerable to both statistical inefficiency and adversarial (Byzantine) behavior. LCA is designed to improve both the empirical and theoretical robustness of model training in federated and distributed self-supervised learning (SSL) pipelines (García-Márquez et al., 27 Mar 2025, &&&1&&&).

1. Mathematical Formulation of Layerwise Cosine Aggregation

Let $n$ denote the number of clients participating in a given FL round, from which each returns an update vector $V_i \in \mathbb{R}^d$. The global parameter vector is decomposed into $L$ disjoint layers of sizes $m_1,\dots,m_L$ (so that $d = \sum_{j=1}^L m_j$), such that for client $i$: $$V_i = (V_{i,1}, V_{i,2}, \ldots, V_{i,L}), \quad V_{i,j} \in \mathbb{R}^{m_j}$$

The aggregation for each layer comprises several key steps:

1. Median-norm Clipping: For each layer $j$, compute the median norm among all clients,

$$r_j = \mathrm{median}\{\|V_{1,j}\|, \ldots, \|V_{n,j}\|\}$$

and clip each client’s update,

$$\widetilde{V}_{i,j} = V_{i,j} \times \min \left\{1, \frac{r_j}{\|V_{i,j}\|} \right\}$$

2. Unit-norm Normalization: Normalize the clipped vectors,

$$\bar{V}_{i,j} = \frac{\widetilde{V}_{i,j}}{\|\widetilde{V}_{i,j}\|}$$

3. Cosine-distance Aggregation: Substitute classical robust aggregation operators (e.g., Krum, GeoMed) with ones that utilize the cosine distance,

$$D_{\cos}(X, Y) = 1 - \frac{\langle X, Y \rangle}{\|X\| \|Y\|}$$

Compute a robust central vector $g_j$ for each layer,

$$g_j = \mathcal{F}(\bar{V}_{1,j},\ldots,\bar{V}_{n,j}), \quad \|g_j\|=1$$

4. Reconstruction and Concatenation: Recover the scaling lost during normalization,

$G_j = r_j g_j, \quad G = (G_1,\dots,G_L)$

The global update is then $G \in \mathbb{R}^d$.

This formalism yields the Layerwise Cosine Aggregator: $$LCA_{\mathcal{F}}(V_1,\ldots,V_n) = (r_1 g_1, \ldots, r_L g_L)$$ (García-Márquez et al., 27 Mar 2025)

2. Theoretical Guarantees and Byzantine Robustness

A robust aggregator $\mathcal{R}$ is said to be $(\alpha, f)$-Byzantine-resilient if, given at most $f$ Byzantine (i.e., fully adversarial) updates among $n$ total, the following holds:

• Angular Bias Bound: The aggregate remains aligned with the true mean $g = \mathbb{E}[G]$ of honest updates,

$$\langle \mathbb{E} \, \mathcal{R}(V_1, \ldots, V_{n-f}, B_1, \ldots, B_f), \, g \rangle \ge (1 - \sin\alpha) \|g\|^2$$

• Moment Control: Higher-order moments $\mathbb{E}\|\mathcal{R}\|^r$ are appropriately bounded for $r = 2,3,4$.

It is shown that if the base aggregator $\mathcal{F}$ is $(\alpha, f)$-resilient in $\mathbb{R}^d$, the layerwise application $L\mathcal{F}$ preserves this resilience. A key insight is that using cosine distance and median-norm clipping per layer tightens worst-case guarantees, as the effective dimension per aggregation step reduces from $d$ to $\max_j m_j$: $$\sin\alpha_{LCA} \approx \max_j c \sqrt{m_j} \ll c \sqrt{d}$$ where $c$ is a constant depending on the specific aggregator (García-Márquez et al., 27 Mar 2025).

3. Algorithmic Structure and Implementation

The LCA algorithm for $n$ clients and $L$ layers using base aggregator $\mathcal{F}$ proceeds as:

1. For $j=1,\ldots,L$:
   • Extract $V_{i,j}$ for all $i$.
   • Compute $r_j = \mathrm{median}_i \|V_{i,j}\|$.
   • Apply clipping and unit-norm normalization: $\widetilde{V}_{i,j}$ and $\bar{V}_{i,j}$.
   • Compute $$g_j = \mathcal{F}(\bar{V}_{1,j},\ldots,\bar{V}_{n,j})$$.
   • Form $G_j = r_j g_j$.
2. Aggregate $G = (G_1,\ldots,G_L)$ as the global model update.

The computation requires $O(n^2 d + n d)$ operations per server round (assuming $O(n^2 m_j)$ for the robust aggregator per layer), with communication overhead comparable to FedAvg ($O(d)$ floats per round) (García-Márquez et al., 27 Mar 2025).

4. Empirical Evaluation and Benchmarks

LCA has been empirically benchmarked on image classification datasets including CIFAR-10, Fashion-MNIST, EMNIST, and CelebA-S, across both IID and non-IID data splits. Models used ranged from small CNNs to EfficientNet-B0. Evaluation under Byzantine attacks (label-flipping, model replacement) and no-attack scenarios reveals:

• Under no attack, LCA achieves up to +16 percentage points improvement over vanilla aggregation (e.g., Krum on CelebA-S Non-IID: 72.4% to 88.5%).
• Under hardest label-flipping attacks, LCA recovers accuracy close to the no-attack case, outperforming all tested baselines by 10–15 points.
• Ablation reveals that both the layerwise and cosine components independently yield performance gains, but their combination is most effective, especially when layer sizes differ markedly.

Comparative analysis indicates that, unlike prior robust aggregation schemes, LCA mitigates overfitting in large dense layers and stabilizes validation loss curves (García-Márquez et al., 27 Mar 2025).

5. Relationship to Other Layerwise Aggregation Methods

A related approach, Layer-wise Divergence Aware Weight Aggregation (L-DAWA), focuses on weighted layerwise averaging using angular alignment between client and global models. The aggregation weight for client $k$ at layer $\ell$ is defined by its cosine similarity with the global weights, normalized over all clients: $$\delta^{(\ell)}_k = \frac{\langle w^{G}_{\ell}, w^k_{\ell} \rangle}{\|w^{G}_{\ell}\| \|w^k_{\ell}\|}, \quad \alpha^{(\ell)}_k = \frac{\delta^{(\ell)}_k}{\sum_j \delta^{(\ell)}_j}$$ The global update at layer $\ell$ is then

$$w^{G, r+1}_{\ell} = \sum_{k=1}^K \alpha^{(\ell)}_k w^{k, r}_{\ell}$$

L-DAWA addresses bias and divergence in non-IID or heterogeneous federated self-supervised learning. Empirical evaluations with SimCLR and Barlow Twins demonstrate improvements over FedAvg, with observed benefits of 4-6% in linear-probe accuracy on standard datasets and 20–30% reduction in required communication rounds (Rehman et al., 2023).

A comparison highlights that LCA extends this philosophy with robust (Byzantine-tolerant) components and median-norm clipping, making it suitable for adversarial and high-dimensional settings (García-Márquez et al., 27 Mar 2025, Rehman et al., 2023).

6. Comparison with Classical Aggregators and Significance

The principal benefits of layerwise cosine aggregation relative to classical Euclidean aggregators (e.g., Krum, Bulyan, FedAvg, GeoMed) are:

• Dimensionality Control: Aggregation in lower-dimensional subspaces per layer reduces curse-of-dimensionality effects on robustness.
• Directional Alignment: Cosine-based metrics better capture meaningful agreement between client updates, avoiding scale-induced biases.
• Empirical Robustness: LCA prevents over-selection of large-dimensional layers and outperforms classical rules under both IID and non-IID, as well as Byzantine, scenarios.
• Computational Feasibility: Maintains overall $O(n^2 d)$ cost, and requires no additional communication or learning-rate hyperparameters beyond base robust rules.

This makes LCA and related layerwise-cosine schemes highly suitable for federated learning experiments that demand robustness, such as industrial and medical distributed ML deployments (García-Márquez et al., 27 Mar 2025, Rehman et al., 2023).

7. Practical Recommendations, Limitations, and Open Directions

Practitioners should adopt the following guidelines for effective LCA deployment:

• Use natural neural network layers as the partition.
• Precede normalization by median-norm clipping to control adversarial influence.
• Carefully select the base robust aggregator $\mathcal{F}$; both Krum and GeoMed are compatible.
• Layerwise cosine aggregation is most effective in architectures with balanced (or not overly skewed) layer sizes.
• Communication and memory costs remain linear in model size, as with standard FL schemes.

Current limitations include the assumption of synchronous rounds, fixed known $f$, and reliance on parameter partitions corresponding to layers. Open problems include adaptive metric learning per layer, extensions to asynchronous or unknown-$f$ settings, and exploration of LCA in non-vision modalities (García-Márquez et al., 27 Mar 2025).

---

References:

• "Improving $(\alpha, f)$-Byzantine Resilience in Federated Learning via layerwise aggregation and cosine distance" (García-Márquez et al., 27 Mar 2025)
• "L-DAWA: Layer-wise Divergence Aware Weight Aggregation in Federated Self-Supervised Visual Representation Learning" (Rehman et al., 2023)