# Stratify: Rethinking Federated Learning for Non-IID Data through Balanced Sampling

arXiv: https://arxiv.org/abs/2504.13462 · HTML: https://mlsanity.com/paper/2504.13462/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 21 of its 32 references that are in the map.
- First posted on arXiv 18 Apr 2025.

## Builds on (most cited first, 21 of 21)

- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://mlsanity.com/paper/1502.03167/) (2015)
- [Federated Optimization in Heterogeneous Networks](https://mlsanity.com/paper/1812.06127/) (2018)
- [Federated Learning on Non-IID Data Silos: An Experimental Study](https://mlsanity.com/paper/2102.02079/) (2021)
- [Learning to Generate Novel Domains for Domain Generalization](https://mlsanity.com/paper/2007.03304/) (2020)
- [Personalized Cross-Silo Federated Learning on Non-IID Data](https://mlsanity.com/paper/2007.03797/) (2020)
- [Federated learning with hierarchical clustering of local updates to improve training on non-IID data](https://mlsanity.com/paper/2004.11791/) (2020)
- [Clustered Sampling: Low-Variance and Improved Representativity for Clients Selection in Federated Learning](https://mlsanity.com/paper/2105.05883/) (2021)
- [Federated Learning for Medical Image Analysis: A Survey](https://mlsanity.com/paper/2306.05980/) (2023)
- [Adapt to Adaptation: Learning Personalization for Cross-Silo Federated Learning](https://mlsanity.com/paper/2110.08394/) (2021)
- [Federated Learning under Distributed Concept Drift](https://mlsanity.com/paper/2206.00799/) (2022)
- [Efficient Distribution Similarity Identification in Clustered Federated Learning via Principal Angles Between Client Data Subspaces](https://mlsanity.com/paper/2209.10526/) (2022)
- [A Multi-agent Reinforcement Learning Approach for Efficient Client Selection in Federated Learning](https://mlsanity.com/paper/2201.02932/) (2022)
- [Why Batch Normalization Damage Federated Learning on Non-IID Data?](https://mlsanity.com/paper/2301.02982/) (2023)
- [FRAug: Tackling Federated Learning with Non-IID Features via Representation Augmentation](https://mlsanity.com/paper/2205.14900/) (2022)
- [Fed-ensemble: Improving Generalization through Model Ensembling in Federated Learning](https://mlsanity.com/paper/2107.10663/) (2021)
- [Convergence Analysis of Sequential Federated Learning on Heterogeneous Data](https://mlsanity.com/paper/2311.03154/) (2023)
- [Optimizing the Collaboration Structure in Cross-Silo Federated Learning](https://mlsanity.com/paper/2306.06508/) (2023)
- [The Cost of Training Machine Learning Models over Distributed Data Sources](https://mlsanity.com/paper/2209.07124/) (2022)
- [Feature Matching Data Synthesis for Non-IID Federated Learning](https://mlsanity.com/paper/2308.04761/) (2023)
- [Coordinating Momenta for Cross-silo Federated Learning](https://mlsanity.com/paper/2102.03970/) (2021)

## Abstract (arXiv)

Federated Learning (FL) on non-independently and identically distributed (non-IID) data remains a critical challenge, as existing approaches struggle with severe data heterogeneity. Current methods primarily address symptoms of non-IID by applying incremental adjustments to Federated Averaging (FedAvg), rather than directly resolving its inherent design limitations. Consequently, performance significantly deteriorates under highly heterogeneous conditions, as the fundamental issue of imbalanced exposure to diverse class and feature distributions remains unresolved. This paper introduces Stratify, a novel FL framework designed to systematically manage class and feature distributions throughout training, effectively tackling the root cause of non-IID challenges. Inspired by classical stratified sampling, our approach employs a Stratified Label Schedule (SLS) to ensure balanced exposure across labels, significantly reducing bias and variance in aggregated gradients. Complementing SLS, we propose a label-aware client selection strategy, restricting participation exclusively to clients possessing data relevant to scheduled labels. Additionally, Stratify incorporates a fine-grained, high-frequency update scheme, accelerating convergence and further mitigating data heterogeneity. To uphold privacy, we implement a secure client selection protocol leveraging homomorphic encryption, enabling precise global label statistics without disclosing sensitive client information. Extensive evaluations on MNIST, CIFAR-10, CIFAR-100, Tiny-ImageNet, COVTYPE, PACS, and Digits-DG demonstrate that Stratify attains performance comparable to IID baselines, accelerates convergence, and reduces client-side computation compared to state-of-the-art methods, underscoring its practical effectiveness in realistic federated learning scenarios.