📊 statistics

A Hierarchical Sampling Framework for bounding the Generalization Error of Federated Learning

This paper proposes a hierarchical sampling procedure for Federated Learning that derives generalization bounds via Wasserstein distance and supersample construction, demonstrating that these bounds strictly improve upon existing results on conditional mutual information and precisely capture the asymptotic error rates in Gaussian models.

Original authors: Dario Filatrella, Ragnar Thobaben, Mikael Skoglund

Published 2026-05-06

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Dario Filatrella, Ragnar Thobaben, Mikael Skoglund

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: Training a Team Without Sharing Secrets

Imagine you are trying to teach a robot to recognize cats. In a normal classroom, you would gather thousands of cat photos from the internet, show them all to the robot at once, and let it learn. This is centralized learning.

But what if these photos belong to different people who do not want to share them? Perhaps they are private, or the internet connection is too slow to send them all to one place. This is Federated Learning (FL). Instead of sending the photos, the robot sends its "brain" (the model) to each person's computer. The computer learns from its own photos and sends back only the changes to the brain, not the photos themselves.

This paper addresses a specific, chaotic version of this problem called Hierarchical Federated Learning (HFL). Imagine that the people are not just individuals; they are organized in a family tree.

Level 1: The whole world (Global).
Level 2: Countries.
Level 3: Cities.
Level 4: Neighborhoods.
Level 5: Individual houses (the actual data).

The data in one neighborhood resembles that of other houses in that neighborhood but differs from that of a house in another city. This creates a "tree" of dependencies. The authors wanted to answer a simple question: How well will this robot actually learn from this chaotic, tree-like structure?

The Problem: Measuring "Generalization"

In machine learning, "generalization" is the ability to perform well on new data it has not seen yet.

The Risk: If the robot memorizes the specific cats in the training photos, it might fail when it sees a new cat.
The Goal: We want a mathematical guarantee (a bound) that states: "The robot's performance on new data will not be much worse than on the training data."

Previous methods tried to measure this with simple math but often ignored the "tree" structure of the data. They treated the data like a random pile of sand, overlooking the fact that data from the same city are interconnected. This paper says: "Let's build a ruler that actually fits the shape of the tree."

The Solution: A "Ghost" Tree and a New Ruler

The authors introduce two main tools to measure this error:

1. The "Ghost" Tree (Supersample Construction)

Imagine you are testing a student's knowledge. Instead of giving them just one test, you give them a "ghost test" that is almost identical to the real one but with a tiny deviation (like swapping one question).

The authors build a ghost tree next to the real data tree.
They create pairs of nodes: a "Real" node and a "Ghost" node.
For every branch of the tree, they flip a coin to decide whether the algorithm learns from the Real data or the Ghost data.
By comparing how much the robot's brain changes when it swaps a Real node for a Ghost node, they can measure how sensitive the robot is to specific data points. If the robot wildly changes its mind over a tiny swap, it is overfitted (memorized). If it stays calm, it is learning well.

2. The "Wasserstein Distance" (The Elastic Ruler)

To measure the difference between the robot's "Real Brain" and its "Ghost Brain," the authors use a metric called Wasserstein distance.

The Analogy: Imagine you have a pile of dirt (Real Brain) and want to move it so that it matches a pile of dirt in a different shape (Ghost Brain).
Old Rulers (Mutual Information): These counted how many dirt clods were different. They are good but can be too strict or too loose.
The Wasserstein Ruler: This measures the effort required to move the dirt. It considers the shape and geometry of the data. It asks: "How far do I have to push this specific dirt clod so that the piles match?"
Since this ruler understands the "shape" of the data distribution, it provides a tighter, more accurate estimate of the error, especially when data is limited (a bound on how large the errors can be).

What They Found

A Better Formula: They derived a new mathematical formula that calculates the maximum possible error. This formula works for the entire tree structure, not just flat data.
It is Tighter: They proved that their new "elastic ruler" (Wasserstein) provides a stricter, more accurate bound for the error than the old "dirt clod counter" methods (conditional mutual information), especially when errors are bounded in size.
Privacy Works: They showed that if you add "noise" to the data to protect privacy (Differential Privacy), their formula still works and can predict how much this privacy noise will affect learning accuracy.
The Test Case (Gaussian Location Model): They tested their math on a specific, simple scenario (the Gaussian location model) where they knew the exact answer.
- Result: Their formula was very close to the true answer. It correctly predicted how the error grows as you add more levels to the tree, although it slightly overestimated the error related to the depth of the tree.

The Conclusion

This paper is like building a better map for a complex, multi-level city. Previous maps treated the city as a flat grid, causing people to get lost. The authors built a map that respects the skyscrapers and subway tunnels (the hierarchy).

By using a "ghost tree" to test sensitivity and a "Wasserstein ruler" to measure distances, they created a more reliable method to predict how well a Federated Learning system will perform. This helps engineers know exactly how much trust to place in a model trained across a complex, hierarchical network of devices without having to see the private data.

