ZMASA: Robust Aggregation for Federated Learning Against Byzantine Attacks

Federated learning (FL) enables collaborative model training without sharing raw data, but its robustness is vulnerable to Byzantine clients, especially under non-identically distributed (non-IID) data. In heterogeneous FL, benign client updates may become multimodal and statistically diverse, making it difficult for majority-based defenses to distinguish malicious deviations from legitimate client diversity. Reference-based defenses, on the other hand, may suffer from distributional mismatch and often require trusted server data. To address these challenges, we propose Z-score and Minkowski distance-based aggregation with multihead self-attention (ZMASA), a two-stage Byzantine-robust aggregation framework for non-IID FL. In the first stage, ZMASA performs dimension-wise Z-score normalization on client update residuals, partitions clients into hierarchical confidence groups, and validates groups using a high-order Minkowski-distance criterion. This design suppresses coordinated malicious updates without requiring prior knowledge of the Byzantine ratio. In the second stage, ZMASA applies multihead self-attention (MHSA) to the retained candidate updates to adaptively assign aggregation weights and accommodate heterogeneous honest behaviors. We provide a theoretical analysis showing that, under bounded retained residuals, the aggregation error is controlled by the candidate-set Byzantine contamination, the magnitude of retained Byzantine residuals, and the stability of attention weights. Experiments on MNIST, CIFAR-10, AG-News, and CIFAR-100 under diverse Byzantine attacks and non-IID settings demonstrate that ZMASA achieves competitive or superior robustness compared with classical and recent robust aggregation baselines. Ablation, parameter sensitivity, scalability, and efficiency studies further validate the effectiveness and practicality of the proposed design.