Geographic Stratified Sampling for Client Cohorts

Proportional sampling from a geospatial fleet reproduces the fleet’s imbalance in the model: 80% of devices are in cities, so 80% of every cohort is urban, and the model learns urban roads well and rural ones badly. Full stratification fixes the coverage and breaks the battery budget of the sparse strata, which get called every round. The workable design is between the two, and it needs its sampling weights, its per-client participation caps and its privacy accounting worked out together. This page provides them, under client selection algorithms in Federated Learning Workflows for Geospatial Data.

Parameter Configuration and Calibration

  • strata — the partition clients are sampled within. Use a stable geographic partition — administrative regions, or H3 cells at a coarse resolution — fixed in advance. Strata derived from the current model’s error are adaptive and leak; strata derived from geography do not.
  • allocation — how the cohort is divided across strata. Three options: proportional (imbalance preserved), equal (sparse strata over-sampled), and square-root or Neyman allocation (proportional to Ns\sqrt{N_s} or to NsσsN_s \sigma_s). Square-root allocation is the practical default; Neyman is better when per-stratum variance is measurable and stable.
  • participation_cap — how often one client may be selected. A stratum with 18 clients and an allocation of 16 per round calls almost everyone every round, which exhausts both battery and privacy budget. Cap participation per client per epoch and let the stratum’s effective allocation fall to what the cap allows.
  • min_stratum_cohort — the floor below which a stratum is merged. A stratum contributing three clients to a round is a disclosure risk in every downstream metric. Merge it with a neighbour for sampling and for reporting, and record the merge.
Allocation Rural share of an 80-client cohort Rural client called every… Comment
Proportional 1.4 13 rounds model stays urban
Square-root 6.8 2.6 rounds usual compromise
Equal 16 every round exhausts the stratum
Neyman (σ-weighted) 9.1 2.0 rounds best when σ is known
Three allocations of an 80-client cohort A grouped bar chart of how many clients each of five density strata contributes under three allocation rules. Proportional allocation gives the remote stratum under two seats; equal allocation gives it sixteen, nearly its entire pool; square-root weighting gives about seven, which is a real voice without exhausting the stratum. Three allocations of an 80-client cohort node pool from 420 urban to 18 remote 0 20 40 clients selected 38.5 27.7 16.0 urban 23.8 21.8 16.0 suburban 11.0 14.8 16.0 exurban 5.0 10.0 16.0 rural 1.6 5.7 16.0 remote proportional √-weighted equal
Equal allocation calls almost the whole remote pool every round. √-weighting is the compromise that survives contact with battery and budget limits.

Reference Implementation

python
from __future__ import annotations

import math
import random
from dataclasses import dataclass, field
from typing import Mapping, Sequence


@dataclass
class StratumState:
    """Per-stratum population and the participation each client has spent."""

    clients: list[str]
    used: dict[str, int] = field(default_factory=dict)

    def eligible(self, cap: int) -> list[str]:
        return [c for c in self.clients if self.used.get(c, 0) < cap]


def allocate(
    populations: Mapping[str, int], cohort: int, *, mode: str = "sqrt",
    sigma: Mapping[str, float] | None = None,
) -> dict[str, int]:
    """Split a cohort across strata under the chosen allocation rule.

    Square-root allocation is the default because it interpolates between the two
    failures: proportional sampling ignores small strata, equal allocation exhausts
    them. Neyman allocation is better when per-stratum variance is measurable, since
    it puts effort where the estimate is noisiest rather than where the population is.
    """
    if mode == "proportional":
        weights = {s: float(n) for s, n in populations.items()}
    elif mode == "equal":
        weights = {s: 1.0 for s in populations}
    elif mode == "sqrt":
        weights = {s: math.sqrt(n) for s, n in populations.items()}
    elif mode == "neyman":
        if not sigma:
            raise ValueError("neyman allocation needs a per-stratum sigma")
        weights = {s: populations[s] * sigma.get(s, 1.0) for s in populations}
    else:
        raise ValueError(f"unknown allocation mode: {mode}")

    total = sum(weights.values())
    raw = {s: cohort * w / total for s, w in weights.items()}
    out = {s: int(v) for s, v in raw.items()}
    # Distribute the remainder to the largest fractional parts, deterministically.
    remainder = cohort - sum(out.values())
    for s, _ in sorted(raw.items(), key=lambda kv: -(kv[1] - int(kv[1])))[:remainder]:
        out[s] += 1
    return out


