Model-heterogeneous federated learning

Overview

In model-heterogeneous federated learning (FL), clients differ in what they can run, so the server cannot give every client the same architecture. Most FL studies assume one model on every client, which shuts out low-end devices and caps the size of the shared model 2212.01548 §1. This page covers the approaches that let each client use a model sized to its capacity.

Two families are named in the source. Knowledge distillation can mix architectures, but it needs public data for competitive accuracy, and it sends client models or outputs to the server, which conflicts with secure aggregation 2212.01548 §1 §2. Partial training avoids both problems: each client trains a sub-model cut from the server's model, and the server merges the results 2212.01548 §1 §3.1. Only partial training is compared in detail here. Only one card covers this slug so far. Other methods are named by name, without links, and their content is not used.

Variants

The variants differ in one choice: which nodes of each layer a client receives. Everything else in the loop is shared (see the implementation sketch).

Static extraction (HeteroFL, named only)

The same leading block of nodes is always chosen for a given capacity. Smaller clients receive prefixes of the larger ones 2212.01548 §3.3.2 2212.01548 fig. 2. The paper's argument: parts of the server model that sit outside the leading block never see low-end clients' data [arxiv:2212.01548 §1] §3.3. Named in the source as HeteroFL; not linked here.

Random extraction (Federated Dropout, named only)

Each round picks a random set of nodes for each client 2212.01548 eq. 5. The paper derives an expected-coverage bound for this scheme: the expected number of rounds until every node is selected at least m times grows roughly as I·log I, with I the number of nodes 2212.01548 §3.3.1 app. A.1. Named in the source as Federated Dropout; not linked here.

Rolling extraction (FedRolex)

For a layer with K nodes, the window starts at j mod K in round j and covers ⌊β K⌋ consecutive nodes, wrapping from K − 1 back to 0, where β is the client's capacity 2212.01548 eq. 4. Each round the window advances by one node, so over time every parameter gets trained on data from many clients 2212.01548 §3.2. The server then averages each parameter over the clients whose window covered it, with equal client weights 2212.01548 eq. 10 app. A.2. The paper's illustration uses five nodes: a large client keeps four and a small one keeps three, over three rounds 2212.01548 fig. 1. Card: FedRolex.

When to use which

Reported results, author-reported and not re-run by us. Global accuracy of the server model 2212.01548 tab. 3:

The ± values in the source are not defined in the visible text, so they are omitted here 2212.01548 §4. The paper reports FedRolex above both partial-training baselines in every column 2212.01548 §4.1. Client participation at 5%, 10% and 20% on CIFAR-10 gives FedRolex 57.90, 69.44 and 71.85, against 48.43, 63.90 and 65.07 for HeteroFL 2212.01548 tab. 6.

Implementation sketch

Common core of partial training: the server holds the full parameters, a client's capacity sets how many nodes it receives, and the server merges results per parameter. Only the select_indices step changes between static, random and rolling variants 2212.01548 §3.3.

# Rolling window over one layer  [arxiv:2212.01548 eq. 4]
# Input:  round j; K nodes in the layer; capacity beta in (0, 1]
# Output: indices of the nodes this client receives

def rolling_indices(j, K, beta):
    # Window start moves one node per round and wraps at K  [eq. 4]
    start = j % K
    # Number of nodes kept, floor of capacity times width  [eq. 4]
    width = floor(beta * K)
    return [(start + t) % K for t in range(width)]

# One server round  [arxiv:2212.01548 alg. 1]
# Input:  global params theta; round j; sampled clients; capacities beta
# Output: updated global params theta

def partial_training_round(theta, j, clients, beta):
    updates = []
    for m in clients:
        # Variant point: rolling here; random or static is swapped in  [§3.3]
        idx = {i: rolling_indices(j, K[i], beta[m]) for i in layers}
        # Send the client only the selected nodes of each layer  [alg. 1]
        local = extract(theta, idx)
        # Client trains its sub-model on its own local data  [alg. 1]
        local = train(local, data[m])
        updates.append((idx, local))
    # Average each parameter over the clients that trained it  [eq. 10]
    return selective_average(theta, updates)

A parameter that no client trained in a round keeps its old value 2212.01548 §3.2. The paper's client update is a plain gradient step on the extracted sub-model 2212.01548 alg. 1. Its architectures were a pre-activated ResNet-18 with static batch normalization and a three-layer Transformer 2212.01548 §4. Static batch normalization uses batch statistics at test time, so evaluation needs care. The paper does not say how it evaluated batch norm 2212.01548 §4.

Contradictions

The paper's text and its released code disagree on several points. Both sides are cited, and we do not resolve them:

• Roll direction. The text moves the window forward 2212.01548 §3.2 eq. 4. The code's window moves backward, which covers the same nodes over time but in a different order. See FedRolex (Paper vs code).
• Stack Overflow cohort. The text says 200 clients per round 2212.01548 §4 tab. 8. The code uses 50 (see the card).
• CIFAR optimizer settings. The table gives weight decay 5e-4 and decay factor 0.1 2212.01548 tab. 8. The code sets 1e-3 and 0.25 for CIFAR (see the card).
• Overlap ablation. The paper's ablation calls r = 1 FedRolex [arxiv:2212.01548 app. A.4]. The code's overlap path draws a random permutation each round instead of following the window (see the card). The paper does not say which path produced its numbers.

Papers

• FedRolex: Model-Heterogeneous Federated Learning with Rolling Sub-Model Extraction: rolling window extraction for clients of different capacity. Each round's window moves one node, and the server averages each parameter over the clients that trained it.