FedRolex: Model-Heterogeneous Federated Learning with Rolling Sub-Model Extraction

In plain words

Federated learning trains one model across many devices without collecting their data. Those devices differ a lot: some are weak and cannot run the same network the server runs. Most methods make every device use the same model, which shuts out the weakest devices and keeps the shared model small.

FedRolex lets each device train a smaller slice of the server's model. The slice moves. Each round the window of parts it takes shifts along, so over time every part of the big server model gets trained on data from many devices. The server then averages each part using only the devices that trained it. The result is a server model larger than any single device could run, and, according to the authors' experiments, better accuracy than the other slicing methods.

TL;DR

• FedRolex is a partial-training method for federated learning with heterogeneous clients §3.2.
• Each client trains a sub-model cut from the server model. The cut is a window over each layer's nodes that advances one step per round, wrapping around §3.2 eq. 4.
• The server averages each parameter over the clients that updated it §3.2 eq. 10.
• Author-reported: beats HeteroFL and Federated Dropout on CIFAR-10/100 and Stack Overflow tab. 3; trains server models larger than the largest client §4.4 fig. 4.
• The paper has no convergence proof, and it does not cover deploying the trained server model to clients §5.

Problem

Most cross-device FL studies assume the server and every client run the same model §1. That excludes low-end clients and blocks large models on resource-limited devices §1.

Knowledge-distillation methods can handle different architectures, but they need public data for competitive accuracy and they send client models or outputs to the server, which conflicts with secure aggregation §1 §2.

Partial-training methods avoid both problems. Random extraction (Federated Dropout) and static extraction (HeteroFL, FjORD) pick the sub-model in a fixed or random way. Under those schemes, parts of the server model are trained unevenly, and some parts never see low-end clients' data §1 §2 §3.3.

Method

Setting §3.1. Each client n holds local data and trains a sub-model whose capacity is β_n. β_n is the fraction of nodes it keeps in each layer. The server model θ is split into client sub-models θ_n, and each client's loss uses its own θ_n eq. 3.

Rolling extraction §3.2. For layer i with K_i nodes, the server selects node indices for client n in round j by the window in eq. 4. The window starts at ĵ = j mod K_i and covers ⌊β_n K_i⌋ consecutive nodes, wrapping from K_i − 1 back to 0 eq. 4. The paper's illustration uses 5 nodes per layer: a large client keeps 4 and a small client keeps 3, over three rounds §3.2 fig. 1.

Server round §3.2 alg. 1.

1. Sample a subset M of clients from N alg. 1.
2. For each client m in M, extract its sub-model with the window of eq. 4 and broadcast it alg. 1.
3. Each client runs local training on its data and returns its updated sub-model alg. 1.
4. Aggregate per parameter with the selective average of eq. 10 §3.2 app. A.2.

Aggregation §3.2 app. A.2. A parameter's new value is the average of the client values for clients whose sub-model contains that parameter. Clients are not weighted: p_m is set to 1 (or 1/N), and the weights cancel app. A.2. A parameter that no client updated in the round stays unchanged §3.2.

Comparison with other extraction schemes §3.3. The paper compares the rolling window against random extraction §3.3.1 and static extraction §3.3.2. In the random scheme, each round selects a random set of nodes eq. 5. In the static scheme, the same leading block of nodes is always selected for a capacity eq. 6, and smaller clients' models are prefixes of larger ones §3.3.2 fig. 2.

Statistical argument §3.3.1 app. A.1. For random extraction, the expected number of rounds until every node is selected at least m times grows as I·log I + I(m−1)·log log I app. A.1. The paper then says the rolling scheme needs about m·I rounds §3.3.1. Its derivation for the rolling scheme is not in the visible source; only the random-scheme bound is derived.

Pseudocode. The first block restates the window of eq. 4 as a function. The second block is one server round alg. 1. The third block is the aggregation rule eq. 10. The client update is a plain gradient step on the extracted sub-model alg. 1.