def select_cohort(
    strata: Mapping[str, StratumState],
    allocation: Mapping[str, int],
    *,
    participation_cap: int,
    min_stratum_cohort: int = 5,
    seed: int = 0,
) -> tuple[dict[str, list[str]], dict[str, str]]:
    """Sample within strata, honouring the per-client participation cap.

    Returns the selected clients per stratum and a note per stratum explaining any
    shortfall, because a silent shortfall is indistinguishable from a stratum that
    simply had nothing to contribute.
    """
    rng = random.Random(seed)
    chosen: dict[str, list[str]] = {}
    notes: dict[str, str] = {}
    for stratum, want in allocation.items():
        state = strata[stratum]
        pool = state.eligible(participation_cap)
        if len(pool) < min_stratum_cohort:
            notes[stratum] = "below the reporting floor — merge with a neighbour"
            continue
        take = min(want, len(pool))
        if take < want:
            notes[stratum] = f"capped: wanted {want}, {len(pool)} eligible"
        picked = rng.sample(pool, take)
        for client in picked:
            state.used[client] = state.used.get(client, 0) + 1
        chosen[stratum] = picked
    return chosen, notes


def coverage_score(
    chosen: Mapping[str, Sequence[str]], populations: Mapping[str, int]
) -> float:
    """1.0 when every stratum is represented in proportion to sqrt(population).

    A single scalar for "did this cohort look like the country?" — useful as a
    per-round health metric, since a drifting score is an early sign that a
    stratum's eligible pool is exhausting.
    """
    ideal = allocate(populations, sum(len(v) for v in chosen.values()), mode="sqrt")
    if not ideal:
        return 0.0
    err = sum(abs(len(chosen.get(s, [])) - ideal[s]) for s in ideal)
    return max(0.0, 1.0 - err / max(1, sum(ideal.values())))

Validation Checkpoint

python
def _validate() -> None:
    populations = {"urban": 420, "suburban": 260, "exurban": 120, "rural": 55, "remote": 18}

    prop = allocate(populations, 80, mode="proportional")
    sqrt = allocate(populations, 80, mode="sqrt")
    equal = allocate(populations, 80, mode="equal")

    # 1. Every allocation must spend the whole cohort, exactly.
    for alloc in (prop, sqrt, equal):
        assert sum(alloc.values()) == 80, alloc

    # 2. Square-root allocation must sit strictly between the two extremes.
    assert prop["remote"] < sqrt["remote"] < equal["remote"]
    assert prop["urban"] > sqrt["urban"] > equal["urban"]

    # 3. Equal allocation over-draws the smallest stratum — the failure it causes.
    assert equal["remote"] >= populations["remote"] * 0.8

    # 4. The participation cap must bind, and the shortfall must be reported.
    strata = {s: StratumState(clients=[f"{s}-{i}" for i in range(n)])
              for s, n in populations.items()}
    chosen, notes = select_cohort(strata, sqrt, participation_cap=1, seed=1)
    again, notes2 = select_cohort(strata, sqrt, participation_cap=1, seed=2)
    assert set(chosen.get("remote", [])) & set(again.get("remote", [])) == set()
    assert "remote" in notes2 or len(again.get("remote", [])) < sqrt["remote"]

    # 5. A stratum below the reporting floor must be skipped with a note, not sampled.
    tiny = {"islet": StratumState(clients=["a", "b"])}
    picked, why = select_cohort(tiny, {"islet": 2}, participation_cap=3)
    assert picked == {} and "islet" in why

    # 6. Coverage must be high for a square-root cohort and low for a proportional one.
    fresh = {s: StratumState(clients=[f"{s}-{i}" for i in range(n)])
             for s, n in populations.items()}
    good, _ = select_cohort(fresh, sqrt, participation_cap=10, seed=3)
    assert coverage_score(good, populations) > 0.9

    print("stratified selection: all assertions passed")


