FLUX: Efficient Descriptor-Driven Clustered Federated Learning under Arbitrary Distribution Shifts

2025-11-27

AI Summary

Federated learning (FL) suffers significant degradation in global model performance under non-independent and identically distributed (Non-IID) data. To address diverse distribution shifts during both training and inference, this paper proposes a prior-free adaptive clustering framework. It jointly performs privacy-preserving client descriptor extraction, unsupervised clustering, and cluster-specific model aggregation—without requiring pre-specified numbers of clusters or prior knowledge of shift types. Crucially, it enables zero-shot, label-free test-time client assignment to the optimal cluster model. The method retains FedAvg’s communication efficiency while enhancing generalization and scalability. Extensive evaluation across four benchmark and two real-world datasets demonstrates up to a 23-percentage-point improvement in average accuracy over baselines, with computational and communication overhead comparable to FedAvg.

Technology Category

• Machine Learning: Distributed Machine Learning & Federated Learning
• Computer Vision: Large Vision Models
• Search and Optimization: Learning to Search

Application Category

• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

Federated Learning (FL) enables collaborative model training across multiple clients while preserving data privacy. Traditional FL methods often use a global model to fit all clients, assuming that clients' data are independent and identically distributed (IID). However, when this assumption does not hold, the global model accuracy may drop significantly, limiting FL applicability in real-world scenarios. To address this gap, we propose FLUX, a novel clustering-based FL (CFL) framework that addresses the four most common types of distribution shifts during both training and test time. To this end, FLUX leverages privacy-preserving client-side descriptor extraction and unsupervised clustering to ensure robust performance and scalability across varying levels and types of distribution shifts. Unlike existing CFL methods addressing non-IID client distribution shifts, FLUX i) does not require any prior knowledge of the types of distribution shifts or the number of client clusters, and ii) supports test-time adaptation, enabling unseen and unlabeled clients to benefit from the most suitable cluster-specific models. Extensive experiments across four standard benchmarks, two real-world datasets and ten state-of-the-art baselines show that FLUX improves performance and stability under diverse distribution shifts, achieving an average accuracy gain of up to 23 percentage points over the best-performing baselines, while maintaining computational and communication overhead comparable to FedAvg.

Problem

• Addresses non-IID data distribution shifts in federated learning
• Proposes a clustering method without prior knowledge of shift types
• Enables test-time adaptation for unseen clients with unlabeled data

Innovation

• Clustering-based FL framework for arbitrary distribution shifts
• Privacy-preserving descriptor extraction and unsupervised clustering
• Test-time adaptation without prior knowledge of shifts or clusters