# Rolling window of node indices for one layer  [eq. 4]
# Input:  round j; layer width K; capacity beta
# Output: list of selected node indices

def rolling_indices(j, K, beta):
    # The window starts at j mod K and advances one node per round  [eq. 4]
    start = j % K
    # Number of nodes this client keeps in the layer  [§3.2] [eq. 4]
    width = floor(beta * K)
    # Consecutive nodes from start, wrapping past K back to 0  [eq. 4]
    return [(start + t) % K for t in range(width)]

# One server round  [alg. 1]
# Input:  global params theta; round j; participants M; capacities beta
# Output: updated global params theta

def server_round(theta, j, participants, beta):
    # Each participant gets a sub-model cut with its own window  [§3.2]
    updates = []
    for m in participants:
        # Node indices for every layer of this client  [eq. 4]
        idx = {i: rolling_indices(j, K[i], beta[m]) for i in layers}
        # Copy the selected slice of the global model to the client  [alg. 1]
        local = extract(theta, idx)
        # Local training on the client's data  [alg. 1]
        local = client_train(local, data[m], lr)
        updates.append((idx, local))
    # Average each parameter over the clients that trained it  [eq. 10]
    return aggregate(theta, updates)

# Selective averaging with equal client weights  [eq. 10] [§3.2]
# Input:  global params theta; list of (idx, local) pairs
# Output: updated global params theta

def aggregate(theta, updates):
    # Visit every scalar parameter of the global model  [app. A.2]
    for each entry e in theta:
        # Local values from clients whose window covers entry e  [eq. 10]
        senders = [local[e] for idx, local in updates if covers(idx, e)]
        # Keep the old value when no client trained this entry  [§3.2]
        if senders:
            theta[e] = mean(senders)
    return theta

# Local training on the extracted sub-model  [alg. 1] [tab. 8]
# Algorithm 1 writes a per-sample step; Table 8 gives batch sizes.
# Input:  local params; data; learning rate lr
# Output: trained local params

def client_train(local, data, lr):
    # Table 8 uses 1 local epoch  [tab. 8]
    for batch in data:
        # Gradient step on the batch loss  [alg. 1]
        local = local - lr * grad_loss(local, batch)
    return local

Equation 4, the window (full):

\mathcal{S}^{(j)}_{n,i} = \begin{cases} \{\hat j, \dots, \hat j + \lfloor \beta_n K_i \rfloor - 1\} & \text{if } \hat j + \lfloor \beta_n K_i \rfloor \le K_i, \\ \{\hat j, \dots, K_i - 1\} \cup \{0, \dots, \hat j + \lfloor \beta_n K_i \rfloor - 1 - K_i\} & \text{otherwise,} \end{cases} \qquad \hat j = j \bmod K_i eq. 4

Equation 10, the aggregation (full):

\theta_{[i,k]} = \frac{1}{\sum_{m \in \mathcal{M}_k} p_m} \sum_{m \in \mathcal{M}_k} p_m \, \theta_{m,[i,k]} eq. 10

Here M_k is the set of clients that updated parameter k of layer i app. A.2.

Architectures §4.

• Small regime: pre-activated ResNet-18 on CIFAR-10 and CIFAR-100. Batch normalization is replaced by static batch normalization, and a scalar module follows each convolution §4.
• Large regime: a 3-layer Transformer on Stack Overflow with a vocabulary of 10,000 words, token embedding size 128, feed-forward hidden size 2048, ReLU, 8 attention heads §4.
• Width levels for clients: β ∈ {1, 1/2, 1/4, 1/8, 1/16}. For ResNet-18 the number of convolution kernels is scaled. For the Transformer, the hidden width of the attention heads is scaled §4.
• Server model: same width as the largest client in the main comparison §4.1. In §4.4 the server is γ ∈ {2, 4, 8, 16} times the client width §4.4.