Technical Conclusion: A Hierarchical Sampling Framework for Bounding the Generalization Error of Federated Learning

Problem Description
Federated Learning (FL) enables multiple clients to jointly train a global model while keeping data local. While generalization bounds exist for centralized supervised learning, their direct application to FL often fails to leverage the specific structure of the FL setup, particularly the non-i.i.d. nature of data resulting from hierarchical sampling strategies. Existing literature lacks a unified framework encompassing various hierarchical sampling strategies and communication structures; it frequently relies on assumptions that conflict with the distributed nature of FL or fails to account for dependencies induced by hierarchical data sources.

Methodology
The authors propose a generalized framework for Hierarchical Federated Learning (HFL) modeled as a multi-layered tree structure. In this setup:

Hierarchical Sampling: The dataset is indexed by a tree with fixed topology. The root represents a meta-distribution $D$ . Nodes in intermediate layers represent clients or clusters (distributions over subordinate nodes), and leaf nodes represent actual data points. Sampling occurs recursively from the root to the leaves, inducing dependencies between client datasets.
Construction of Supersamples: To analyze generalization, the authors adapt the "supersample" technique (originally by Steinke et al., 2020). They construct a "ghost" dataset $\bar{\mu}$ with identical statistical properties to the original dataset $\mu$ . For each node in the tree, a binary variable $U$ selects either the original node or its ghost copy. This enables measuring the algorithm's sensitivity to changes in a single node within the sampling tree.
Wasserstein Distance: The generalization error is bounded using the Wasserstein-1 distance ( $W$ ) between the distribution of the output hypothesis, conditioned on a specific node selection, and the unconditional distribution. This approach relies on the assumption that the loss function $\ell(w, z)$ is $L$ -Lipschitz continuous with respect to the hypothesis $w$ .

Main Contributions

Generalized HFL Framework: The work presents a framework that encompasses various hierarchical sampling strategies and prior FL frameworks as special cases, explicitly modeling dependencies induced by the tree structure.
Wasserstein Generalization Bound: The authors derive a new generalization bound in terms of Wasserstein distance. This bound applies to a broad range of loss functions (under the Lipschitz assumption) and proves strictly tighter than existing state-of-the-art conditions for conditional mutual information (CMI) for bounded losses.
Privacy-Aware Bounds: The framework is extended to incorporate assumptions regarding Differential Privacy (DP). By imposing local privacy constraints on aggregation steps within the tree, the authors derive generalization bounds based on algorithmic privacy.
Case Study on Gaussian Location Model (GLM): The tightness of the derived bounds is evaluated using a Gaussian location model. The authors demonstrate that their bound recovers the correct asymptotic rate of the generalization error with respect to the number of data points, although it exhibits a scaling deviation regarding the depth of the sampling tree compared to the true error.

Results

Main Theorem (Theorem 5): Under the Lipschitz assumption, the expected generalization error is bounded by a sum over tree layers of the expected Wasserstein distance between the hypothesis distribution, conditioned on a selected node, and the unconditional hypothesis distribution.
Relationship to CMI (Corollary 6 & 9): For bounded losses, the Wasserstein bound implies existing CMI bounds. The authors show that their derivation avoids the looseness introduced by chains of inequalities (such as the repeatedly applied Pinsker inequality) in previous CMI-based proofs, resulting in a strictly tighter bound.
Implications for Privacy (Theorem 13): Under $\epsilon$ -differential privacy assumptions for local aggregation, the generalization error is bounded by a function of the privacy parameters $\epsilon_l$ across layers, specifically involving $\sqrt{\min(\epsilon_l, \epsilon_l(e^{\epsilon_l}-1))}$ .
GLM Analysis: In the Gaussian location model, the true generalization error scales as $\sqrt{\frac{\sum \sigma_l^2/N_l}{\sum \sigma_l^2}}$ . The derived Wasserstein bound scales as $\sum \frac{\sigma_l}{N_l}$ . While the bound captures scaling with the number of data points ( $N_l$ ), the authors note that it fails to capture the specific scaling with tree depth ( $L$ ) observed in the true error, particularly when variances are homogeneous across layers.

Significance and Claims
The work claims to offer a more general and tighter theoretical foundation for understanding generalization in Federated Learning by explicitly accounting for hierarchical data dependencies. By leveraging Wasserstein distance, the framework captures the geometry of distributions and the sensitivity of the learning algorithm to the sampling structure more effectively than previous information-theoretic approaches. The authors claim that their bound strictly implies existing CMI bounds for bounded losses, thus representing a refinement rather than merely an alternative. Furthermore, the framework's ability to seamlessly integrate with differential privacy assumptions demonstrates its utility in analyzing privacy-preserving FL systems. The authors acknowledge limitations, noting that the current bound does not explicitly account for communication constraints or computational constraints of individual nodes, but suggest these can be integrated via independence conditions, as shown in their derivation of Corollary 10. Future work is proposed to extend the framework to general (non-fixed) tree structures.