Cohort Selection

Cohort selection operates on the principle of group-based training, where clients are not selected individually but as pre-defined or dynamically formed cohorts. A cohort is a set of clients that train together and have their model updates aggregated before contributing to the global model. This is distinct from per-client selection strategies like random selection or power-of-choice.

• Core Mechanism: The server selects a cohort (or multiple cohorts) for a training round. Each client within the selected cohort performs local training, and their updates are aggregated (e.g., via federated averaging) at a cohort level before being sent to the global server.
• Typical Use: Common in cross-silo federated learning, where clients are large, stable organizations (e.g., hospitals, banks) that can be logically grouped by domain, geography, or data type.

This method provides inherent structural benefits for privacy-preserving machine learning. By treating a cohort as a single logical participant, it adds a layer of aggregation and obfuscation.

• Privacy Amplification: Individual client updates are aggregated within the cohort, making it harder to isolate any single client's contribution. This can complement formal techniques like secure aggregation and differential privacy.
• Trust Boundaries: Cohorts can be formed within trusted enclaves or secure multi-party computation (MPC) clusters, allowing for more sensitive intra-cohort aggregation before a less-trusted global server sees the result.
• Mitigates Inference Attacks: The cohort aggregate acts as a buffer, reducing the risk of model inversion or membership inference attacks targeting a specific device's data.

Cohort selection is a key tool for managing system scalability in federated networks with thousands or millions of potential clients (cross-device FL).

• Reduced Orchestration Overhead: The server coordinates with a manageable number of cohort heads or representatives instead of every individual device, drastically cutting communication overhead.
• Hierarchical Aggregation: Enables a natural hierarchical federated learning architecture. Updates are aggregated at the cohort level, and then cohort-level updates are aggregated globally. This mirrors organizational or network topology.
• Handles Client Churn: Client dropout within a cohort can be managed locally. The cohort can report a partial aggregate if some members fail, making the system more robust than direct global aggregation with many unreliable participants.

Cohorts can be constructed to achieve specific statistical properties, directly influencing model quality and fairness.

• Controlled Data Distribution: Cohorts are often formed to have internally IID (Independent and Identically Distributed) data or a controlled mix of non-IID distributions. This can stabilize training compared to fully random selection from a highly heterogeneous population.
• Stratified Cohort Formation: Similar to stratified sampling, clients can be grouped into cohorts based on attributes like data distribution, label skew, or device type. The server then samples cohorts to ensure all strata are represented over time, aiding in bias mitigation.
• Facilitates Personalization: Cohort-level models can serve as an intermediate personalization step. A personalized federated learning approach can first train a cohort-specific model, which is then fine-tuned by individual clients within that cohort.

Cohort selection provides a framework for handling the vast heterogeneity in edge device capabilities, network conditions, and availability.

• Cohort Profiling: Client profiling data (compute, memory, bandwidth) can be aggregated at the cohort level. The server selects cohorts that meet minimum resource thresholds, implementing resource-aware selection efficiently.
• Straggler Mitigation: Slow devices within a cohort can be managed by the cohort's own scheduling or aggregation logic, preventing them from becoming global stragglers. The cohort only communicates when a sufficient fraction of its members have completed work.
• Tiered Learning: Frameworks like TiFL (Tier-based Federated Learning) are a form of cohort selection, where clients are grouped into tiers based on performance. Selecting cohorts from different tiers each round balances speed and model quality.

Cohort selection is not mutually exclusive with other strategies; it often serves as a higher-level organizational layer.

• Foundation for Sampling: Within a selected cohort, further random selection, importance sampling, or power-of-choice can be applied to choose which specific clients train.
• Complement to Client Clustering: Client clustering algorithms (based on data similarity, gradients, or capabilities) are a direct method to form the cohorts. The selection policy then chooses which clusters (cohorts) to activate.
• Enables Advanced Protocols: It is a prerequisite for protocols like federated coreset selection, where a synthetic dataset representing a cohort's data is used for efficient server-side updates. It also aligns with asynchronous FL paradigms, where cohorts submit updates as they are ready.

In cross-silo federated learning, cohorts consist of a small number of reliable, high-capacity organizational clients, such as hospitals or financial institutions. Selection focuses on data complementarity and regulatory alignment.

• Key Drivers: Compliance with data sovereignty laws (e.g., GDPR, HIPAA) and the need to train on vertically partitioned data.
• Example: Three hospitals, each holding different patient feature sets, form a cohort to collaboratively train a diagnostic model without sharing raw records.
• Technical Consideration: Cohorts are often static or semi-static, with selection governed by legal agreements and secure multi-party computation protocols.

