Breaking the Learning-Rate Dilemma in Clustered Federated Learning

Clustered federated learning (CFL) jointly learns the client partition and the model for each cluster, where the cluster models are derived via model averaging. Existing CFL methods have shown that CFL depends on good initialization, and these methods often use small or diminishing training learning rates for stable model optimization. If the initialization is inaccurate in CFL, a client may be initially assigned to the wrong cluster; in this way, the small or diminishing learning rate and model averaging can restrict the updates that the client makes to the wrong cluster model, leaving it trapped in the wrong cluster persistently. We refer to this as the *learning-rate dilemma* in CFL. To break this learning-rate dilemma, we propose Probe-based Cluster Recovery for CFL (PCR-CFL), which calculates the accumulated distances between the responses of the clients to a sequence of probe models, and uses spectral clustering to estimate the number of clusters and obtain the client partition. This process allows a misassigned client to leave the wrong cluster since its responses to the same sequence of probe models are close to those of clients in its ground-truth cluster. We theoretically prove that PCR-CFL recovers the ground-truth cluster structure with high probability after a finite number of probes, and further establish model convergence under both constant and diminishing learning rates. Experiments on extensive datasets, together with extensive ablations, demonstrate the effectiveness and robustness of PCR-CFL.