FedProc: Prototypical Contrastive Federated Learning on Non-IID data

In plain words

In federated learning, many devices each train a model on their own private photos, and a server combines what they learn without ever seeing the photos. The trouble is that each device's photos are often very different from the others': one phone mostly has cats, another mostly has cars. Each device then pulls the shared model in its own direction, and the combined model gets worse.

FedProc keeps the usual "train locally, then average" loop but gives every device a shared reference point for each class. The server keeps one reference point per class, computed as the average of what the devices see for that class. During local training, each device is nudged so that its images of a class move toward that class's reference point and away from the other classes' points. Training also starts with this "learn good features" goal and slowly shifts toward the ordinary "learn to label" goal as the rounds go on.

The point is to make every device's local goal agree with the shared goal, so the averaged model does better on uneven data.

TL;DR

• FedProc adds a supervised prototypical contrastive loss to FedAvg-style local training on non-IID image data §3.3.
• The server averages per-class prototypes from the clients and sends them back each round alg. 1.
• A round-based weight moves training from the contrastive loss to cross-entropy eq. 2 §3.3.
• Author-reported gains over MOON are 1.6 to 7.9 points on three image datasets tab. 2. Several details that matter for implementation are not stated (see Pitfalls).

Problem

Under non-IID client data, local training drifts away from the global optimum, and the authors argue that FedAvg, FedProx, SCAFFOLD and MOON do not use the per-class knowledge that clients hold §1. The authors also show, with t-SNE plots on two CIFAR-10 clients trained alone, that the features of the same class differ between clients fig. 2 §3.2. The problem is stated as minimising a data-size-weighted sum of client losses eq. 1 §3.1.

Method

FedProc is FedAvg with a changed local objective §3.3 alg. 1. The server sends the global model and the global class prototypes to each client. Each client trains for E epochs on its local data with the combined loss. After training, the client recomputes its per-class mean features. The server averages the prototypes uniformly and averages the models weighted by local dataset size.

Local network. The local network has three parts §3.3 fig. 3:

1. A base encoder that produces a representation r from input x.
2. A projection head, an MLP with one hidden layer, that maps r to z in R^Q §3.3. The output size is 256 §4.1.
3. A single linear output layer f_c that maps z to K class logits s §3.3.

The classifier reads z, not r §3.3. The feature extractor f_e is the encoder plus projection head §3.3.

Local objective. The loss is a convex mix of the two parts, with weight α set per round eq. 2 §3.3:

loss = α · loss_gpc + (1 − α) · loss_ce     [eq. 2]
α = 1 − t / T                               [§3.3] [alg. 1 l.21]

Here t is the current round and T the total number of rounds §3.3. The local objective is the expectation of this loss over the client's data eq. 3.

Global prototypical contrastive loss. For a sample with label k, the loss is a softmax over cosine similarities to the class prototypes, with the sample's own class prototype as the positive eq. 4 §3.3:

loss_gpc = −log( exp(sim(z, c_k)) /
                 ( exp(sim(z, c_k)) + Σ_{k' ≠ k} exp(sim(z, c_k')) ) )
                                             [eq. 4] [§3.3]

sim(z, c) = z·c / (‖z‖₂ ‖c‖₂)               [§3.3] [eq. 4]

The prototype of class k on client i is the mean of f_e over that client's samples of class k eq. 5 §3.3. The text describes c_{i,k} as the mean "of the samples belonging to class k in the client C_i" §3.3.

Algorithm 1 alg. 1 (structure as printed; line numbers as printed):

# Server loop of FedProc  [alg. 1 l.1-12]
# Input:  clients, rounds T, local epochs E, learning rate eta
# Output: final global model

def fedproc_server(clients, T, E, eta):
    # initialize weights and class prototypes  [alg. 1 l.3]
    global_model = init_model()
    # c^0 initialisation is not stated in the paper
    global_protos = init_prototypes()
    for t in range(T):
        # every client trains in parallel from the same state  [alg. 1 l.4-8]
        updates = [
            client_update(c, t, T, global_model, global_protos, E, eta)
            for c in clients
        ]
        # prototypes are averaged uniformly over clients  [alg. 1 l.9]
        global_protos = uniform_mean([u.protos for u in updates])
        # models are averaged weighted by local dataset size  [alg. 1 l.10]
        global_model = weighted_mean(
            [u.model for u in updates],
            [u.num_samples for u in updates],
        )
    return global_model

# Client local training  [alg. 1 l.13-24]
# Input:  local data, round t of T, weights, prototypes, epochs E, eta
# Output: updated weights and recomputed class prototypes

def client_update(data, t, T, model, protos, E, eta):
    # split weights into feature extractor and output layer  [§3.3]
    extractor, classifier = split(model)
    for epoch in range(E):
        for images, labels in batches(data):
            # features after the projection head  [§3.3] [alg. 1 l.17]
            z = extractor(images)                 # (B, Q)
            # prototypical contrastive term  [eq. 4] [alg. 1 l.18]
            loss_gpc = prototype_contrastive(z, labels, protos)
            # cross-entropy on the output layer  [§3.3] [alg. 1 l.19-20]
            logits = classifier(z)                # (B, K)
            loss_ce = cross_entropy(logits, labels)
            # weight shifts from features to classifier over rounds
            # [eq. 2] [alg. 1 l.21-22]
            alpha = 1 - t / T
            loss = alpha * loss_gpc + (1 - alpha) * loss_ce
            # one SGD step  [§3.3] [alg. 1 l.23]
            model = sgd_step(model, loss, eta)
    # recompute class means on local data after training  [alg. 1 l.26-30]
    protos = class_means(extractor, data)
    return model, protos