Data heterogeneity §4. On CIFAR-10 each client holds L = 2 labels (high heterogeneity) or L = 5 (low). On CIFAR-100, L = 20 or L = 50. These levels approximately match Dirichlet α = 0.1 and 0.5 §4. Stack Overflow is partitioned by user ID §4.

Participation §4. CIFAR: 100 clients, 10% (10 clients) per round. Stack Overflow: 342,477 clients, 200 per round per the text §4 (see Paper vs code).

Hyperparameters

Training settings for the FedRolex and partial-training rows of tab. 3 tab. 8:

Homogeneous baselines use different settings tab. 9:

Settings for the ρ sweep fig. 3 tab. 10 tab. 11: ρ is the share of large-capacity clients. For CIFAR-10 (high and low heterogeneity alike), decay points are 500 and 1000 with 1250 rounds for ρ ∈ {0.0, 0.1, 0.2, 0.8, 0.9, 1.0}, and 700 and 1200 with 1500 rounds for ρ ∈ {0.3, 0.4, 0.6, 0.7}. CIFAR-100 uses decay points 1000 and 1500 with 2000 rounds for all ρ. Stack Overflow uses decay point 800 and 1500 rounds for all ρ. Other values are as in tab. 8.

Settings for the server-size sweep fig. 4 tab. 12: decay points 800 and 1200 with 1500 rounds for CIFAR-10 and CIFAR-100 (both heterogeneity levels); Stack Overflow uses decay point 800 and 1500 rounds. The table gives only these values; batch size and client count are not stated there.

Settings for the emulated real-world distribution fig. 5 tab. 13: identical to tab. 8.

Client-capacity distribution for the real-world experiment fig. 7: the paper maps household income to capacity. The pie shares are 6, 10, 11, 18 and 55 (sum 100) for 1×, 1/2×, 1/4×, 1/8× and 1/16×. Income bands: 1/16× below $75,000; 1/8× $75,000–$100,000; 1/4× $100,000–$150,000; 1/2× $150,000–$200,000; 1× above $200,000 fig. 7.

Clients per round: 10% of 100 clients is ceil(0.1 · 100) = 10 config.yml main_resnet.py.

Results

Author-reported. All numbers are mean ± a spread the text does not define, from 5 seeds (CIFAR) or 3 seeds (Stack Overflow) §4. Global accuracy is the server model on the test set. Stack Overflow reports next-word accuracy, including OOV and end-of-sentence tokens §4.

Table 3 (global accuracy, tab. 3).

KD-based rows are copied from another paper tab. 3 note, not re-run here. The paper says KD results are not from identical settings §4.1.

Headline numbers §4.1 tab. 3: FedRolex beats both PT baselines in all columns. It is below FedDF and Fed-ET on CIFAR-10 high heterogeneity, and above all KD baselines on CIFAR-100. On Stack Overflow it is 29.22 vs. 29.79 for the largest homogeneous model tab. 3.

Table 4 (emulated real-world distribution, tab. 4).

The FedRolex low-heterogeneity CIFAR-10 global cell is missing an opening parenthesis in the source; the value is read as 80.14 (±0.52). Local accuracy is the server model on each client's local data §4.

Claims §4.5 fig. 5: FedRolex beats the smallest homogeneous model on global accuracy for all three datasets. On local accuracy it beats the smallest model on CIFAR-10 and CIFAR-100. The per-client distribution shift is shown in fig. 5 and described in words only.

Table 5 (client weighting ablation, high heterogeneity, tab. 5).

The paper finds no weighting scheme significantly better and keeps non-weighting app. A.3.

Table 6 (client participation, CIFAR-10, global accuracy, tab. 6).

Table 7 (per-client cost, tab. 7). Averages over participating clients per round app. A.6.

Figures, described in words.

