Preservation of the Global Knowledge by Not-True Distillation in Federated Learning

In plain words

In federated learning, many devices each train a copy of a model on their own data, and a server averages the copies. When each device's data is very different from the others, the shared model tends to do worse. The authors suggest one reason: after each round, the shared model forgets what it had learned about the kinds of data a device does not hold. Local training makes this worse, because each device only pushes its own kinds of examples.

Their fix is to keep a memory of the old shared model while a device trains. On its own data, the device learns the correct label as usual. It also copies the old shared model's opinions about the wrong labels, and ignores the correct label when doing so. This keeps the local model from drifting away from what the whole group knows. Nothing extra is sent over the network. The authors report that this helps across several image datasets.

TL;DR

• FedNTD trains each client with cross-entropy plus a distillation loss that matches the frozen global model only on the not-true classes §3.
• The paper argues that forgetting of out-of-local knowledge is a main cause of data heterogeneity problems in federated learning §1, §2.
• Author-reported: best or near-best accuracy on most of 14 settings (4 datasets, two partition schemes) against FedAvg and five other methods tab. 1.
• Costs an extra forward pass of the global model on every client batch, and a copy of the model per client round. Adds no communication §4.2, code.

Problem

Non-IID data across clients hurts federated averaging. The authors argue this is mainly forgetting: the global model loses knowledge from earlier rounds, and local training erases knowledge about classes a client does not hold §1, §2.1, §2.2. They measure this with class-wise accuracy across rounds and with an analysis of which classes the local model forgets §2, fig. 2, fig. 4. They treat the forgetting measure as a diagnostic, and do not prove the causal link §2.

Method

FedNTD changes only the local loss. The server keeps the usual averaging round. Each round, the server samples clients, sends them the current global weights, and each client trains its local copy alg. 1 §3.

Local loss §3, eq. labelled FedNTD. For a labelled example with true class y, the local model produces logits z^l. The frozen global model from the start of the round produces logits z^g. The loss is cross-entropy plus a weighted distillation term:

\mathcal{L} = \mathcal{L}_{\text{CE}}(q^{l}, \mathbb{1}_y) + \beta \cdot \mathcal{L}_{\text{NTD}}(\tilde{q}_{\tau}^{l}, \tilde{q}_{\tau}^{g}) §3

Not-true distillation loss §3. Softmax is taken only over the not-true logits, with temperature \tau. For every class c \neq y:

\mathcal{L}_{\text{NTD}} = -\sum_{c \neq y} \tilde{q}_{\tau}^{g}(c) \log\frac{\tilde{q}_{\tau}^{l}(c)}{\tilde{q}_{\tau}^{g}(c)} §3

\tilde{q}_{\tau}^{l}(c) = \frac{\exp(z_c^{l}/\tau)} {\sum_{\tilde{c} \neq y} \exp(z_{\tilde{c}}^{l}/\tau)}, \qquad \tilde{q}_{\tau}^{g}(c) = \frac{\exp(z_c^{g}/\tau)} {\sum_{\tilde{c} \neq y} \exp(z_{\tilde{c}}^{g}/\tau)} §3

The true-class logit is excluded from both softmaxes. Its gradient from the distillation term is therefore zero §3. The paper says \beta sets how strongly out-of-local knowledge is kept, and that the trade-off resembles the stability-plasticity dilemma in continual learning §3.

Why the not-true form §3, prop. 2. Proposition 2 splits the KL term into a part on the true class and a part on the not-true classes. It shows that the not-true part is a weighted sum over the local label distribution and the out-of-local distribution. The derivation is in the appendix, "Proof of Proposition 2".

Round procedure alg. 1. Inputs: T rounds, E local epochs, sampled clients K^{(t)}, learning rate \gamma.