Cohort selection directly addresses Non-IID (Non-Independent and Identically Distributed) data across devices by forming groups with complementary distributions.

• Strategy: Actively cluster clients based on data distribution similarity (e.g., using metadata or proxy data) to create more IID-like cohorts for local training.
• Benefit: Reduces client drift—the divergence of local models—leading to faster convergence and a more stable global model.
• Example: In a next-word prediction model, forming separate cohorts for mobile keyboard users in different geographic regions to account for local slang and language patterns.

Cohorting is used to batch clients with similar resource profiles, dramatically improving training round completion times and system throughput.

• Process: Devices are profiled and grouped into tiers based on compute, memory, bandwidth, and battery. Selection occurs within tiers.
• Impact: Minimizes stragglers (slow devices) that delay aggregation. Frameworks like TiFL (Tier-based Federated Learning) use this approach.
• Real-World Metric: In mobile device FL, separating high-end smartphones from legacy models can reduce per-round time by over 60%.

Cohorts serve as logical units for specialized aggregation algorithms and personalized model creation.

• Multi-Model Learning: Different cohorts can train slightly different model variants (personalized FL) tailored to their shared data characteristics.
• Hierarchical Aggregation: Model updates are first aggregated within a cohort, then across cohorts, allowing for cluster-based model tuning.
• Use Case: A smart thermostat manufacturer uses geographic/ climatic cohorts to create regionally personalized energy-saving models before a final global aggregation.

For deployments with millions of potential clients (e.g., smartphones), cohort selection is a scalability necessity. The server selects a manageable number of cohorts per round, not individual devices.

• Mechanism: Uses client clustering as a pre-processing step. The server selects representative cohorts from these clusters.
• Benefit: Reduces orchestration overhead and communication complexity while maintaining statistical representation of the massive population.
• Architecture Link: Often integrated with federated coreset techniques to approximate the global data distribution with selected cohorts.

Cohort structures strengthen privacy guarantees by enabling group-based privacy mechanisms and limiting exposure.

• Differential Privacy (DP): Applying DP noise can be more efficient at the cohort level, providing strong privacy for individuals within the group while preserving utility.
• Secure Aggregation: Cryptographic protocols like Secure Multi-Party Computation (MPC) are more computationally feasible within a small, managed cohort than across thousands of devices.
• Threat Mitigation: Limits the influence and visibility of any single client, providing a layer of defense against model inversion and membership inference attacks.

The overarching process of determining which edge devices or data silos participate in a given training round. It is the parent category for all specific strategies like cohort selection. Key objectives include:

• Improving convergence speed by selecting informative clients.
• Managing system resources (bandwidth, compute, battery).
• Ensuring fairness and representativeness across the client population.
• Mitigating stragglers and Byzantine clients.

A selection method that divides the client population into non-overlapping subgroups (strata) based on key attributes (e.g., geographic region, device type, data distribution). Participants are then sampled from each stratum to form a cohort. This ensures the selected group is statistically representative of the overall population, which is critical for building unbiased and generalizable global models, especially with non-IID data.

A strategy that prioritizes clients based on their available system resources to improve training efficiency and reliability. Selection metrics include:

• Computational power and memory.
• Network bandwidth and latency.
• Battery level (for mobile devices).
• Expected completion time. Frameworks like FedCS (Federated Learning with Client Selection) formalize this approach by selecting clients that can meet a training deadline, directly addressing straggler mitigation.

A pre-processing or ongoing strategy that groups clients based on similarity, often used to inform cohort formation. Clustering dimensions can include:

• Data distribution similarity (via embeddings or statistical tests).
• Hardware capability tiers.
• Network conditions. Clusters enable efficient cohort selection by allowing the server to sample from predefined, homogeneous groups. This can also facilitate personalized federated learning by training specialized models per cluster.

An approach that incorporates explicit fairness constraints to prevent the systematic over- or under-selection of certain client groups. This is a form of bias mitigation. Techniques include:

• Enforcing minimum selection quotas for defined demographic or data strata.
• Using fairness-regularized utility functions in selection policies.
• Applying counterfactual evaluation to assess selection bias. Without it, selection can amplify existing biases in the global model, leading to poor performance for underrepresented groups.

The core algorithm or rule set that governs how clients are chosen. Policies balance exploration and exploitation and can be:

• Rule-based: e.g., "select all clients with >100 data samples."
• Heuristic-based: e.g., Power-of-Choice or Oort's joint optimization of statistical utility and system efficiency.
• Learning-based: e.g., framing selection as a Multi-Armed Bandit problem to learn which client types yield the best model improvement. The policy directly implements the strategy (e.g., cohort, random) and is executed by the federated learning orchestrator.