In Federated Learning, a global model is learned by aggregating model updates
computed at a set of independent client nodes, to reduce communication costs
multiple gradient steps are performed at each node prior to aggregation. A key
challenge in this setting is data heterogeneity across clients resulting in
differing local objectives which can lead clients to overly minimize their own
local objective, diverging from the global solution. We demonstrate that
individual client models experience a catastrophic forgetting with respect to
data from other clients and propose an efficient approach that modifies the
cross-entropy objective on a per-client basis by re-weighting the softmax
logits prior to computing the loss. This approach shields classes outside a
client's label set from abrupt representation change and we empirically
demonstrate it can alleviate client forgetting and provide consistent
improvements to standard federated learning algorithms. Our method is
particularly beneficial under the most challenging federated learning settings
where data heterogeneity is high and client participation in each round is low