• fig. 3 (§4.3): global accuracy vs. ρ, the share of large-capacity clients (β = 1) against small ones (β = 1/16). On CIFAR-10 the gap between high and low heterogeneity stays large for ρ from 0.1 to 1. On CIFAR-100 the heterogeneity gap is much smaller, and the largest capacity is the bottleneck. On all three datasets a small share of large clients gives most of the gain, and more of them add little §4.3.
• fig. 4 (§4.4): global accuracy as γ (server/client width ratio) grows from 2 to 16. Accuracy drops for both FedRolex and Federated Dropout, mostly from γ = 1 to 4. FedRolex is higher than Federated Dropout at every γ on CIFAR-10 and CIFAR-100. On Stack Overflow the drop is much smaller.
• fig. 6 (app. A.4): effect of the overlap ratio r on accuracy; the effect is non-linear and inconsistent.

Paper vs code

Official repository: AIoT-MLSys-Lab/FedRolex, linked from the paper's abstract (arXiv metadata) README.md. Commit 813510997e1802eb756d53baa4229c90ce5ac008. License: Apache-2.0 (LICENSE file). The repo is built on HeteroFL README.md. I read the ResNet and Transformer servers, the client, the data settings, and the configs.

Differences:

1. Roll direction. Paper: the window starts at ĵ and moves forward, for example {a,b,c,d} → {b,c,d,e} §3.2 eq. 4 fig. 1. Code: torch.roll with a positive shift, so the start index is (−round) mod K and the window moves backwards resnet_server.py transformer_server.py. The set of nodes covered over time is the same, but the order differs from the paper's illustration.

2. --overlap path is not the paper's scheme. The default (overlap unset, the README command) follows eq. 4 with the roll rounds % K resnet_server.py README.md. When --overlap r is set, the code (a) uses a random permutation of the K nodes per layer resnet_server.py; (b) advances the offset by 1 + ⌊L(1−r)⌋ per round and takes it modulo the local width L, not K resnet_server.py; and (c) calls reshuffle_params() whenever rounds % 512 is non-zero, which is almost every round, and this also runs on the default path, where the permutation is unused resnet_server.py resnet_server.py. The reshuffle resets the offset to 0 resnet_server.py. So the offset never accumulates, and each round selects a fresh random set of nodes, closer to Federated Dropout than to eq. 4. This reading is from the source and was not executed. The paper's ablation app. A.4 describes r = 1 as FedRolex, but the code path for r is different. The paper does not say which path produced its reported numbers.

3. CIFAR weight decay and decay factor. Paper: weight decay 5e-4 and decay factor 0.1 tab. 8. Code: process_control sets weight decay 1e-3 and decay factor 0.25 for CIFAR-10/100 utils.py utils.py. It runs after argument parsing, so the command line cannot override these values main_resnet.py main_resnet.py. Momentum 0.9 and SGD agree utils.py.

4. Stack Overflow cohort. Paper: 200 clients per round §4 tab. 8. Code: active_user = 50 for Stack Overflow utils.py; the server uses it as the cohort size main_transformer.py.

5. Learning rate, schedule and epochs. Paper: 2e-4 (CIFAR-10), 1e-4 (CIFAR-100), 2e-4 (Stack Overflow); decay points and round counts in tab. 8. Code defaults: lr 3e-4 and Adam in config.yml config.yml, 200 epochs and milestones 100 and 150 config.yml. The CIFAR and Stack Overflow branches override the optimizer to SGD utils.py utils.py. The README example uses --lr 2e-4, --g_epoch 3200, and --schedule 1200, which match none of tab. 8 README.md. Use the flags to reproduce tab. 8; the defaults do not.

6. Batch size. Paper: 10 (CIFAR-10), 24 (CIFAR-100 and Stack Overflow) tab. 8. Code agrees for CIFAR-10 (10) and Stack Overflow (24) utils.py utils.py. CIFAR-100 uses 10 in the code too utils.py; the paper says 24 tab. 8.

