Momentum-Based Byzantine-Resilient Non-Convex Stochastic Optimization with Different Batch Sizes Across Different Clients

In recent years, many studies have focused on the robustness of federated learning (FL) systems against Byzantine attacks. Prior work has shown that integrating stochastic optimization algorithms with momentum can effectively enhance the robustness, particularly in non-convex scenarios. However, most studies assume a fixed batch size across all workers, which may not be suitable for practical deployments. In this study, we investigate the impact of different batch sizes across different clients on Byzantine robustness in FL. Specifically, we propose a weighted aggregation framework (WAggF), which assigns different weights to gradients received from different clients. WAggF can be easily embedded into various aggregation rules to boost Byzantine resilience. We further establish theoretical convergence guarantees for two momentum-based stochastic optimization algorithms—Byz-WNSGDM and Byz-EF21-WSGDM—under WAggF. To reduce communication overhead and enable synchronous training, we incorporate a contractive compressor (which encompasses Top-K as a special case) into the stochastic optimization algorithm. Extensive experiments demonstrate that our proposed algorithms consistently improve robustness against a range of Byzantine attacks.