# One communication round t  [alg. 1]
# global_weights: w^(t-1); data_of[k]: client k's local dataset
def fedntd_round(global_weights, sampled, data_of, lr):
    # Each sampled client starts from the global weights  [alg. 1]
    uploads = []
    for client in sampled:
        local_weights = copy(global_weights)
        # Frozen copy of the round-start global model  [§3]
        reference = copy(global_weights)
        # E local epochs of mini-batch SGD on the FedNTD loss  [alg. 1]
        for epoch in range(local_epochs):
            for batch in data_of[client]:
                logits = forward(local_weights, batch.inputs)
                ref_logits = forward(reference, batch.inputs)
                loss = fedntd_loss(logits, ref_logits, batch.targets)
                local_weights = sgd_step(local_weights, loss, lr)
        uploads.append(local_weights)
    # Average the uploaded weights  [alg. 1]
    # The code weights each client by its data size; see Paper vs code
    return mean(uploads)

# FedNTD loss on one mini-batch  [§3, eq. FedNTD]
# logits, ref_logits: (B, C); targets: (B,)
def fedntd_loss(logits, ref_logits, targets, beta=1.0, tau=1.0):
    # Ordinary cross-entropy over all C classes  [§3]
    ce_loss = cross_entropy(logits, targets)
    # Remove the true-class column from both logit matrices  [§3]
    local_nt = drop_true_class(logits, targets)      # (B, C-1)
    ref_nt = drop_true_class(ref_logits, targets)    # (B, C-1)
    # Temperature-softened softmax over the not-true classes  [§3]
    local_probs = softmax(local_nt / tau)
    # The reference side is a constant, so no gradient flows into it  [§3]
    ref_probs = softmax(ref_nt / tau)
    # KL(reference || local), summed over classes, mean over batch  [§3]
    ntd_loss = kl_divergence(ref_probs, local_probs)
    return ce_loss + beta * ntd_loss

Hyperparameters

Values are as stated in the paper unless marked as code. The paper's main text and appendix disagree on the initial learning rate; both are listed.

Results

Author-reported. Setup: four datasets, 100 clients (200 for CINIC-10), two non-IID partition schemes, the FedAvg CNN §4.1 app. Experimental Setups. Each cell is test accuracy (%) with the forgetting measure \mathcal{F} in parentheses. Source: tab. 1. The paper reports standard deviations in a separate appendix table app. Experiment Table with Standard Deviation.

Sharding partition

LDA partition (CIFAR-10 columns give \alpha; CINIC-10 uses \alpha = 0.1)

Headline claims from the text §4.2: FedNTD is best in most cases, and the paper says it "consistently outperforms the baselines on all setups". The table above does not fully support the second claim (see Pitfalls). In a separate comparison on the released FedAlign code, the authors report that FedAlign's gains disappear at severe heterogeneity while FedNTD improves consistently app. Comparison to FedAlign. The paper states that FedNTD adds no communication burden, while several baselines need statefulness, extra communication or auxiliary data §4.2, app. Conceptual comparison to prior works.

Paper vs code

Official repository: https://github.com/Lee-Gihun/FedNTD, owned by the first author and linked from the paper's abstract. Commit be00ee598139654abe650c446a7ba3bc4e340233, MIT license LICENSE.

