Efficient Image Representation Learning with Federated Sampled Softmax

Sagar M. Waghmare, Hang Qi, Huizhong Chen, Mikhail Sirotenko, Tomer Meron

Keywords: Federated learning, sampled softmax

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 computation and communication increases proportionally to the number of classes. In this work we introduce Federated Sampled Softmax, a novel resource-efficient approach for learning image representation with Federated Learning. Specifically, the FL clients sample a set of negative classes and optimize only the corresponding model parameters with respect to a sampled softmax objective that approximates the global full softmax objective. We analytically 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 in a privacy preserving way.

One-sentence Summary: We propose an algorithm for using sampled softmax in federated setting.