Efficient Image Representation Learning with Federated Sampled Softmax

TL;DR

This paper introduces FedSS, a resource-efficient federated learning method that approximates full softmax for image representation learning, significantly reducing communication costs while maintaining performance.

Contribution

FedSS is the first federated sampled softmax approach that enables efficient learning of image representations with many classes in decentralized data settings.

Findings

Reduces communication and computation costs in federated learning.

Achieves comparable performance to full softmax methods.

Enables scalable image representation learning on decentralized data.

Abstract

Learning image representations on decentralized data can bring many benefits in cases where data cannot be aggregated across data silos. Softmax cross entropy loss is highly effective and commonly used for learning image representations. Using a large number of classes has proven to be particularly beneficial for the descriptive power of such representations in centralized learning. However, doing so on decentralized data with Federated Learning is not straightforward as the demand on FL clients' computation and communication increases proportionally to the number of classes. In this work we introduce federated sampled softmax (FedSS), a resource-efficient approach for learning image representation with Federated Learning. Specifically, the FL clients sample a set of classes and optimize only the corresponding model parameters with respect to a sampled softmax objective that approximates the global full softmax objective. We examine the loss formulation and empirically show that our method significantly reduces the number of parameters transferred to and optimized by the client devices, while performing on par with the standard full softmax method. This work creates a possibility for efficiently learning image representations on decentralized data with a large number of classes under the federated setting.