Efficient Image Representation Learning with Federated Sampled Softmax
This work addresses the computational and communication bottlenecks in federated learning for image representation, enabling efficient training with large class sets on decentralized data.
The paper tackled the problem of learning image representations on decentralized data with many classes in federated learning, where standard softmax is computationally expensive, and introduced federated sampled softmax (FedSS), which reduces parameter transfer and optimization while performing comparably to full softmax.
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.