7. Masked classifier, not in the paper. The code masks the classifier output to the labels of each client during training (mask is 1 in the control string), by setting other logits to 0 rather than −∞ resnet.py utils.py. For the linear layer, aggregation uses only the rows for the client's labels resnet_server.py resnet_server.py. The paper does not describe either step.

8. Aggregation default. Code default weighting = 'avg' main_resnet.py sets K = 1 resnet_server.py, so the average is an unweighted mean over the clients that updated each entry resnet_server.py. This matches eq. 10 with p_m = 1. Batch-norm and bias entries are multiplied by a running update count, which cancels within the average because all clients share the count per entry. No difference in effect.

9. Transformer heads. Paper: 8 heads with 12-dimensional query, key and value vectors §4. Code: embedding 128 and 8 heads utils.py. The server splits the output width by 8 heads transformer_server.py, so each head gets 16 for width 128. The paper's 12 × 8 = 96 does not match 128 in the paper itself.

10. Client step. Algorithm 1 writes a plain per-sample step alg. 1. Table 8 gives batch sizes, and the code iterates over batches with optimizer.step() resnet_client.py.

11. Static batch norm. Paper: static batch normalization §4. Code: BatchNorm2d with track_running_stats=False by default resnet.py resnet.py. This matches the paper's "static" BN; see Pitfalls for the test-time effect.

12. Not compared: FjORD, because its code was not open source §4 fn.. The paper says so.

Search steps: the paper's own link to the repository (step 1) was enough, so no GitHub search was run. No other repository was considered.

Pitfalls

• Use the flags for Table 8. The defaults in config.yml (Adam, lr 3e-4, 200 epochs) are not the paper's settings. The CIFAR branch then overrides weight decay and decay factor (see Paper vs code 3), so the paper's CIFAR values cannot be set from the command line.
• Overlap ablation is not FedRolex. If you implement the --overlap idea, implement eq. 4 directly: take the window of nodes in a fixed order and advance it by a step. Do not use the code path described in Paper vs code 2.
• Roll direction. If you match the paper's illustration (forward), use (j + t) % K. If you match the code's order, use (−j + t) % K. The coverage over rounds is the same.
• Static batch norm changes inference. With track_running_stats=False, batch norm uses batch statistics even at test time resnet.py. Evaluate with batches that are representative, or plan a separate statistics step. The paper does not say how evaluation handled BN.
• Masked classifier. The code masks the output layer to each client's labels during training and aggregates the linear layer per label (Paper vs code 7). The paper does not describe either step. To match the reported numbers, follow the code and mark this as a code-only detail.
• Participation count. Stack Overflow cohort: 200 in the paper, 50 in the code. Set it explicitly.
• Rolling statistics. The "rolling needs m·I rounds" claim in §3.3.1 has no visible derivation; only the random-scheme bound is derived in app. A.1. Do not rely on the rolling number without re-deriving it.
• Stated limits. The paper gives no full convergence analysis §5. It does not cover deploying a trained server model to devices of the lower capacities §5.
• ± is undefined. The visible text does not define the ± in the tables. A source comment says "mean and standard deviation", but that line is commented out §4.
• KD numbers are borrowed. FedDF, DS-FL and Fed-ET results come from Cho et al. tab. 3 note. Their settings are described as similar, not identical §4.1. Do not treat them as re-run baselines.
• Stack Overflow claim. The text says FedRolex is "on par" with the largest homogeneous model. The numbers are 29.22 and 29.79, with the uncertainty given in tab. 3.
• Communication and compute costs are client-side averages app. A.6. Server-side cost is not reported.

Lineage

• Builds on: partial training as in HeteroFL (static extraction) and Federated Dropout (random extraction), both cited in §2 §2. Named here without arXiv ids; this card uses no content from those papers.
• Compared against: FedDF, DS-FL and Fed-ET (knowledge distillation); HeteroFL and Federated Dropout (partial training); homogeneous FL at the smallest and largest client size §4 tab. 3.
• Related concurrent work named in the introduction: Federated Select §1.