[Paper Review] Re-Weighted Softmax Cross-Entropy to Control Forgetting in Federated Learning

Gwen Legate, Lucas Caccia+1 more|arXiv (Cornell University)|Apr 11, 2023

Privacy-Preserving Technologies in Data4 citations

TL;DR

This paper proposes Weighted Softmax (WSM), a client-level re-weighting of softmax logits in federated learning to reduce catastrophic forgetting caused by data heterogeneity. By suppressing gradients for out-of-distribution classes during local training, WSM mitigates client drift and improves global model performance, especially under high data heterogeneity and low client participation, with gains up to 7.6% on CIFAR-100 and 1.02% on CIFAR-10.

ABSTRACT

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.

Motivation & Objective

• Address the problem of local client forgetting in federated learning, where clients overfit to their local data and forget representations for other clients' data.
• Identify that data heterogeneity and limited client participation exacerbate client drift and degrade global model performance.
• Propose a lightweight, communication-efficient method to reduce forgetting without requiring data sharing or additional model parameters.
• Demonstrate that modifying the cross-entropy objective via per-client softmax re-weighting improves generalization across clients and enhances convergence.
• Evaluate the method across multiple FL algorithms (FedAvg, SCAFFOLD, FedProx) and datasets (CIFAR-10, CIFAR-100) under varying heterogeneity and client participation levels.

Proposed method

• Introduce a client-specific re-weighting of softmax logits before computing cross-entropy loss, using a class-specific weighting scheme.
• Apply a modified cross-entropy objective where logits for classes outside a client’s label set are down-weighted to reduce their influence on local updates.
• Use a temperature-scaled softmax with class-specific weights to stabilize training and reduce gradient updates for out-of-distribution classes.
• Integrate the weighted softmax (WSM) objective into standard federated optimization algorithms like FedAvg, SCAFFOLD, and FedProx without altering aggregation or communication protocols.
• Design the weighting function to prioritize in-distribution classes while suppressing updates for out-of-distribution classes, minimizing interference during local optimization.
• Ensure the method is communication-efficient and compatible with privacy-preserving FL by avoiding data sharing or model regularization beyond the loss function.

Figure 1: Illustration of catastrophic forgetting within client rounds. A global model with knowledge of all classes is sent to all clients participating in a given FL round. Local training increases the client model performance on the client’s local distribution but tends to simultaneously decrease

Experimental results

Research questions

• RQ1 To what extent does local client forgetting degrade global model performance in non-i.i.d. federated learning settings?
• RQ2 Can a per-client re-weighting of the softmax logits in the cross-entropy loss reduce catastrophic forgetting without requiring data sharing or model regularization?
• RQ3 How does the proposed WSM method perform under high data heterogeneity and low client participation—key challenges in real-world FL?
• RQ4 Does WSM improve generalization across clients and enhance convergence when applied to multiple FL algorithms beyond FedAvg?
• RQ5 What is the impact of local training steps and client participation rate on the effectiveness of WSM in mitigating forgetting?

Key findings

• WSM reduces local client forgetting by suppressing gradient updates for out-of-distribution classes, leading to more stable and generalizable client models.
• On CIFAR-100 with high data heterogeneity (α=0.1), FedAvg+WSM achieves 27.4% accuracy, a 7.6% improvement over FedAvg’s 19.8%.
• On CIFAR-10, FedAvg+WSM reaches 62.4% accuracy, a 1.02% improvement over FedAvg’s 60.8% under optimal learning rates.
• The performance gap between FedAvg and FedAvg+WSM is most pronounced at low client participation (e.g., 1%), with gains decreasing as more clients participate.
• Increasing local iterations from 7 to 21 causes a sharper accuracy drop in FedAvg, while FedAvg+WSM shows significantly more stable performance, indicating better forgetting control.
• WSM improves performance not only for FedAvg but also for SCAFFOLD and FedProx, demonstrating broad applicability across federated optimization frameworks.

Figure 2: CIFAR-10 - Distributions of ten randomly selected clients with data partitioned according to a Dirichlet distribution parameterized by $\alpha=0.1$ .