Efficient Image Representation Learning with Federated Sampled Softmax

1 INTRODUCTION . The success of many computer vision applications , such as classification ( Kolesnikov et al. , 2020 ; Yao et al. , 2019 ; Huang et al. , 2016 ) , detection ( Lin et al. , 2014 ; Zhao et al. , 2019 ; Ouyang et al. , 2016 ) , and retrieval ( Sohn , 2016 ; Song et al. , 2016 ; Musgrave et al. , 2020 ) , relies heavily on the quality of the learned image representation . Many methods have been proposed to learn better image representation from centrally stored datasets . For example , the contrastive ( Chopra et al. , 2005 ) and the triplet losses ( Weinberger & Saul , 2009 ; Qian et al. , 2019 ) enforce local constraints among individual instances while taking a long time to train on O ( N2 ) pairs and O ( N3 ) triplets for N labeled training examples in a minibatch , respectively . A more efficient loss function for training image representations is the softmax cross entropy loss which involves only O ( N ) inputs . Today ’ s top performing computer vision models ( Kolesnikov et al. , 2020 ; Mahajan et al. , 2018 ; Sun et al. , 2017 ) are trained on centrally stored large-scale datasets using the classification loss . In particular , using an extremely large number of classes has proven to be beneficial for learning universal feature representations ( Sun et al. , 2017 ) . However , a few challenges arise when learning such image representations with the classification loss under the cross-device federated learning scenario ( Kairouz et al. , 2019 ) where the clients are edge devices with limited computational resources , such as smartphones . First , a typical client holds data from only a small subset of the classes due to the nature of non-IID data distribution among clients ( Hsieh et al. , 2020 ; Hsu et al. , 2019 ) . Second , as the size of the label space increase , the communication cost and computation operations required to train the model will grow proportionally . Particularly for ConvNets the total number of parameters in the model will be dominated by those in its classification layer ( Krizhevsky , 2014 ) . Given these constraints , for an FL algorithm to be practical it needs to be resilient to the growth of the problem scale . In this paper , we propose a method called federated sampled softmax ( FedSS ) for using the classification loss efficiently in the federated setting . Inspired by sampled softmax ( Bengio & Senécal , 2008 ) , which uses only a subset of the classes for training , we devise a client-driven negative class sampling mechanism and formulate a sampled softmax loss for federated learning . Figure 1 illustrates the core idea . The FL clients sample negative classes and request a sub network from the FL server by sending a set of class labels that anonymizes the clients ’ positive class labels in its local dataset . The clients then optimize a sampled softmax loss that involves both the clients ’ sampled negative classes as well as its local positive classes to approximate the global full softmax objective . To the best of our knowledge , this is the first work addressing the intersection of representation learning with Federated Learning and resource efficient sampled softmax training . Our contributions are : 1 . We propose a novel federated sampled softmax algorithm , which extends the image representation learning via large-scale classification loss to the federated learning scenario . 2 . Our method performs on-par with full softmax training , while requiring only a fraction of its cost . We evaluate our method empirically and show that less than 10 % of the parameters from the classification layer can be sufficient to get comparable performance . 3 . Our method is resilient to the growth of the label space and makes it feasible for applying Federated Learning to train image representation and classification models with large label spaces . 2 RELATED WORK . Large scale classification . The scale of a classification problem could be defined by the total number of classes involved , number of training samples available or both . Large vocabulary text classification is well studied in the natural language processing domain ( Bengio & Senécal , 2008 ; Liu et al. , 2017 ; Jean et al. , 2015 ; Zhang et al. , 2018 ) . On the contrary , image classification is well studied with small to medium number of classes ( LeCun et al. , 1998 ; Krizhevsky et al . ; Russakovsky et al. , 2015 ) while only a handful of works ( Kolesnikov et al. , 2020 ; Hinton et al. , 2015 ; Mahajan et al. , 2018 ; Sun et al. , 2017 ) address training with large number of classes . Training image classification with a significant number of classes requires a large amount of computational resources . For example , Sun et al . ( 2017 ) splits the last fully connected layer into sub layers , distributes them on multiple parameter servers and uses asynchronous SGD for distributed training on 50 GPUs . In this work , we focus on a cross-device FL scenario and adopt sampled softmax to make the problem affordable for the edge devices . Representation learning . Majority of works in learning image representation are based on classification loss ( Kolesnikov et al. , 2020 ; Hinton et al. , 2015 ; Mahajan et al. , 2018 ) and metric learning objectives ( Oh Song et al. , 2016 ; Qian et al. , 2019 ) . Using full softmax loss with a large number of classes in the FL setting can be very expensive and sometimes infeasible for two main reasons : ( i ) exorbitant cost of communication and storage on the clients can be imposed by the classification layer ’ s weight matrix ; ( ii ) edge devices like smartphones typically do not have computational resources required to train on such scale . On the other hand , for metric learning methods ( Oh Song et al. , 2016 ; Qian et al. , 2019 ) to be effective , extensive hard sample mining from quadratic/cubic combinations of the samples ( Sheng et al. , 2020 ; Schroff et al. , 2015 ; Qian et al. , 2019 ) is typically needed . This requires considerable computational resources as well . Our federated sampled softmax method addresses these issues by efficiently approximating the full softmax objective . Federated learning for large scale classification . The closest related work to ours is Yu et al . ( 2020 ) , which considers the classification problem with large number of classes in the FL setting . They make two assumptions : ( a ) every client holds data for a single fixed class label ( e.g . user identity ) ; ( b ) along with the feature extractor only the class representation corresponding to the client ’ s class label is transmitted to and optimized by the clients . We relax these assumptions in our work since we focus on learning generic image representation rather than individually sensitive users ’ embedding . We assume that the clients hold data from multiple classes and the full label space is known to all the clients as well as the FL server . In addition , instead of training individual class representations we formulate a sampled softmax objective to approximate the global full softmax cross-entropy objective . 3 METHOD . 3.1 BACKGROUND AND MOTIVATION . Softmax cross-entropy and the parameter dominance . Consider a multi-class classification problem with n classes where for a given input x only one class is correct y ∈ [ 0 , 1 ] n with∑n i=1 yi = 1 . We learn a classifier that computes a d-dimensional feature representation f ( x ) ∈ Rd and logit score oi = wTi f ( x ) + b ∈ R for every class i ∈ [ n ] . A softmax distribution is formed by the class probabilities computed from the logit scores using the softmax function pi = exp ( oi ) ∑n j=1 exp ( oj ) , i ∈ [ n ] . ( 1 ) Let t ∈ [ n ] be the target class label for the input x such that yt = 1 , the softmax cross-entropy loss for the training example ( x , y ) is defined as L ( x , y ) = − n∑ i=1 yi log pi = −ot + log n∑ j=1 exp ( oj ) . ( 2 ) The second term involves computing the logit score for all the n classes . As the number of classes n increase so does the number of columns in the weight matrix W ≡ [ w1 , w2 , . . . , wn ] ∈ Rd×n of the classification layer . The complexity of computing this full softmax loss also grows linearly . Moreover , for a typical ConvNet classifier for n classes , the classification layer dominates the total number of parameters in the model as n increases , because the convolutional layers typically have small filters and the total number of parameters ( See Figure 8 in A.1 for concrete examples ) . This motivates us to use an alternative loss function to overcome the growing compute and communication complexity in the cross-device federated learning scenario . Sampled softmax . Sampled softmax ( Bengio & Senécal , 2008 ) was originally proposed for training probabilistic language models on datasets with large vocabularies . It reduces the computation and memory requirement by approximating the class probabilities using a subset N of negative classes whose size is m ≡ |N | n. These negative classes are sampled from a proposal distribution Q , with qi being the sampling probability of the class i . Using the adjusted logits o′j = oj − log ( mqj ) , ∀j ∈ N , the target class probability can be approximated with p′t = exp ( o′t ) exp ( o′t ) + ∑ j∈N exp ( o ′ j ) . ( 3 ) This leads to the sampled softmax cross-entropy loss Lsampled ( x , y ) = −o′t + log ∑ j∈N∪ { t } exp ( o′j ) . ( 4 ) Note that the sampled softmax gradient is a biased estimator of the full softmax gradient . The bias decreases as m increases . The estimator is unbiased only when the negatives are sampled from the full softmax distribution ( Blanc & Rendle , 2018 ) or m→∞ ( Bengio & Senécal , 2008 ) . 3.2 FEDERATED SAMPLED SOFTMAX ( FEDSS ) . Now we discuss our proposed federated sampled softmax ( FedSS ) algorithm listed in Algorithm 1 , which adopts sampled softmax in the federated setting by incorporating negative sampling under FedAvg ( McMahan et al. , 2017 ) framework , the standard algorithm framework in federated learning . One of the main characteristics of FedAvg is that all the clients receive and optimize the exact same model . To allow efficient communication and local computing , our federated sampled softmax algorithm transmits a much smaller sub network to the FL clients for local optimization . Specifcally , we view ConvNet classifiers parameterized by θ = ( ϕ , W ) as two parts : a feature extractor f ( x ; ϕ ) : Rh×w×c → Rd parameterized by ϕ that computes a d-dimensional feature given an input image , and a linear classifier parameterized by a matrixW ∈ Rd×n that outputs logits for class prediction 1 . The FL clients , indexed by k , train sub networks parameterized by ( ϕ , WSk ) where WSk contains a subset of columns in W , rather than training the full model . With this design , federated sampled softmax is more communication-efficient than FedAvg since the full model is never transmitted to the clients , and more computation-efficient because the clients never compute gradients of the full model . In every FL round , every participating client first samples a set of negative classes Nk ⊂ [ n ] /Pk that does not overlap with the class labels Pk = { t : ( x , y ) ∈ Dk , yt = 1 , t ∈ [ n ] } in its local dataset Dk . The client then communicates the union of these two disjoint sets Sk = Pk ∪ Nk to the FL server for requesting a model for local optimization . The server subsequently sends back the sub network ( ϕ , WSk ) with all the parameters of the feature extractor together with a classification matrix that consists of class vectors corresponding to the labels in Sk . Algorithm 1 : Federated sampled softmax ( FEDSS ) . The key differences to the FedAvg are lines 5–7 where the clients request and optimize different sub networks locally . η and α are the client and server learning rates , respectively . 1 Initialize θ0 = ( ϕ , W ) , where ϕ is the parameter of the feature extractor and W is the classification matrix . 2 for each round t = 0 , 1 , . . . do 3 Select K participating clients . 4 for each client k = 1 , 2 , . . . , K do in parallel 5 Client k samples negativesNk . 6 Client k requests the model wrt Sk = Pk ∪Nk . 7 The server sends back model θ ( k ) t = ( ϕ , WSk ) . 8 Start local optimization with θ ( k ) ← θ ( k ) t . 9 for each local mini-batch b over E epochs do 10 θ ( k ) ← θ ( k ) − η∇L ( k ) sampled ( b ; θ ( k ) ) 11 ∆θ ( k ) ← θ ( k ) − θ ( k ) 0 12 ḡt ← ∑K k=1 nk n ∆θ ( k ) , where n = ∑K k=1 nk 13 θt+1 ← θt − αḡt Then every client trains its sub network by minimizing the following sampled softmax loss with its local dataset L ( k ) FedSS ( x , y ) = −o ′ t + log ∑ j∈Sk exp ( o′j ) , ( 5 ) 1We omit the bias term in discussion without loss of generality . after which the same procedure as FedAvg is used for aggregating model updates from all the participating clients . In our federated sampled softmax algorithm , the set of positive classes Pk is naturally constituted by all the class labels from the client ’ s local dataset , whereas the negative classes Nk are sampled by each client individually . Next we discuss negative sampling and the use of positive classes in the following two subsections respectively .

Summary: In this paper, the authors proposed federated sampled softmax (FedSS) for resource-efficient federated image representation learning. When the number of classes is large, the final classification layer could take up a large part of the communication cost during training. FedSS allows subsampling the weights of negative classes to reduce data transfer and leads to similar accuracy compared to the full softmax.