Addressing Label Shift in Distributed Learning via Entropy Regularization

TL;DR

This paper introduces VRLS, a novel entropy regularization method for mitigating label shift in distributed learning, significantly improving model accuracy under challenging label distribution changes across nodes.

Contribution

The paper proposes VRLS, a new entropy-based regularization technique that enhances label shift adaptation in distributed learning environments, with theoretical guarantees and practical improvements.

Findings

VRLS outperforms baselines by up to 20% on benchmark datasets.

It effectively mitigates label shift in multi-node distributed learning.

Theoretical bounds support the robustness of VRLS.

Abstract

We address the challenge of minimizing true risk in multi-node distributed learning. These systems are frequently exposed to both inter-node and intra-node label shifts, which present a critical obstacle to effectively optimizing model performance while ensuring that data remains confined to each node. To tackle this, we propose the Versatile Robust Label Shift (VRLS) method, which enhances the maximum likelihood estimation of the test-to-train label density ratio. VRLS incorporates Shannon entropy-based regularization and adjusts the density ratio during training to better handle label shifts at the test time. In multi-node learning environments, VRLS further extends its capabilities by learning and adapting density ratios across nodes, effectively mitigating label shifts and improving overall model performance. Experiments conducted on MNIST, Fashion MNIST, and CIFAR-10 demonstrate the effectiveness of VRLS, outperforming baselines by up to 20% in imbalanced settings. These results highlight the significant improvements VRLS offers in addressing label shifts. Our theoretical analysis further supports this by establishing high-probability bounds on estimation errors.