Fed-Focal Loss for imbalanced data classification in Federated Learning

TL;DR

This paper introduces Fed-Focal Loss, a novel loss function for federated learning that improves classification performance on imbalanced datasets by down-weighting well-classified examples and leveraging client contribution sampling.

Contribution

The paper proposes Fed-Focal Loss, a new loss function combined with a sampling framework to enhance robustness and accuracy in federated learning with imbalanced data.

Findings

Achieves over 9% improvement on unbalanced MNIST.

Demonstrates consistent performance gains across multiple benchmarks.

Applicable across various federated learning algorithms.

Abstract

The Federated Learning setting has a central server coordinating the training of a model on a network of devices. One of the challenges is variable training performance when the dataset has a class imbalance. In this paper, we address this by introducing a new loss function called Fed-Focal Loss. We propose to address the class imbalance by reshaping cross-entropy loss such that it down-weights the loss assigned to well-classified examples along the lines of focal loss. Additionally, by leveraging a tunable sampling framework, we take into account selective client model contributions on the central server to further focus the detector during training and hence improve its robustness. Using a detailed experimental analysis with the VIRTUAL (Variational Federated Multi-Task Learning) approach, we demonstrate consistently superior performance in both the balanced and unbalanced scenarios for MNIST, FEMNIST, VSN and HAR benchmarks. We obtain a more than 9% (absolute percentage) improvement in the unbalanced MNIST benchmark. We further show that our technique can be adopted across multiple Federated Learning algorithms to get improvements.