pFedBBN: A Personalized Federated Test-Time Adaptation with Balanced Batch Normalization for Class-Imbalanced Data

Abstract

Test-time adaptation (TTA) in federated learning (FL) is crucial for handling unseen data distributions across clients, particularly when faced with domain shifts and skewed class distributions. Class Imbalance (CI) remains a fundamental challenge in FL, where rare but critical classes are often severely underrepresented in individual client datasets. Although prior work has addressed CI during training through reliable aggregation and local class distribution alignment, these methods typically rely on access to labeled data or coordination among clients, and none address class unsupervised adaptation to dynamic domains or distribution shifts at inference time under federated CI constraints. Revealing the failure of state-of-the-art TTA in federated client adaptation in CI scenario, we propose pFedBBN,a personalized federated test-time adaptation framework that employs balanced batch normalization (BBN) during local client adaptation to mitigate prediction bias by treating all classes equally, while also enabling client collaboration guided by BBN similarity, ensuring that clients with similar balanced representations reinforce each other and that adaptation remains aligned with domain-specific characteristics. pFedBBN supports fully unsupervised local adaptation and introduces a class-aware model aggregation strategy that enables personalized inference without compromising privacy. It addresses both distribution shifts and class imbalance through balanced feature normalization and domain-aware collaboration, without requiring any labeled or raw data from clients. Extensive experiments across diverse baselines show that pFedBBN consistently enhances robustness and minority-class performance over state-of-the-art FL and TTA methods.

Figures (20)

• Figure 1: The overall framework of pFedBBN where each client performs unsupervised local adaptation using class-wise balanced batch normalization (BBN) and confidence-filtered distillation. Adapted batch normalization statistics are then used to compute client similarities, enabling personalized aggregation without sharing raw data.
• Figure 2: Average global class accuracies (%) of different TTA methods under different Federated learning Schemes. Our method (BBN) consistently delivers best performance under different federated setups for both CIFAR-10-C and CIFAR-100-C datasets.
• Figure 3: Domain-wise accuracy comparison of BBN under different Federated setups (Dirichlet $\delta=0.01$) across CIFAR-10-C corruptions.
• Figure 4: Major and Minor Class accuracy per client for CIFAR-10-C (Dirichlet $\delta$ = 0.01) using pFedBBN strategy. The ten classes of CIFAR-10-C are denoted by the symbols from C0 to C9. The accuracies for major and minor classes are similar for almost all the clients, demonstrating mitigation of the class-imbalance issue.
• Figure 5: Collaboration Matrix (round = 40, 75) of our pFedBBN (total 75 federated rounds) which indicates which client give more priority while aggregating and it has clearly be seen that similar domains are aggregated more than others