Beyond Aggregation: Guiding Clients in Heterogeneous Federated Learning

Federated learning (FL) is increasingly adopted in domains like healthcare, where data privacy is paramount. A fundamental challenge in these systems is statistical heterogeneity-the fact that data distributions vary significantly across clients (e.g., different hospitals may treat distinct patient demographics). While current FL algorithms focus on aggregating model updates from these heterogeneous clients, the potential of the central server remains under-explored. This paper is motivated by a healthcare scenario: could a central server not only coordinate model training but also guide a new patient to the hospital best equipped for their specific condition? We generalize this idea to propose a novel paradigm for FL systems where the server actively guides the allocation of new tasks or queries to the most appropriate client. To enable this, we introduce a density ratio model and empirical likelihood-based framework that simultaneously addresses two goals: (1) learning effective local models on each client, and (2) finding the best matching client for a new query. Empirical results demonstrate the framework’s effectiveness on benchmark datasets, showing improvements in both model accuracy and the precision of client guidance compared to standard FL approaches. This work opens a new direction for building more intelligent and resource-efficient FL systems that leverage heterogeneity as a feature, not just a bug. Code is available at https://github.com/zijianwang0510/FedDRM.git.

💡 Research Summary

This paper introduces FedDRM, a novel federated learning (FL) framework that transforms statistical heterogeneity from a drawback into a useful resource. Traditional FL methods focus on aggregating client updates into a single global model and, more recently, on personalizing models for each client. Neither line of work addresses the problem of directing new queries or tasks to the most appropriate client. FedDRM fills this gap by jointly learning (i) accurate local predictive models and (ii) a client‑identification model that enables the central server to route incoming queries to the client whose data distribution best matches the query.

The technical core combines a Density Ratio Model (DRM) with Empirical Likelihood (EL). Each client’s joint distribution (P_i(X,Y)) is expressed as a multiplicative tilt of a reference distribution (P_0(X)). Conditional label probabilities are modeled with a soft‑max parameterized by (\alpha_k,\beta_k) on a shared embedding (g_\theta(x)). The marginal feature distribution of client (i) is linked to the reference via an exponential density‑ratio (\exp{\gamma_i + \xi_i^\top h_\tau(g_\theta(x))}). EL treats the reference distribution as an atomic measure with probabilities (p_{ij}) assigned to each observed sample, imposing two constraints: (1) probabilities sum to one, and (2) the weighted exponential tilts also sum to one. Solving the EL problem yields closed‑form Lagrange multipliers and a profile log‑likelihood that decomposes into two cross‑entropy terms: (a) a standard classification loss for the target task, and (b) a client‑classification loss that predicts the originating client of each sample.

Training proceeds in the usual FL round‑based fashion. The server broadcasts the current global embedding parameters (\theta) and DRM parameters ((\gamma,\xi)). Each client updates its local embedding, a target‑class head, and a shared client‑class head using its private data, then sends the updated parameters back. The server aggregates these updates, solves the EL constraints to obtain the Lagrange multipliers, and refines the global parameters. When label shift is present, the authors propose a simple re‑weighting of the client‑classification loss to mitigate extreme label imbalance.

Empirical evaluation on benchmark vision datasets (CIFAR‑10, FEMNIST) and a real‑world healthcare dataset (eICU) with 20–100 simulated clients demonstrates that FedDRM consistently outperforms strong baselines such as FedAvg, FedProx, pFedMe, and Ditto. Accuracy improvements range from 3 % to 7 % over baselines, while the routing precision—measured as the top‑1 match rate of a new query to the correct client—exceeds 85 %, a several‑fold gain over random assignment. Ablation studies confirm that both the DRM component and the client‑identification loss are essential; removing either degrades performance markedly.

The authors acknowledge two limitations. First, solving the EL constraints requires computing Lagrange multipliers for each client, leading to linear scaling of computational cost with the number of clients, which may become burdensome in very large‑scale deployments. Second, the DRM assumption of an exponential tilt may be violated when client distributions are extremely divergent, potentially reducing the quality of the density‑ratio estimates. Future work is suggested to explore approximate solvers, non‑linear density‑ratio models (e.g., neural density‑ratio estimation), and real‑time deployment in clinical settings for patient‑to‑hospital routing.

In summary, FedDRM pioneers an “intelligent router” role for the FL server, jointly learning predictive models and distributional fingerprints of clients. By leveraging heterogeneity rather than suppressing it, the framework opens a new direction for FL systems that are simultaneously privacy‑preserving, personalized, and capable of task‑specific client allocation.