Pith. sign in

REVIEW 5 major objections 5 minor 23 references

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

T0 review · 5 major / 5 minor · reviewed 2026-08-16 · deepseek-v4-flash

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗ Review JSON

Pith's one-line read Stratify claims that a Stratified Label Schedule, which balances and shuffles label exposure during federated training, restores accuracy close to IID baselines under non-IID data and converges in far fewer rounds.

desk verdict A genuinely new label-balancing schedule, but the baselines are run in a protocol that strips their mechanisms, so the headline gains are not credible evidence of superiority. read the letter →

arxiv 2504.13462 v1 pith:UBMNY2GI submitted 2025-04-18 cs.LG

Hui Yeok Wong , Chee Kau Lim , Chee Seng Chan This is my paper · ORCID
classification cs.LG
keywords federatedlearningnon-IIDdatastratifiedlabelschedulelabel-awareclientselectionheterogeneousskewfeaturehomomorphicencryption
verification ladder T0 review T1 audit T2 compute T3 formal

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

Federated learning degrades when client data are non-IID because the global model sees labels and features unevenly. Stratify claims to fix the root cause rather than patch the aggregation step: the server builds a Stratified Label Schedule that repeats each label a controlled number of times, selects only clients holding the scheduled label, and applies fine-grained updates one sample or small batch at a time. The paper argues this produces an unbiased gradient estimator with lower variance, and reports accuracy close to the IID baseline across label-skew and feature-skew benchmarks while converging in many fewer rounds. A sympathetic reader would care because it offers a way to make federated learning robust to extreme heterogeneity without asking clients to share data.

What carries the argument

The central mechanism is the Stratified Label Schedule (SLS): a shuffled list in which every label $l$ appears $f_l$ times, with $f_l$ either uniform or proportional to a capped global label count, so the global model is exposed to all classes at controlled frequencies. Around it the paper places two supporting mechanisms: label-aware client selection, which restricts each training step to clients whose local data contain the scheduled label, and a fine-grained update scheme that applies one sample or a small batch per global update instead of multiple local epochs. The theoretical claim is carried by a variance-decomposition identity: writing the aggregated gradient as $\hat{g} = \frac{1}{|\mathrm{SLS}|} \sum_{l \in \mathrm{SLS}} g_l$, balanced repetition makes the per-label contribution probability $P(l) = f_l / \sum_l f_l$ uniform, so $\mathbb{E}[\hat{g}]$ is unbiased and the cross-label variance term $\mathrm{Var}(\mathbb{E}[\hat{g} \mid l])$ is minimized.

What would settle it

Run a label-skew experiment, for example CIFAR-10 with one class per client, where FedAvg, FedProx, and SCAFFOLD are allowed their standard multiple local epochs per round and the same total communication budget as Stratify; if one of them matches or exceeds Stratify's accuracy, the claim that balanced label scheduling is what drives the gains is undercut. Alternatively, measure the empirical variance of the aggregated gradient under SLS versus random client sampling; finding no variance reduction would directly contradict the paper's theoretical claim.

Watch

Extended reading notes

Core claim

On the paper's own terms, the discovery is that the imbalance that breaks federated learning is not primarily in the aggregation rule but in the model's exposure: when a global model is trained in rounds where each client optimizes its own skewed local data, the aggregate update drifts. Stratify's response is to turn training into a server-orchestrated schedule of labels. The SLS is built by repeating each label a controlled number of times and shuffling; at each step the server assigns the next label only to clients that actually hold that label, and the model is updated immediately on that label's data rather than after full local epochs. The paper claims this yields an unbiased gradient estimate, reduces the between-label variance term of the law of total variance, and lets each label contribute equally to aggregated updates. Empirically it reports that Stratify reaches accuracy comparable to its own IID runs on MNIST, CIFAR-10, CIFAR-100, Tiny-ImageNet, COVTYPE, PACS, and Digits-DG, often in a small fraction of the baselines' rounds.

Load-bearing premise

