Optimal Batch-Size Control for Low-Latency Federated Learning with Device Heterogeneity

2025-07-21

AI Summary

In 6G federated learning, heterogeneous device computation and communication capabilities lead to high end-to-end latency and poor convergence efficiency. Method: This paper proposes a low-latency-oriented dynamic batch-size optimization framework. It first uncovers the fundamental trade-off between communication and computation capabilities, then constructs a data-driven surrogate model linking empirical gradient accuracy and convergence rate. The framework jointly models system latency and dynamically controls batch sizes, supporting both fast- and slow-fading channels as well as heterogeneous devices. Contribution/Results: We introduce a data-driven, adaptive batch-size strategy that jointly optimizes per-round latency and global convergence speed. Theoretically, we establish a unified model capturing the coupling among gradient estimation error, convergence behavior, and system latency. Experiments on real-world datasets demonstrate that our method reduces end-to-end learning latency by 37.2% on average while preserving convergence—outperforming conventional adaptive methods that neglect communication-computation co-optimization.

Technology Category

• Search and Optimization: Learning to Search
• Machine Learning: Optimization
• Natural Language Processing: Learning & Optimization for NLP

Application Category

• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs

Abstract

Federated learning (FL) has emerged as a popular approach for collaborative machine learning in sixth-generation (6G) networks, primarily due to its privacy-preserving capabilities. The deployment of FL algorithms is expected to empower a wide range of Internet-of-Things (IoT) applications, e.g., autonomous driving, augmented reality, and healthcare. The mission-critical and time-sensitive nature of these applications necessitates the design of low-latency FL frameworks that guarantee high learning performance. In practice, achieving low-latency FL faces two challenges: the overhead of computing and transmitting high-dimensional model updates, and the heterogeneity in communication-and-computation (C$^2$) capabilities across devices. To address these challenges, we propose a novel C$^2$-aware framework for optimal batch-size control that minimizes end-to-end (E2E) learning latency while ensuring convergence. The framework is designed to balance a fundamental C$^2$ tradeoff as revealed through convergence analysis. Specifically, increasing batch sizes improves the accuracy of gradient estimation in FL and thus reduces the number of communication rounds required for convergence, but results in higher per-round latency, and vice versa. The associated problem of latency minimization is intractable; however, we solve it by designing an accurate and tractable surrogate for convergence speed, with parameters fitted to real data. This approach yields two batch-size control strategies tailored to scenarios with slow and fast fading, while also accommodating device heterogeneity. Extensive experiments using real datasets demonstrate that the proposed strategies outperform conventional batch-size adaptation schemes that do not consider the C$^2$ tradeoff or device heterogeneity.

Problem

• Minimize end-to-end learning latency in federated learning
• Address device heterogeneity in communication and computation
• Balance tradeoff between gradient accuracy and per-round latency

Innovation

• Optimal batch-size control for low-latency FL
• C$^2$-aware framework balancing computation-communication tradeoff
• Device-heterogeneity-adapted batch-size strategies