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: This paper investigates the problem of training a good computer vision model in the cross-device federated setting, focusing on classification. For deep networks, when the number of classes gets large ($\geq 10^3$), the number of parameters is dominated by the classification layers, hence scaling proportionally to the number of classes in the network. The authors propose to alleviate the resulting computation and communication burden by using sub-networks for each client, with a shared feature extractor but with only a smaller number of classes, re-using the sampled softmax proposed by Bengio & Senécal 2008. The resulting algorithm, Federated Sampled Softmax (fedSS), is benchmarked on image classification and image retrieval tasks along with baselines and variants. Experimental results demonstrate the validity of the approach.