# Global prototypical contrastive loss for one sample  [eq. 4] [§3.3]
# Input:  feature z (Q,), true class label, K prototypes (K, Q)
# Output: scalar loss

def prototype_contrastive(z, label, protos):
    # cosine similarity to every class prototype  [eq. 4]
    sims = [cosine(z, protos[k]) for k in range(K)]
    # positive term: the sample's own class  [eq. 4]
    positive = exp(sims[label])
    # denominator: positive plus the other classes (k' != k)  [eq. 4]
    others = [exp(sims[k]) for k in range(K) if k != label]
    return -log(positive / (positive + sum(others)))

Hyperparameters

Scalability runs §4.4: CIFAR-100 with m = 50 and sampling rate γ = 1, and m = 100 with γ = 0.2 (20 of 100 clients per round). Round counts are in tab. 6.

Results

Author-reported. Setup: CIFAR-10, CIFAR-100 and Tiny-ImageNet with the non-IID Dirichlet split above, default configuration tab. 1 §4.1. Baselines are FedAvg, FedProx, SCAFFOLD, MOON and SOLO. The paper says that, unless stated, comparisons "use reported results from respective papers" §4, and does not mark which baseline rows were re-run.

Top-1 test accuracy, mean ± std tab. 2:

Heterogeneity on CIFAR-100, top-1 tab. 3:

Loss-schedule ablation tab. 4:

Scalability on CIFAR-100, top-1 tab. 6:

Round time, seconds per round on the stated hardware tab. 5: FedAvg 8m34s / 24m36s / 104m11s (CIFAR-10 / CIFAR-100 / Tiny-ImageNet); FedProc 11m57s / 27m41s / 111m12s; MOON 9m06s / 45m05s / 186m10s. Taken together with the accuracy results, these numbers are the basis for the claim of "acceptable computation cost" §4.3.

Accuracy over rounds and local epochs is shown only as curves fig. 4 fig. 5, so no exact values are given for those sweeps.

Paper vs code

No official code. Searched: (1) the paper, abstract and arXiv metadata for links (none); (2) GitHub repository search for "FedProc federated" and "prototypical contrastive federated learning non-IID" (no repository owned by an author or author lab); (3) GitHub search for "FedProc Xidian" (no results). No repository was cloned and none is recorded on this card.

Pitfalls

• Which prototypes does the loss use? §3.3 calls the prototypes "global" and says the server sends them to clients §3.3 alg. 1 l.5-6. Eq. 5 defines c_{i,k} as a mean over one client's own class-k samples, and the loss in Algorithm 1 line 18 uses c^t_{i,k}, with a client index. The paper does not say whether the positives and negatives are the server average from the previous round or the client's own means. Implementers need to choose and document the choice.
• Two averaging rules. Prototypes are averaged uniformly over clients (1/m) alg. 1 l.9, while models are averaged weighted by |D_i|/|D| [alg. 1 l.10]. This is written explicitly, so implement it as printed.
• Empty classes. If a client has no samples of class k, D_i^k is empty and eq. 5 is undefined. The paper does not say how such a client enters the prototype average alg. 1 l.9.
• Denominator range. The text says the sum runs over k' ≠ k §3.3 eq. 4, but Algorithm 1 line 18 writes Σ_{k'} with no exclusion. Read literally, the positive term appears twice in the denominator. The first reading is the one the text states.
• No temperature. Eq. 4 uses raw cosine similarities, which lie in [−1, 1], inside exp, with no temperature. Results depend on this scale; the paper does not discuss it.
• Schedule details. Algorithm 1 indexes t from 0. So α = 1 in round 0, and the classifier term has zero weight that round alg. 1 l.21, l.3-4. α is set once per round, not per step or per epoch, and reaches 1/T in the last round.
• Classifier input. The output layer reads the projection output z, not the encoder output r §3.3. Using r would change the model.
• Prototype initialisation and refresh. c^0 is not specified alg. 1 l.3. Within a round, the loss uses the prototypes sent at the start of the round, and the new prototypes are only computed after local training [alg. 1 l.7, l.26-30]. Whether the prototypes are detached from the gradient, and whether BatchNorm is in train or eval mode when computing them, is not stated.
• Hidden cost. Each client runs an extra forward pass over its full local dataset to compute prototypes after every round alg. 1 l.26-30. Table 5 shows FedProc's CIFAR-10 round at 11m57s versus 8m34s for FedAvg, the slowest of the listed methods on that dataset tab. 5. The text describes overhead on CIFAR-10 and CIFAR-100 as "nearly the same as most of the methods" §4.3.
• Sampling. Algorithm 1 trains all m clients in parallel alg. 1 l.4. The scalability runs use γ = 0.2 on 100 clients §4.4. The paper does not say how prototypes are averaged over the sampled subset.
• Table 6 round counts. For m = 100 the table reports T = 250 and T = 500 tab. 6. The body text mentions only T = 200 for m = 50. The reason for the T = 250 column is not stated.
• Statistics. The ± values in tab. 2 have no stated number of runs or seeds. The partition seed is not stated either. Baseline rows come from the comparison caveat in §4, so their setup may not match exactly.
• Figure-only claims. The best-at-E=10 claim and the convergence claims rest on curves fig. 4 fig. 5, not tables.
• Units. The abstract's "1.6% ~ 7.9%" are absolute points over MOON tab. 2: 70.7 − 69.1 = 1.6, 74.6 − 67.5 = 7.1, 35.4 − 27.5 = 7.9.

Lineage

• Builds on: FedAvg §3.3 (McMahan et al., 2017); prototypes as class means §1; prototypical contrastive learning §1 §2.2; the cumulative-learning weighting idea for α §3.3.
• Compared against: MOON, FedAvg, FedProx, SCAFFOLD, SOLO §4 tab. 2.
• Superseded by: not stated in this wiki.