[Paper Review] High Dimensional Restrictive Federated Model Selection with multi-objective Bayesian Optimization over shifted distributions

Xudong Sun, Andrea Bommert+4 more|arXiv (Cornell University)|Feb 24, 2019

Machine Learning and Data Classification40 references3 citations

TL;DR

This paper proposes Restrictive Federated Model Selection (RFMS), a novel framework for hyperparameter optimization in high-dimensional, privacy-sensitive clinical data scenarios where models are trained locally but evaluated remotely. Using multi-objective Bayesian optimization over shifted data distributions, RFMS achieves superior generalization across data sites compared to single-objective or random search baselines, with fmo outperforming other methods in hypervolume metrics.

ABSTRACT

A novel machine learning optimization process coined Restrictive Federated Model Selection (RFMS) is proposed under the scenario, for example, when data from healthcare units can not leave the site it is situated on and it is forbidden to carry out training algorithms on remote data sites due to either technical or privacy and trust concerns. To carry out a clinical research under this scenario, an analyst could train a machine learning model only on local data site, but it is still possible to execute a statistical query at a certain cost in the form of sending a machine learning model to some of the remote data sites and get the performance measures as feedback, maybe due to prediction being usually much cheaper. Compared to federated learning, which is optimizing the model parameters directly by carrying out training across all data sites, RFMS trains model parameters only on one local data site but optimizes hyper-parameters across other data sites jointly since hyper-parameters play an important role in machine learning performance. The aim is to get a Pareto optimal model with respective to both local and remote unseen prediction losses, which could generalize well across data sites. In this work, we specifically consider high dimensional data with shifted distributions over data sites. As an initial investigation, Bayesian Optimization especially multi-objective Bayesian Optimization is used to guide an adaptive hyper-parameter optimization process to select models under the RFMS scenario. Empirical results show that solely using the local data site to tune hyper-parameters generalizes poorly across data sites, compared to methods that utilize the local and remote performances. Furthermore, in terms of dominated hypervolumes, multi-objective Bayesian Optimization algorithms show increased performance across multiple data sites among other candidates.

Motivation & Objective

• Address the challenge of training machine learning models in clinical research when data cannot leave local sites due to privacy, trust, or technical constraints.
• Develop a model selection process that generalizes well across data sites with different, potentially shifted, feature distributions.
• Enable efficient hyperparameter tuning without requiring remote data access for training, only for performance evaluation.
• Overcome limitations of standard federated learning by avoiding synchronized, communication-heavy training across sites.
• Improve model generalization by jointly optimizing for both local and remote prediction performance using Bayesian optimization.

Proposed method

• Propose a new learning paradigm, Restrictive Federated Model Selection (RFMS), where model training occurs only on local data, but performance feedback is obtained from remote data sites.
• Use multi-objective Bayesian optimization (MOBO) to jointly optimize for local and remote prediction losses, aiming for Pareto-optimal hyperparameters.
• Model remote performance as a scalar feedback signal sent back after deploying a model, minimizing communication and privacy overhead.
• Apply dimensionality reduction (PCA with 10% variance retention) and clustering to simulate realistic data distribution shifts across clinical data sites.
• Use stratified random splitting to simulate more balanced data distribution scenarios for robust evaluation.
• Employ Gaussian Process regression with a nugget value of 1e-6 to model the objective function, handling performance platet boundaries in hyperparameter search.

Experimental results

Research questions

• RQ1 Can multi-objective Bayesian optimization outperform single-objective and random search in hyperparameter tuning under restrictive federated model selection?
• RQ2 How does model generalization across data sites with shifted distributions compare when using only local data for hyperparameter tuning versus incorporating remote feedback?
• RQ3 Does the performance of MOBO methods improve under different data distribution scenarios (e.g., uneven vs. stratified splits)?
• RQ4 How do distribution shifts in high-dimensional clinical data affect the stability and convergence of hyperparameter optimization?
• RQ5 What is the impact of communication efficiency and asynchronous deployment on the feasibility and performance of model selection in privacy-constrained settings?

Key findings

• Models tuned solely on local data generalize poorly across data sites with distribution shifts, confirming the need for remote performance feedback.
• Multi-objective Bayesian optimization methods (especially fmo) achieve significantly higher hypervolume scores than single-objective and random search baselines.
• In the DRC (Dimension Reduction and Clustering) scenario, fmo outperforms rand_mo by a large margin, winning more than half of the experiments.
• Under the SRS (Stratified Random Split) scenario, all methods show improved performance due to more balanced data distribution, but fmo still dominates.
• The Winner-vs-Loser plots demonstrate consistent superiority of fmo over other candidates, with statistical confidence, across multiple datasets and scenarios.
• Only ~5% of experiments were excluded due to algorithmic failures (e.g., numerical issues in GP regression), indicating robustness of the evaluation pipeline.