The load-bearing premise is that the baselines can be fairly compared when they are forced into the same single-sample or batch-per-iteration update protocol that Stratify uses, even though FedAvg, FedProx, SCAFFOLD, and SFL were designed for multiple local epochs over each client's full dataset; if that protocol constraint is removed, the reported accuracy gap may shrink or reverse.

Editorial extensions

If this is right

  • Under label skew as severe as one class per client, Stratify reports accuracy comparable to the IID setting, while FedAvg, FedProx, SCAFFOLD, and SFL degrade sharply.
  • Because the schedule controls what the global model sees and when, the method converges in far fewer rounds, for example 8 rounds on CIFAR-10 batch learning versus 50 for FedAvg, lowering cumulative local training time per client.
  • The label-aware selection mechanism lets clients with only a few samples contribute meaningfully, since each global update needs only data for the scheduled label.
  • For feature skew, weighted client selection prevents one large-domain client from dominating later iterations, while uniform selection preserves label privacy.
  • With Batch Normalization, Stratify aggregates sum-based batch statistics rather than averaging local statistics, avoiding the mismatch that hurts FedTAN-style BN in non-IID settings.

Reading between the lines

Editorial extensions of the paper, not claims the author makes directly.

  • The authors leave implicit that the SLS mechanism does not depend on homomorphic encryption: any protocol that yields aggregate label counts, including trusted-server aggregation, would preserve the core balancing effect; the encryption layer is an implementation choice.
  • A testable extension would apply the same label-scheduling idea to centralized or asynchronous training under class imbalance, where a server already knows global label frequencies and no privacy machinery is needed.
  • The single-sample sequential mode is effectively a continual-learning protocol, so a natural extension is a head-to-head comparison with replay-based continual-learning methods under distribution shift; the paper does not evaluate that regime.
Share X Bluesky LinkedIn Reddit HN

Signed reviews

No signed human review yet.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, and a circularity audit.

Referee Report

5 major / 5 minor

Summary. The paper proposes Stratify, a federated learning framework that replaces the usual multi-epoch local training protocol with a Stratified Label Schedule (SLS), label-aware client selection, and fine-grained single-sample or batch-data global updates. The authors claim that this balanced exposure to labels yields an unbiased, lower-variance gradient estimator, that a secure client selection protocol based on CKKS homomorphic encryption preserves privacy, and that Stratify matches IID performance, converges faster, and reduces client-side computation relative to FedAvg, FedProx, SCAFFOLD, and SFL on MNIST, CIFAR-10, CIFAR-100, Tiny-ImageNet, COVTYPE, PACS, and Digits-DG. The central evidence is empirical, presented in Tables I and II, where all baselines are evaluated under Stratify's update protocol.

Significance. If the claims were validated, the idea of explicitly balancing label exposure through a stratified schedule could be a useful addition to the federated learning toolbox, especially for extreme label skew, and the paper includes concrete algorithms and ablations for client selection and chunk size. However, the current evaluation does not support the headline claims. The baselines are forced into a protocol they were not designed for, the theoretical argument is largely definitional and does not establish the claimed variance reduction, and the reported numbers are selected optimistically without standard deviations. The privacy contribution is also narrower than stated, since secure aggregation is explicitly left to future work. Strengths of the manuscript include the detailed algorithm pseudocode and the attempt to analyze communication and implementation overhead, but these do not compensate for the load-bearing problems in the comparison and theory.