• Averaging weights. Algorithm 1 averages the uploads with weight 1/|K^{(t)}| alg. 1. The code weights each client by its local dataset size BaseServer.py. This matters under LDA, where client sizes differ. The paper does not say which averaging produced the reported numbers.
• \tau^2 factor. The code multiplies the KL term by \tau^2 criterion.py. The paper's loss has no such factor §3. The two agree at the default \tau = 1. For other \tau the code effectively rescales \beta by \tau^2, so the paper's \tau ablation can only be read under one convention, and the paper does not say which it used.
• Initial learning rate. Main text says 0.1 §4.1; appendix and code say 0.01 app. Learning Setups fedntd.json. The code agrees with the appendix. Treat 0.01 as the value used, but note the conflict.
• Reference model. The code keeps a frozen copy of the model at the start of each client's round BaseClientTrainer.py and uses its logits in the loss ClientTrainer.py. This matches §3.
• True-class removal. The code drops the true-class column with a gather over the other C-1 indices utils.py. The softmax is taken over those C-1 logits criterion.py. This matches the paper.
• Momentum. The paper says momentum is local only app. Learning Setups. The code loads the server optimizer state into each client [code: algorithms/BaseClientTrainer.py#L105-L108]. The server optimizer is never stepped, so no momentum buffer is carried between rounds. This is consistent with the paper, but inferred from the code rather than stated.
• LR schedule. The scheduler steps once per round after aggregation BaseServer.py, and clients inherit the current learning rate through the optimizer state BaseServer.py. This matches the 0.99 per-round decay.
• LDA resampling. The code retries the Dirichlet split until every client has at least 10 samples datasetter.py. The paper does not mention this minimum.
• Client sampling. The code seeds NumPy with the round index before sampling, so the same clients are chosen across methods BaseServer.py. The paper does not state this.
• Config vs paper. The shipped config covers CIFAR-10 with 200 rounds, 5 local epochs, sampling ratio 0.1, and \beta = \tau = 1 fedntd.json. These match the paper for CIFAR-10. Only the FedNTD config was checked; the other configs in the repository were not compared.

No other discrepancies found in the local-loss and aggregation logic beyond the items above.

Pitfalls

• Claim vs table. §4.2 says FedNTD "consistently outperforms the baselines on all setups". In tab. 1, SCAFFOLD is higher on CIFAR-10 s=5 (65.74 vs 64.93), on CIFAR-10 s=10 (68.97 vs 68.56), on CIFAR-10 LDA α=0.5 (66.16 vs 66.12) and on CIFAR-100 LDA (33.68 vs 32.37). The bold marks are also inconsistent: both FedNTD and SCAFFOLD are bold for CIFAR-10 s=10.
• Baseline tuning is uneven. The paper grid-searches the baselines but not FedNTD, which uses fixed \beta = \tau = 1 [app. Algorithm Implementation Details]. The table lists one hyperparameter value per baseline and the grid candidates, but does not say which value produced each reported number.
• Failures. FedNova and SCAFFOLD report 10.00% ("failed") on CIFAR-10 LDA \alpha = 0.05 tab. 1. The paper does not explain these failures.
• \tau is sensitive. Too small \tau "significantly drops the accuracy"; \beta is less sensitive app. Effect of FedNTD Hyperparameters. Use \tau = 1 as the default unless you tune it.
• MSE is not a drop-in replacement. The MSE version of the loss gets worse as \beta rises and diverges at \beta = 0.3, while the KL form reaches 52.61% on CIFAR-10 s=2 at \beta = 1.0 app. MSE Loss for Not-True Distillation. Use the softmax-KL form.
• Cost per step. Each batch needs a forward pass through the frozen global model ClientTrainer.py, so local compute is roughly doubled in the forward pass. The paper does not report this cost, and the "no extra communication" claim does not cover compute.
• Memory. The client makes a full deep copy of the model every round BaseClientTrainer.py.
• Proof reference. The main text points the reader to "proof_proposition2" for the proposition derivation §3. The appendix proof labelled "Proof of Proposition 2" is the correct one. Do not confuse it with "Proof of Proposition 1", which is a different result.
• Rounds are not stated for the main experiments §4. Use the config value of 200 if you need to reproduce, and say so.
• Table typo. In the LDA table, FedNova's CIFAR-100 entry has a malformed forgetting parenthesis in the source; the value 0.28 is used here as written.
• Figures are not reproduced. fig. 1 shows the concept, fig. 2 the global-model forgetting, fig. 3 the in-local and out-local distributions, fig. 4 local forgetting. Describe them in words only.

Lineage

• Builds on: FedAvg (federated averaging), and the continual-learning stability-plasticity framing §3. The paper frames forgetting using continual learning §1, §6. Not linked: no paper card exists for these yet.
• Compared against: FedCurv, FedProx, FedNova, SCAFFOLD, MOON tab. 1, and FedAlign as an add-on app. Comparison to FedAlign. Also knowledge distillation over the full softmax app. Comparison to KD.
• Superseded by: not known from content in this wiki.