_validate()
A small stratum's eligible pool, round by round A descending staircase of the eligible client pool in an eighteen-client stratum across twelve rounds, under a cap of one participation per client per epoch. The pool is exhausted after the third round, so from round four onward the stratum contributes nothing and the model stops seeing it for the rest of the epoch. A small stratum's eligible pool, round by round 18 clients, cap of one participation per epoch, 7 wanted per round 2.5 5 7.5 10 0 5 10 round eligible clients remaining wanted per round
The stratum goes silent from round 4, and nothing in the aggregate says so. Alert on unfilled allocations, or merge the stratum.

Assertion 4 encodes the tension the whole design exists to manage: with a participation cap of one per epoch, the remote stratum simply cannot fill its allocation twice in a row. The right response is to see that in the notes and decide — merge strata, extend the epoch, or accept a smaller cohort — rather than to have the sampler quietly return fewer clients.

Incident Response and Edge Cases

  • A stratum’s eligible pool empties mid-epoch. Its allocation goes unfilled and the model stops seeing it. Either the cap is too tight for the allocation, or the stratum is too small to be a stratum. Merge it, and record the merge in the model card so nobody claims coverage it does not have.
  • Coverage score drifts downward over an epoch. Normal as caps bind, and worth alerting on when it crosses a threshold: it means later rounds in the epoch are systematically more urban than earlier ones, which introduces a within-epoch bias that is easy to mistake for drift.
  • A stratum is over-represented and its metrics look excellent. Over-sampling reduces that stratum’s metric variance, which can read as “that region is doing well”. Compare metrics against each region’s own error bar, not against each other’s point estimates.
  • Selection is deterministic across rounds. A fixed seed reuses the same clients, which is a privacy problem as well as a statistical one. Seed from the round id, and verify that consecutive rounds have low overlap in small strata.
  • Someone proposes stratifying by model error. That makes the strata a function of the data and turns selection into an adaptive, leaking procedure. If error-driven targeting is genuinely needed, derive it from published DP metrics and accept the accounting.
Coverage and rural error move together A grouped bar chart of the coverage score and the resulting rural model error for four selection policies. Proportional selection scores worst on both. Coverage-aware selection scores best on both. The two series move together, showing that better coverage is not a fairness cost paid out of accuracy but a route to it. Coverage and rural error move together same cohort size, same budget, four selection policies 0 0.5 1 score / error 0.42 0.31 proportional 0.91 0.14 √-weighted 0.78 0.11 equal 0.96 0.09 coverage-aware coverage score rural error
Better coverage is not charity paid for out of accuracy. The excluded distribution keeps pulling the model back, so covering it converges faster.

Frequently Asked Questions

Why square-root allocation rather than equal?

Equal allocation calls almost every client in a small stratum every round, exhausting battery and per-client privacy budget within days. Square-root allocation gives sparse strata a much larger share than their population would justify while keeping per-client participation sustainable — it is the compromise that survives contact with a real fleet.

Does the cohort composition itself need protecting?

Yes. "This round drew 16 clients from the remote stratum" is a statement about who is where, and publishing cohort composition per round at fine granularity is a slow leak. Coarsen it, publish it on a delay, or report it only at the level of merged regions.

How does the participation cap interact with the privacy budget?

Directly: each participation is a release against that client's budget, so the cap and the per-round epsilon together determine the per-epoch guarantee. Choosing the cap for battery reasons and the epsilon for privacy reasons independently is how deployments end up with a per-client epsilon nobody can state.

Should strata follow administrative or density boundaries?

Density boundaries usually model better, since the thing that varies is the data distribution rather than the jurisdiction. Administrative boundaries report better, because that is how obligations and stakeholders are organised. Many deployments stratify on density and report on administrative regions, which is fine as long as the mapping between them is fixed and published.

Up one level: Client Selection Algorithms · Section: Federated Learning Workflows for Geospatial Data.