major comments (5)
  1. [Section V, Tables I and II] The comparison is not apples-to-apples: FedAvg, FedProx, SCAFFOLD, and SFL are evaluated under the proposed single-sample or batch-per-iteration protocol, whereas these methods are defined for multiple local epochs over each client's full local dataset. Under batch size 1, SCAFFOLD's control variates cannot be meaningfully updated, FedProx's proximal term degenerates, and FedAvg's aggregation reduces to averaging single-sample stochastic gradients; the reported gaps (e.g., CIFAR-10 #C=1: 81.18% vs 11.56% in Table I) may therefore be artifacts of disabling the baselines' core mechanisms. A fair evaluation must either use the baselines' native protocols or justify why the modified protocol is the correct one.
  2. [Section III-B] The 'unbiased and balanced gradient estimator' claim is definitional: if f_l is uniform, the estimator is an equally weighted average over labels, so P(l)=1/|L| by construction. But the standard FL objective is the sample-average empirical risk; an equally weighted label average is unbiased only under balanced class sizes or a deliberately reweighted objective, neither of which is stated. The variance-reduction argument is also incomplete: the law of total variance decomposition does not show that uniform f_l minimizes Var(E[g|l]); classical stratified sampling would call for allocation proportional to stratum sizes and variances (Neyman allocation). The paper should compare against that baseline and state the objective being estimated.
  3. [Section V, Tables I, II, and V] The reported numbers are selected optimistically: the epoch with highest test accuracy is chosen for every method, and for feature skew the better of uniform and weighted client selection is reported for Stratify (Section V). No standard deviations or number of seeds are given. This can easily explain small margins (e.g., PACS #D=2 single-sample: 95.35% vs 95.61% for SFL). Please report mean ± std over multiple runs under a single pre-specified selection rule.
  4. [Section V-E, Table III] The efficiency comparison is internally inconsistent. Table III shows Stratify's per-round cost is two to five times larger (e.g., 145s vs 53s for single-sample FedAvg), while the text claims a per-client total of 199s for Stratify vs 575s for FedAvg; these figures do not follow from the table. A valid comparison should report end-to-end training time to a fixed accuracy target, including serialization and model-passing overhead.
  5. [Section IV-A and Algorithm 4] The privacy claim is stronger than what is implemented. The secure client selection protocol protects label identities behind placeholders, but the server obtains decrypted global counts N(p) for each placeholder, and secure aggregation of gradients is explicitly deferred ('We leave the exploration of security aspects as an important direction for future work' near Algorithm 4). The text should state that only client selection is privacy-preserving, not the full training procedure.
minor comments (5)
  1. [Section VI] There is a typo in the Conclusion: 'Fedvg' should be 'FedAvg'.
  2. [Section V-F, Table V] The claim that 'chunk size 5 results in around 35% reduction in model transfer frequency' should define the baseline for that percentage and clarify how the reduction is computed.
  3. [Section V] The sentence 'The algorithms are trained for a maximum epoch took in their respective IID case' is unclear; please clarify whether E# denotes an epoch or a communication round and how the maximum is determined.
  4. [Section IV-B] The communication overhead formula T=R×Cupdate×ffreq assumes a single client; please clarify how it applies to the sequential single-sample protocol where model passing occurs between clients.
  5. [Table IV] The units in Table IV mix seconds and milliseconds in the header and body; please make the units consistent and define the timing conventions clearly.

Circularity Check

1 steps flagged · score 6.0 of 10

Unbiased-gradient "theorem" reduces to the definition of SLS; empirical results are independent but the theoretical claim is self-definitional.

  1. self definitional [Section III-A (Stratified Label Schedule) and Section III-B (Theoretical Analysis), equations defining SLS and P(l)]
    "We define the Stratified Label Schedule (SLS) as SLS = Shuffle [l∈L f_l{l}] ... By enforcing a uniform frequency across all labels l∈L, P(l) becomes uniform. As a result, each label contributes equally to the aggregated gradient. This confirms that SLS yields an unbiased and balanced gradient estimator by constructing an equal representation across all labels."

    The SLS is constructed by choosing repetition counts f_l, and P(l) is defined as f_l/Σf_l. 'Enforcing a uniform frequency' is just choosing f_l equal, which makes P(l) uniform by definition. The conclusion that each label contributes equally is identical to the construction of the schedule. The estimator is 'unbiased' only with respect to the artificial uniform label distribution inserted by the authors; no independent target distribution is estimated. Hence the claimed theorem is an analytic restatement of the definition, not a first-principles derivation.

full rationale

The only step that reduces to its inputs by construction is the unbiased-gradient claim in Section III-B. The SLS is defined as a shuffle over f_l repetitions, and the inclusion probability is defined as f_l divided by the sum of f_l. Setting f_l uniform then trivially yields uniform P(l), so 'balanced exposure' and 'unbiased balanced estimator' are restatements of the schedule definition rather than derived results. The variance-reduction remark is asserted from the law of total variance without comparing against an optimal stratified allocation; that is a proof gap, but not itself circular. The empirical performance claims are independent experiments, although the baseline comparison is questionable because FedAvg, FedProx, SCAFFOLD, and SFL are forced into Stratify's single-sample or batch-per-iteration protocol instead of their native local-epoch protocols; that is a fairness/comparability issue, not circularity. No load-bearing self-citations are present. Because the theoretical 'prediction' reduces by construction but the central experimental evaluation retains independent content, the overall circularity is partial, yielding a score of 6.

Assumptions & free parameters 4 free parameters · 3 assumptions · 0 invented entities

The method depends on several hyperparameters (f_l, chunkSize, batchSize) and on the choice of client selection strategy, none of which are derived from theory. The theoretical claim of unbiasedness is built into the definition of the estimator. No new physical entities are introduced.

free parameters (4)
  • Label frequency f_l in SLS = uniform or capped proportional
    Controls how many times each label appears in the schedule; chosen by the authors, not derived from theory.
  • chunkSize (CS) = 5, 50, 150 in ablation
    Controls grouping of labels for consecutive updates on the same client; the main experiments do not specify the value used.
  • batchSize = not reported
    Controls number of labels per global update in batch-data learning; no value is stated in the main experiments.
  • Client selection strategy (uniform vs weighted) = selected per setting, better reported
    For feature skew the authors report the better of two strategies, a post-hoc choice that inflates reported accuracy.
assumptions (3)
  • standard math Linearity of expectation and the law of total variance apply to the gradient estimator.
    Used in Section III-B; valid but do not by themselves prove variance reduction.
  • domain assumption The server can compute global label statistics without learning label identities.
    Required by the secure client selection protocol; the CKKS-based construction is approximate and the decryption procedure can reveal label differences, so this assumption is questionable.
  • ad hoc to paper A uniform label frequency f_l yields an unbiased estimator of the FL objective.
    The estimator is unbiased only with respect to the SLS sampling design, not the true data distribution; this is a modeling choice, not an established property.

how reviews work

0 comments
Cite this review

Pith. "Pith review of Stratify: Rethinking Federated Learning for Non-IID Data through Balanced Sampling." pith.science (2026). https://pith.science/paper/UBMNY2GI

@misc{pith2026250413462,
  author       = {Pith},
  title        = {Pith review of: Stratify: Rethinking Federated Learning for Non-IID Data through Balanced Sampling},
  year         = {2026},
  howpublished = {\url{https://pith.science/paper/UBMNY2GI}},
  note         = {Machine review of arXiv:2504.13462}
}
read the original abstract

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.

Figures

Figures reproduced from arXiv: 2504.13462 by the authors.

Figure 1
Figure 1. Overview of Stratify Training Process. (a) Single-sample learning: (a1) Send initialized global model parameter (only first client), masked labels to train, and next client address to current client, (a2) Convert masked label to real label, (a3) Update model sequentially on the required labels, (a4) Signal server to send next masked labels to train to next client, (a5) Send updated model to next client, Repeat step … view at source ↗
Figure 2
Figure 2. Comparison of algorithms’ performance with increasing client numbers on different datasets and data partitions in single-sample learning [PITH_FULL_IMAGE:figures/full_fig_p010_2.png] view at source ↗
Figure 3
Figure 3. Comparison of algorithms’ performance with increasing client numbers on different datasets and data partitions in batch-data learning [PITH_FULL_IMAGE:figures/full_fig_p011_3.png] view at source ↗
Figures from the paper (1 more)
Figure 4
Figure 4. Figure 4: FedTAN and Stratify performance on CIFAR-10 across various non￾IID settings selection, especially with a smaller chunk size. The batch learning setting, on the other hand, shows slower convergence with uniform selection, but still able to achieve accuracy close to weig…

Discussion (0). Continue with ORCID to comment.

Reference graph

Works this paper leans on

23 extracted references · 13 canonical work pages

  1. [1]

    The cost of training machine learning models over distributed data sources,

    E. Guerra, F. Wilhelmi, M. Miozzo, and P. Dini, “The cost of training machine learning models over distributed data sources,” IEEE Open Journal of the Communications Society , vol. 4, pp. 1111–1126, 2023

    work page 2023

  2. [2]

    Federated learning with hierarchical clustering of local updates to improve training on non-IID data,

    C. Briggs, Z. Fan, and P. Andras, “Federated learning with hierarchical clustering of local updates to improve training on non-IID data,” in 2020 International Joint Conference on Neural Networks (IJCNN) . IEEE, 2020, pp. 1–9

    2020

  3. [3]

    Federated learning under distributed concept drift,

    E. Jothimurugesan, K. Hsieh, J. Wang, G. Joshi, and P. B. Gibbons, “Federated learning under distributed concept drift,” in International Conference on Artificial Intelligence and Statistics . PMLR, 2023, pp. 5834–5853

    work page 2023

  4. [4]

    Efficient distribution similarity identification in clustered federated learning via principal angles between client data subspaces,

    S. Vahidian, M. Morafah, W. Wang, V . Kungurtsev, C. Chen, M. Shah, and B. Lin, “Efficient distribution similarity identification in clustered federated learning via principal angles between client data subspaces,” in Proceedings of the AAAI Conference on Artificial Intelligence , vol. 37, no. 8, 2023, pp. 10 043–10 052

    2023

  5. [5]

    Clustered federated multi-task learning with non-IID data,

    Y . Xiao, J. Shu, X. Jia, and H. Huang, “Clustered federated multi-task learning with non-IID data,” in2021 IEEE 27th International Conference on Parallel and Distributed Systems (ICPADS). IEEE, 2021, pp. 50–57

    work page 2021

  6. [6]

    Per- sonalized cross-silo federated learning on non-iid data,

    Y . Huang, L. Chu, Z. Zhou, L. Wang, J. Liu, J. Pei, and Y . Zhang, “Per- sonalized cross-silo federated learning on non-iid data,” in Proceedings of the AAAI conference on artificial intelligence , vol. 35, no. 9, 2021, pp. 7865–7873

    2021

  7. [7]

    Adapt to adaptation: Learning personalization for cross-silo federated learning,

    J. Luo and S. Wu, “Adapt to adaptation: Learning personalization for cross-silo federated learning,” in IJCAI: proceedings of the conference , vol. 2022. NIH Public Access, 2022, p. 2166

    2022

  8. [8]

    FRAug: Tackling federated learning with Non-IID features via representation augmentation,

    H. Chen, A. Frikha, D. Krompass, J. Gu, and V . Tresp, “FRAug: Tackling federated learning with Non-IID features via representation augmentation,” in Proceedings of the IEEE/CVF International Confer- ence on Computer Vision , 2023, pp. 4849–4859

    work page 2023

Show all 23 references
  1. [9]

    Feature matching data synthesis for non-iid federated learning,

    Z. Li, Y . Sun, J. Shao, Y . Mao, J. H. Wang, and J. Zhang, “Feature matching data synthesis for non-iid federated learning,” IEEE Transac- tions on Mobile Computing , vol. 23, no. 10, pp. 9352–9367, 2024

    2024

  2. [10]

    Federated optimization in heterogeneous networks,

    T. Li, A. K. Sahu, M. Zaheer, M. Sanjabi, A. Talwalkar, and V . Smith, “Federated optimization in heterogeneous networks,” Proceedings of Machine learning and systems , vol. 2, pp. 429–450, 2020

    2020

  3. [11]

    Coordinating momenta for cross-silo federated learning,

    A. Xu and H. Huang, “Coordinating momenta for cross-silo federated learning,” in Proceedings of the AAAI Conference on Artificial Intelli- gence, vol. 36, no. 8, 2022, pp. 8735–8743

    2022

  4. [12]

    Scaffold: Stochastic controlled averaging for federated learn- ing,

    S. P. Karimireddy, S. Kale, M. Mohri, S. Reddi, S. Stich, and A. T. Suresh, “Scaffold: Stochastic controlled averaging for federated learn- ing,” in International conference on machine learning . PMLR, 2020, pp. 5132–5143

    2020

  5. [13]

    Ensemble federated learning with non-iid data in wireless networks,

    Z. Zhao, J. Wang, W. Hong, T. Q. Quek, Z. Ding, and M. Peng, “Ensemble federated learning with non-iid data in wireless networks,” IEEE Transactions on Wireless Communications , vol. 23, no. 4, pp. 3557–3571, 2023

    2023

  6. [14]

    Fed-ensemble: Ensemble models in federated learning for improved generalization and uncertainty quantification,

    N. Shi, F. Lai, R. Al Kontar, and M. Chowdhury, “Fed-ensemble: Ensemble models in federated learning for improved generalization and uncertainty quantification,” IEEE Transactions on Automation Science and Engineering, 2023

    2023

  7. [15]

    Communication-efficient learning of deep networks from decentralized data,

    B. McMahan, E. Moore, D. Ramage, S. Hampson, and B. A. y Arcas, “Communication-efficient learning of deep networks from decentralized data,” in Artificial intelligence and statistics . PMLR, 2017, pp. 1273– 1282

    2017

  8. [16]

    S. K. Thompson, Sampling. John Wiley & Sons, 2012, vol. 755

    2012

  9. [17]

    J. D. Hamilton, Time series analysis. Princeton university press, 2020

    2020

  10. [18]

    Optimizing the collaboration structure in cross-silo federated learning,

    W. Bao, H. Wang, J. Wu, and J. He, “Optimizing the collaboration structure in cross-silo federated learning,” in International Conference on Machine Learning . PMLR, 2023, pp. 1718–1736

    2023

  11. [19]

    Clustered sampling: Low-variance and improved representativity for clients selection in federated learning,

    Y . Fraboni, R. Vidal, L. Kameni, and M. Lorenzi, “Clustered sampling: Low-variance and improved representativity for clients selection in federated learning,” in International Conference on Machine Learning . PMLR, 2021, pp. 3407–3416

    2021

  12. [20]

    Fedsts: A stratified client selection framework for consistently fast federated learning,

    D. Gao, D. Song, G. Shen, X. Cai, L. Yang, G. Liu, X. Li, and Z. Wang, “Fedsts: A stratified client selection framework for consistently fast federated learning,” IEEE Transactions on Neural Networks and Learning Systems, 2024

    2024

  13. [21]

    Convergence analysis of sequential federated learning on heterogeneous data,

    Y . Li and X. Lyu, “Convergence analysis of sequential federated learning on heterogeneous data,” Advances in Neural Information Processing Systems, vol. 36, pp. 56 700–56 755, 2023

    2023

  14. [22]

    Why batch normalization damage federated learning on non-iid data?

    Y . Wang, Q. Shi, and T.-H. Chang, “Why batch normalization damage federated learning on non-iid data?” IEEE transactions on neural networks and learning systems , 2023

    2023

  15. [23]

    Federated learning on non-iid data silos: An experimental study,

    Q. Li, Y . Diao, Q. Chen, and B. He, “Federated learning on non-iid data silos: An experimental study,” in 2022 IEEE 38th international conference on data engineering (ICDE) . IEEE, 2022, pp. 965–978

    2022

Pith tools

Reviewed August 16, 2026 · model on record in the stance chip above.