🤖 AI Summary
This study addresses the dilemma in federated learning on memory-constrained edge devices, where first-order optimization incurs prohibitive memory overhead while zero-order methods suffer from slow convergence. To overcome this, we propose a mixed-order optimization framework that dynamically partitions network layers according to each client's memory budget, applying zero-order estimation to lower layers and first-order training to upper layers. Furthermore, by revealing the trade-off between precision and data representativeness, we construct a dimension-aware sampling strategy to optimize gradient estimation. Both theoretical analysis and empirical results demonstrate that the proposed framework substantially reduces client-side memory requirements while maintaining task performance comparable to full first-order optimization. This work provides an efficient solution for federated learning in resource-constrained environments.
📝 Abstract
Federated learning (FL) on memory-constrained edge devices faces a dilemma: first-order (FO) optimization (i.e., backpropagation) demands substantial memory, whereas zeroth-order (ZO) optimization suffers from severe convergence slowdown. To resolve this dilemma, we introduce HO-FL, a hybrid-order FL framework that trains a model's bottom segment with ZO optimization and its top segment with FO optimization. Each device can flexibly select its order boundary according to its memory budget while participating in the training of the same global model. Moreover, our convergence analysis reveals a new, fundamental trade-off: clients with larger FO-trained segments can provide more accurate updates, but favoring them can underrepresent other clients' data. We connect this trade-off to the bias and variance of actual multi-step local updates, yielding a sampling optimization problem and a practical dimension-aware approximation with direct model averaging. Experiments on language tasks examine task performance, client memory, and sampling under data heterogeneity. The results show that hybrid-order local training can retain much of the full-FO performance with substantially lower client memory requirements. Our code is available at https://github.com/HKU-WILL-Lab/HO-FL.