Efficient Image Representation Learning with Federated Sampled Softmax
Authors: Sagar M. Waghmare
Hang Qi
Huizhong Chen
Mikhail Sirotenko
Tomer Meron
Venue: ICLR 2022 Submitted (2022)
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.

Under review as a conference paper at ICLR 2022
EFFICIENT IMAGE REPRESENTATION LEARNING WITH FEDERATED SAMPLED SOFTMAX
Anonymous authors Paper under double-blind review
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 large number of classes has proven to be particularly beneficial for the descriptive power of such representations in centralized learning. However, do- ing so on decentralized data with Federated Learning is not straightforward as the demand FL clients' computation and communication increases propor- tionally to the number of classes. In this work we introduce federated sampled softmax (FedSS), 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 sampled soft- max 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 pos- sibility for efficienlly learning image representations on decentralized data with a large number of classes under the federated setting.
INTRODUCTION
The success of many computer vision applications, such as classification (Kolesnikov et al., 2020; Yao et al., 2019; Huang et al., 2016), detection (Linn 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 pairs and O(N3) triplets for N labeled training examples in 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 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, 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 small subset of the classes due to the nature of non-IlI) 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 proportion- ally. 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 FL algorithm to be practical it needs to be resilient to the growth the problem scale.
In this paper, we propose method called federated sampled softmax (FedSS) for using the classi- fication loss efficiently in the federated setting. Inspired by sampled softmax (Bengio & Senécal, 2008), which uses only subset of the classes for training, we devise client-driven negative class

---

Under review as a conference paper at ICLR 2022
Server
ConvNet ((
W
Client requests model wro disclosing data.
Sk
((K)
(9, Ws.)
(g(
model perameters wrt St.
for server
Sk = Pk
U
local positives
sampled negatives
ConvNet f(z;0)
(() sumpled
babby
{x, y)}
Ws.
FL client
Figure 1:
An FedSS training round: The client sends set of obfuscated class labels Sg the FL server and receives the feature extractor and few columns Ws,, corresponding to classes in Sk. from the weight matrix of the classification layer. The client optimizes this sub network with the sampled softmax loss and then communicates back the model update to the server. The server aggregates the model updates from all the selected clients to construct a new global model for the next round.
sampling mechanism and formulate sampled softmax loss for federated learning. Figure 1 illus- trates the core idea. The FL clients sample negative classes and request sub network from the FL server by sending a set of class labels that anonymizes the clients' positive class labels in local dataset. The clients then optimize 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 knowledge, this is the first work addressing the intersection of representation learning with Federated Learning and resource efficient sampled softmax training. Our contributions are:
We propose a novel federated sampled softmax algorithm, which extends the image repre- sentation 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 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.
RELATED WORK
Large scale classification. The scale of 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 & Sené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 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 clas- sification with 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 multiple parameter servers and uses asynchronous SGD for distributed training 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.
2

---

Under review as a conference paper at ICLR 2022
Representation learning. Majority of works in learning image representation based on classifi- cation 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 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 the clients can be imposed by the classification layer's weight matrix; (ii) edge devices like smartphones typically do not have computational re- sources 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 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 multi-class classification problem with n classes where for a given input æ only one class is correct y E [0. 1]" with Ex11% 3/: I. We learn a classifier that computes d-dimensional feature representation f(x) Rd
and logit score o, wi f()) + IR for every class E [n]. A softmax distribution is formed by the class probabilities computed from the logit scores using the softmax function
exp(o:) Pj Σド= exp(o,)
i [n].
(I)
Let t E [n] be the target class label for the input such that y 1, the softmax cross-entropy loss for the training example (x. y) is defined as
f(x, v) = Σ logn: = -0, + log exp(o,).
logo=
(2)
J=1
The second term involves computing the logit score for all the n classes. As the number of classes increase so does the number of columns in the weight matrix W = [Wi, W2.... Wn] Rdx= of the classification layer. The complexity of computing this full softmax loss also grows linearly.
Moreover, for a typical ConvNet classifier for 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.I for concrete examples). This motivates us to use an alternative loss function to overcome the growing compute and communica- tion complexity in the cross-device federated learning scenario.
Sampled softmax. Sampled softmax (Bengio & Senécal, 2008) was originally proposed for training probabilistic language models on datasets with large vocabularies. It reduces the com- putation and memory requirement by approximating the class probabilities using a subset W of negative classes whose size is | & n. These negative classes are sampled from a pro- posal distribution Q, with being the sampling probability of the class i. Using the adjusted logits og o log(mq;), Vj N, the target class probability can be approximated with
exp(o))
1/ exp(o)) EyEN exp(o))
(3)
3

---

Under review as a conference paper at ICLR 2022
This leads to the sampled softmax cross-entropy loss
Sampled (2. y) -o; log E
exp(o))-
JENU(H)
Note that the sampled softmax gradient 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 - so (Bengio & Sené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 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. Specifeally, we view ConvNet classifiers parameterized by 0 = (vp, W) as two parts: feature extractor fæ; 4p) hxxwxx xe + Rd parameterized by yp that computes a d-dimensional feature given an input image, and linear classifier parameterized by matrix W E Rd that outputs logits for class prediction ! The FL clients, indexed by k, train sub networks parameterized by (4p, Ws.) where Ws, contains subset of columns in W, rather than training the full model. With this design, federated sampled softmax is more coaamuiiiaatoonefffiiie than FedAvg since the full model is never transmitted to the clients, and more computation eefficcent because the clients never compute gradients of the full model.
In every FL round, every participating client first samples set of negative classes N₂. C [n] /Px that does not overlap with the class labels Pa = (t (x. y) E Dr, Vr = 1, E [n]] in its local dataset Dz. The client then communicates the union of these two disjoint sets S2 Ps U Na. to the FL server for requesting a model for local optimization. The server subsequently sends back the sub network (ip, Ws.) with all the parameters of the feature extractor together with a classification matrix that consists of class vectors corresponding to the labels in Sy.
Algorithm I: Federated sampled softmax (FEDSS). The key differences to the FedAvg are lines 5-7 where the clients request and optimize diffcrent sub networks locally. n and are the client and server leamning rates, respectively.
Initialize 0, = (1p, W), where &p is the parameter of the feature extractor and W is the classification matrix.
for each round t l....dd.. do
Select participating clients.
for each client 1, 2, K do in parallel
Client & samples negatives Ng.
Client & requests the model wrt S₁ Px U NN.
The server sends back model 0(*)) (1p, Ws,).
Start local optimization with g(h) +- (!4)
for each local mini-batch b over E epochs do
g(h) + (()) v/V ((t)
sampled (b; g(t)
10
(K(K) - g(k) o
+- ΣK k=1 (0(k), where ΣK =k=1
12
M++1= + agt 13
Then every client trains its sub network by minimizing the following sampled softmax loss with its local dataset
(k) L (x, y) =-6 log Σ
exp(o)),
(5)
jESz
We omit the hias in discussion without loss of generality.
4

---

Under review as a conference paper at ICLR 2022
after which the same procedure FedAvg is used for aggregating model updates from all the par- ticipating 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 's local dataset, whereas the negative classes Ngg are sampled by each client individually. Next we discuss negative sampling and the use of positive classes in the following two subsections respectively.
3.3 CLIENT-DRIVEN UNIFORM SAMPLING OF NEGATIVE CLASSES
For centralized learning, proposal distributions and sampling algorithms are designed for efficient sampling of negatives or high quality estimations of the full softmax gradients. For example, Jean et al. (2015) partition the training corpus and define non-overlapping subsets of class labels as sam- pling pools. The algorithm is efficient once implemented, but the proposal distribution imposes sampling bias which is not mitigable even as m + oo. Alternatively, efficient kernel-based algo- rithms (Blane & Rendle, 2018; Rawat et al., 2019) yield unbiased estimators of the full softmax gradients by sampling from the softmax distribution. These algorithms depend on both the current model parameters (4p, W) and the current raw input for computing feature vectors and logit scores. However, this is not feasible in the FL. scenario, one the one hand due to lack of resources on FL. clients for receiving the full model, on the other hand due to the constraint of keeping raw inputs only on the devices.
In the FedSS algorithm, we assume the label space is known and take client-driven approach, where every participating FL client uniformly samples negative classes N from [n]/Pk- Using a uniform distribution over the entire label space is simple yet effective choice that does not incur sampling bias. The bias on the gradient estimation can be mitigated by increasing m (See A.3 for an empirical analysis). Moreover, N₁ can be viewed as noisy samples from the maximum entropy dis- tribution over [n]/PP PP that mask the client's positive class labels. From the server's 's perspective, is not able to identify which labels in Sy belong to the client's 's dataset. In practice, private information retrieval techniques (Chor et al., 1995) can further be used such that no identity information about the set is revealed to the server. The sampling procedure can be performed on every client locally and independently without requiring peer information or the current latest model from the server.
3.4 INCLUSION OF POSITIVES LOCAL OPTIMIZATION
When computing the federated sampled softmax loss, including the set of positive class labels Pk in Eq. is crucial. To see this, Eq. 5 can be equivalently written as follows (shown in A.6)
Fedss (x, 3) log 1+
Σ jES+/(r)
explo - oi) +))
(6)
Minimizing this loss function pulls the input image representation f(æ;; (p) and target class repre- sentation w, closer, while pushing the representations of the negative classes Ws./tg away from f(x;; ()) Utilizing Pk/ {t} as an additional set of negatives to compute this loss encourages the separation of classes in Pk with respect to each other as well as with respect the classes in Ag (Figure 2d).
Alternatively, not using Pa/(t) as additional negatives leads negatives-only loss function
Clk) (x. v) = logg 1 Σ JENL
exp(o' - of)
where t E PE only contributes to computing the true logit for individual inputs, while the same Ng is shared across all inputs (Figure 2b). Minimizing this negatives-only loss, trivial solutions can be found for a client's local optimization. Because it encourages separation of target class representations Wp, from the negative class representations WN,, which can be easily achieved by increasing the magnitudes of the former and reducing those of the latter. In addition, the learned representations can collapse, as the local optimization is reduced to a binary classification problem between the on-client classes Ps and the off-client classes Ni-
5

---

Under review as a conference paper at ICLR 2022
(a) Input-dependent (1x y)
(b) NegOnly N
(c) PasOnly
(d) FedSS (Ours) Px W&/((()
Pk/(t)
Ps
The set of classes providing pushing forces for the local training under different sampled softmax loss formulations. (a) Input-dependent negative classes (depicted by the red squares) are sampled wrt to the inputs and current model, not feasible the FL setting.
Figure 2:
(b) Only using the sampled negatives reduces the problem to binary classification. (c) Using only the local positives lets the local objectives diverge from the global one. (d) FedSS approximates the global objective with sampled negative classes together with local positives.
In contrast, using only the local positives P3 without the sampled negative classes N₂ gives
c(k) Prn (xr y) log 1
Σ
explo - -d)
jEPL/(4)
Minimizing this loss function solves the client's local classification problem which diverges from the global objective (Figure 2c), especially when Pa remains fixed over FL rounds and [Px] < n.
EXPERIMENTS
4.1
SETUP
Notations and Baseline methods. We denote our proposed algorithm as FedSS where both the sampled negatives and the local positives are used in computing the client's sampled softmax loss.
We compare our method with the following alternatives:
NegOnly: The client's objective is defined by sampled negative classes only (Eq. 7).
PosOnly: The client's objective is defined by the local positive classes only, no negative classes is sampled (Eq. 8).
FedAwS (Yu et al., 2020): client optimization is same as the PosOnly, but a spreadout regularization is applied server.
In addition, we also provide two reference baselines:
The client's objective is the full softmax cross-entropy loss (Eq. 2), serving as performance references when it is affordable for clients to compute the full model.
FullSoftmax:
Centralized: model is trained with the full softmax cross-entropy loss (Eq. 2) in tralized fashion using IID data batches.
Evaluation protocol. We conduct experiments on two computer vision tasks: multi-class image classification and image retrieval. Performance is evaluated on the test splits of the datasets, which have no sample overlap with the corresponding training splits. We report the mean and standard devi- ation of the performance metrics from three independent runs. For the FullSoftmax and Centralized baselines, we report the best result from three independent runs. Please see A.2 for implementation details.
4.2 MULTI-CLASS IMAGE CLASSIFICATION
For multi- class classification we use the Landmarks- Londdmarks-User--ee 160K (Hsu et al., 2020) and report top-1 accuracy its test split. Landmarks-User--Iee 160k is landmark recognition dataset created for FL.

---

Under review as a conference paper at ICLR 2022
(a) Landmarks, [5 110
(b) SOP, |5 = 40
Tap-1 Accuracy
0.6
Fed55 (Ours)
NegOnly
0.25
PasOnily
MAP@10
0.20
0.15
Full Softmax
0.10
0.05
0.0
1000
3000
0.00
500
750
1000
1250
1500
1750
FL rounds
FL rounds
Figure 3: Learning curve for different methods for an average value of number of classes |S&| on the clients. The PosOnly, FedAwS and FullSoftmax methods have [Pal. [Pyl and n classes respectively,, on the clients.
simulations. consists of 1,262 natural clients based on image authorship. Collectively, every client contains 130 images distributed across 90 class labels. For our experiments K 64 clients are randomly selected to participate in each FL round. We train for total 5,000 rounds, which is sufficient for reaching convergence.
[S
95
100
110
130
170
% of n
(4.7%)
(4.9%)
(5.4%)
(6.4%)
(8,4%)
FedSS (Ours)
51.7 +004
53.3
4:06
54.9
+0.3
55.3
+0.6
56.0
+-0.06
NegOnly
7.1 +:3.7
18.7
+0.4
22.0 +0.8
25.0 +0.4
26.5 +1.4
PosOnly
43.1 #0.2
FedAwS (Yu et al., 2020)
42.5 10.4
FullSoftmax
56.8
Centralized
59.5
Table I: Top-1 accuracy (%) on Landmarks-Usen--166l at the end of 5k FL rounds. PosOnly and FedAwS have ~4.4% of class representations on the clients, whereas, FullSoftmax has all the class representations.
Table summarizes the top-1 accuracy the test split. For FedSS and NegOnly we report aceu- racy across different [Sg]. Overall, we observe that our method performs similar to the FullSoftmax baseline while requiring only fraction of the classes the clients. Our FedSS formulation also outperforms the alternative NegOnly, PosOnly and FedAwS formulations by a large margin. Ap- proximating the full softmax loss with FedSS does not degrade the rate of convergence either seen in Figure 3a. Additionally. Figure 4a shows learning curves for FedSS with different [Sy]. Learning with sufficiently large [S&| follows closely the performance the FullSoftmax baseline, We also report performance ImageNet-21k (Deng et al., 2009) in A.4.
4.3
IMAGE RETRIEVAL
25 (0.22%)
40
100 (0.88%)
% of n
(0.27%)
(0.35%)
(0.53%)
FedSS (Ours)
25.2 +0.2
25.8 +0.2
26.1 +0.1
26.4 +0.12
26.5 +0.03
NegOnly
15.5
+0.2
16.2 +0.1
16.3 40.1
16.5 40.01
16.7 40.17
PosOnly
19.7
+0.09
FedAwS (Yu et al., 2020)
20.0 +0.04
FullSoftmax
25.7
Centralized
25.4
Table 2: MAP@ 10 the SOP dataset the end of 2k FL rounds.
The Stanford Online Products dataset (Song et al., 2016) has 120,053 images of 22,634 online products the classes. The train split includes 59,551 images from 11,318 classes, while the test

---

Under review as a conference paper at ICLR 2022
0.27
Top-1 Accuracy
0.55
0.26
0.50
95
25
100
0.24
0.23
130
60
0.22
0.35
FullSoftmax
FullSoftmax
0.30
1000
2000
3000
4000
5000
0.20
250
500
750
1000
1250
1500
1750
2000
FL rounds
FL rounds
FedSS (ours) learning for different |Sy
Figure 4: Convergence curves for the proposed FedSS method at different cardinalities of Sg. Given that Pk is fixed for a client, the increase in |S& is caused by increase in |W&l)- The estimate of softmax probability via sampled softmax improves with the increase in [Sk], and therefore improving the efficacy of the method.
[a] Landmarks
(b) SDP
0.300
FullSoftmax
0.275
FullSffma
0.5
Top-1 Accuracy
0.4
0.225
0.3
0.175
FedSS (Ours)
(Ours)
0.2
NegOnly
0.125
NegOnly
PasOnly
PosOnly
0.1
FedAwS
FedAwS
5.8 6.4 7.0
8.3
10.9
x
1.3 1.9 2.6
3.8
6,4
424.4(x10"]
# of params in the classification layer
# of params in the classification layer
Figure 5: Performance vs number of parameters in the classification layer transmitted to and opti- mized by the clients for Landmarks- -aaSS--UUUS6666 (a) and the SOP (b) datasets, respectively.
split includes 11,316 different classes with 60,502 images total. For FL experiments, we partition the train split into 596 clients, each containing 100 images distributed across 20 class labels. For each FL round, nearest neighbor retrieval to evaluate the models. Every image in the test split is used as query image against the remaining ones. We use normalized euclidean distance to compare two image representations. We report MAP@ Rt (R 10) as the evaluation metric (Musgrave et al., 2020), which is defined as follows:
32 clients are randomly selected. Similar metric learning literature, we use
R MAP@
precision i, Rh retrieval correct 0.
P(), R
where P(t)
(9)
otherwise.
1=1
Table 2 summarizes MAP@ 10 the SOP test split at the end of FL rounds. Our FedSS formu- lation consistently outperforms the alternative methods while requiring less than 1% of the classes on the clients. This reduces the overall communication cost by 16% when [Sy] 100 for every client round. For reasonably small value of |Sy our method has a similar rate of convergence to the FullSoftmax baseline, as seen Figure 3b and Figure 4b.
Using the MobilenetV3 (Howard et al., 2019) architecture with embedding size 64, the classification layer contributes to 16% of the total number of parameters in the SOP experiment and 3.4% in the Landarks- - 160k experiment. In the former, our FedSS method requires only 84% of the model parameters on every client per round when Syl = 100. In the latter, it reduces the model parameters transmitted by 3.38% per client per round when S& 170 (summarized in Figure 5). These savings will increase as the embedding size or the total number of classes increases (Figure 8 in A.1). For example with embedding size of 1280, which is default embedding size of MobileNetV3, above setup will result 79% and 38% reduction in the communication cost per client per round for the SOP and Landarks-User- 160k datasets, respectively.
8

---

Under review as a conference paper at ICLR 2022
4.4 ON IMPORTANCE OF PA LOCAL OPTIMIZATION
One may note that the NegOnly loss (Eq. 7) involves fewer terms inside the logarithm than FedSS (Eq. 6). To show that the NegOnly is unfairly penalized, we compare the FedSS with NegOnly such that the number of classes providing pushing forces for every input is the same. This is done by sampling additional [P&] negative classes for the NegOnly method. As seen in Figure 6, using the on-client classes (Px) as additional negatives instead of the additional off-client negatives crucial to the learning.
round
Top-1 Accuracy
0.4
MAPgg10
0.20
0.30
0.3
0.0
20
40
0.00
Number sampled negatives (n=2028)
Number of sampled regatives (m=11318)
FedSs (Ours):
[P] negative
NegOnly: [AV]] negative classes
Negonly: [P negative classes
Figure 6: Performance of the FedSS (Ours) and NegOnly methods with different compositions of the negative classes used for computing the sampled softmax loss. Utilizing on-client classes additional negatives i.e, FedSS method, has superior performance to the NegOnly method with equivalent number of negatives.
FedSS
NegOnly
1.0
True
0.4
-0.0
Predicted
Predicted
Figure 7: Confusion matrices for Pa the same client from Landmarks-User 160k dataset. both the FedSS and NegOnly formulations we used [S 95. In the former, the class representations are learned and well-separated, but are collapsed in the latter.
This boost can be attributed to better approximation of the global objective by the clients. Figure plots client's confusion matrix corresponding to the FedSS and NegOnly methods. The NegOnly loss leads to trivial solution for the client's local optimization problem such that the client's 's positive class representations collapse onto one representation, as reasoned in section 3.4.
CONCLUSION
Federated Learning is becoming prominent field of research. Major contributing factors this trend are: rise in privacy awareness among the general users, surge in amount of data generated by edge devices, and the noteworthy increase in computing capabilities of edge devices. In this work we presented novel federated sampled softmax method which facilitates efficient training of large models on edge devices with Federated Learning. The clients solve small subproblems approximating the global problem by sampling negative classes and optimizing a sampled softmax objective. Our method significantly reduces the number of parameters transferred to and optimized by the clients, while performing par with the standard full softmax method. We hope that this en- couraging result can inform future research on efficient local optimization beyond the classification layer.

---

Under review as a conference paper at ICLR 2022
REFERENCES
Yoshua Bengio and ccan-Sébasticn Senécal. Adaptive importance sampling accelerate training of neural probabilistic language model. IEEE Transactions on Neural Networks, 19(4):713-722, 2008.
Guy Blanc and Steffen Rendle, Adaptive sampled softmax with kernel hased sampling. In JCMI., 2018.
Sumit Chopra, Raia Hadsell, and Yann LeCun. Learning similarity metric discriminatively, with application to face verification. In CVPR, 2005.
Benny Chor, Oded Goldreich, Eyal Kushilevitz, and Madhu Sudan. Private information retrieval. In Annual Foundations of Computer Science, 1995.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. ImageNet: A large-scale hierarchical image database. In CVPR, 2009.
Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in neural network. stat, 1050:9, 2015.
Andrew Howard, Mark Sandler, Grace Chu, Liang-Chieh Chen, Bo Chen, et al. Searching for MobileNetV3. In ICCV, 2019.
Kevin Hsieh, Amar Phanishayee, Onur Mutlu, and Phillip B. Gibbons. The non-IID data quagmire of decen- tralized machine learning. arXiv preprint urXiv:1404.599) 2020.
Tzu-Ming Harry Hsu, Hang Qi, and Matthew Brown. Measuring the effects of non-identical data distribution for federated visual classification. arXiv preprint 1rX.v:1909.06335, 2019.
Tzu-Ming Harry Hsu, Hang Qi, and Matthew Brown. distribution. In ECCV, 2020.
Federated visual classification with real-world data
Chen Huang, Yining Li, Chen Change Loy, and Xiaoou Tang. Learning deep representation for imbalanced classification, In CVPR, 2016.
Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In ICML, 2015.
Sébastien Jean, Kyunghyun Cho, Roland Memisevic, and Yoshua Bengio. On using very large target vocabulary for neural machine translation. In ACL and IJCNPP, pp. 1-10, 2015.
Peter Kairouz, H Brendan McMahan, et al. Advances and open problems in federated learning. arXiv preprint arXiv:1912.04977, 2019.
Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Joan Puigcerver, Jessica Yung. Sylvain Gelly, and Neil Houlsby. Big transfer (BiT): General visual representation learning. In ECCV, 2020.
Alex Krizhevsky. One weird trick for parallelizing convolutional neural networks. arXiv preprint arXiv:1404.5997, 2014.
Alex Krizhevsky, Vinod Nair, and Geoffrey Hinton. CIFAR-100 (canadian institute for advanced research). URL http:/ttp//c /ww. toronto......o.oo...... edu/ kri 2 ci siirar... html.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-basse learning applied to document 2324, 1998.
recognition. Proceedings of the IEEE,
Tsung-Yi Yi Lin, Michael Maire, Serge Belongic, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C Lawrence Zimick. Microsoft COCO; Common objects in context. In ECCV, 2014.
Jingzhou Liu, Wei-Cheng Chang, Yuexin Wu, and Yiming Yang. Deep learning for extreme multi- labti-label text classification. In SIGIR, 2017.
Dhruv Mahajan, Ross Girshick, Vignesh Ramanathan, Kaiming He, Manohar Paluri, Yixuan Li, Ashwin Bharambe, and Laurens van der Maaten. Exploring the limits of weakly supervised pretraining. In ECCV, 2018.
Brendan MeMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-cfficionnn learning of deep networks from decentralized data. AISTATS, 2017.
Kevin Musgrave, Serge Belongie, and Ser-Nam Lim. A metric learning reality check. arXiv preprint arXiv:2003.08500 2020.
10

---

Under review as a conference paper at ICLR 2022
Hyun Oh Song. Yu Xiang, Stefanie Jegelka, and Silvio Savarese. Deep metric learning via lifted structured feature embedding. In CVPR, 2016.
Wanli Ouyang, Xiaogang Wang, Cong Zhang, and Xiaokang Yang. Factors in finetuning deep model for object detection with long-tail distribution. CVPR, 2016.
Qi Qian, Lei Shang, Baigui Sun, Juhua Hu, Hao Li, and Rong Jin. SoffTriple loss: Deep metric learning without triplet sampling. In ICCV, 2019.
Ankit Singh Rawat, Jiecao Chen, Felix Xinman X Yu, Ananda Theertha Suresh, and Sanjiv Kumar. Sampled softmax with random fourier features. In NeurIPS, 2019.
Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, ImageNet large scale visual recognition challenge. IJCV, 115(3), 2015.
Florian Schroff, Dmitry Kalenichenko, and James Philbin. FaceNet: unified embedding for face recognition and clustering. In CVPR, 2015.
Hao Sheng. Yanwei Zheng. Wei Ke, Dongxiao Yu, Xiuzhen Cheng. Weifeng Lyu, and Zhang Xiong, Mining hard samples globally and efficiently for person reidentification. IEEE Internet of Things Journal, 7(10): 9611-9622, 2020.
Kihyuk Sohn. Improved deep metric learning with multi-class n-pair loss objective. In NIPS, pp. 1857-1865, 2016.
Hyun Oh Song, Yu Xiang, Stefanie Jegelka, and Silvio Savarese. Deep metric learning via lifted structured feature embedding. In CVPR, 2016.
Chen Sun, Abhinav Shrivastava, Saurabh Singh, and Abhinav Gupta. Revisiting unreasonable effectiveness of data in deep learning era. In ICCV, 2017.
Feng Wang, Xiang Xiang, Jian Cheng, and Alan Loddon Yuille. NormFace: L2 hypersphere embedding for face verification. In ACM Multimedia, 2017.
Kilian Q Weinberger and Lawrence Saul. Distance metric learning for large margin nearest neighbor classi- fication. JMLR, 10(2), 2009.
Yuxin Wu and Kaiming He. Group normalization. In ECCV, 2018.
Hantao Yao, Shiliang Zhang, Richang Hong, Yongdong Zhang, Changsheng Xu, and Qi Tian. Deep represen- lation learning with part loss for person re-identification. IEEE Transactions on Image Processing, 28(6): 2860-2871, 2019.
Felix Yu, Ankit Singh Rawat, Aditya Menon, and Sanjiv Kumar. Federated learning with only positive lahels. In ICML, 2020.
Wenjie Zhang, Junchi Yan, Xiangfeng Wang, and Hongyuan Zha. Deep extreme multi-label learning. In ACM International Conference on Multimedia Retrieval, pp. 100-107, 2018.
Zhong- Qiu Zhao, Peng Zheng, Shou-tao Xu, and Xindong Wu. Object detection with deep learning: A review. IEEE Transactions on Neural Networks and Learning Systems, 30(1 1):3212-3232, 2019.

---

Under review as a conference paper at ICLR 2022
A SUPPLEMENTARY MATERIAL
A.1 PARAMETERS IN THE LAST LAYER
The number of parameters in the classification layer grows linearly with respect to the number of classes and typically dominates the total number of parameters in the model. Figure 8 shows the number of parameters in the classification layer as the percentage of total number of parameters in the MobileNerV3 model. Each curve shows the percentage for different number of target classes for a fixed embedding size.
100%
% paramelers clasiiication
60%
40%
d-64
d-128
d-256
20%
d=512
d=1280
0%
1000
4000
8000
16000
32000
64000
Number of classes (n)
Figure 8: The number of parameters in the classification layer dominates the model as the number of classes n grows. We show the percentage of parameters in the last layer using the MobileNetV3 architecture (Howard et al., 2019) while varying the number classes n and dimension d of the feature (d = 1280 is the default dimensionality of MobileNerV3)..
is obvious that as the number of classes or the size of image representation increases so does the communication and local optimization cost for the full softmax training in the federated setting. In either of these situations our proposed method will facilitate training at significantly lower cost.
A.2 IMPLEMENTATION DETAILS
For all the datasets we use the default MobileNetV3 architecture (Howard et al., 2019), except that instead of 1280 dimensional embedding we output 64 dimensional embedding. We replace Batch Normalization (loffe & Szegedy, 2015) with Group Normalization (Wu & He, 2018) to improve the stability of federated learning (Hsu et al., 2019; Hsich et al., 2020). Input images are resized to 256x256 from which random crop of size 224x224 is taken. All ImageNet-21k trainings start from scratch, whereas, for Landmarks- Uses---U--16000 and the SOP we start from a ImageNet- k (Rus- sakovsky al., 2015) pretrained checkpoint. For client side optimization data once and use stochastic gradient descent optimizer with batchsize of 32. We use the learn- ing rate of 0.01 for the SOP and Landmarks-User-166 160k. All ImageNet-21k experiments start from scratch and use the same learning rate of 0.001. To have fair comparison with FedAwS method we do hyperparameter search to find the best spreadout weight and report the performances corre- sponding to it. For all the experiments, we use scaled cosine similarity with fixed scale value (Wang al., 2017) 7) of 20 for computing the logits; the server side optimization is done using Momentum optimizer with learning rate of 1.0 and momentum of 0.9. All Centralized baselines are trained with stochastic gradient descent.
go through the local
For given dataset, all the FL methods are trained for a fixed number of rounds. The corresponding centralized experiment trained for an equivalent number of model updates.
A.3 FEDSS GRADIENT NOISE ANALYSIS
Bengio & Senécal (2008) provides theoretical analysis of convergence of the sampled softmax loss. Doing so for the porposed federated sampled softmax within the FedAvg framework is beyond
12

---

Under review as a conference paper at ICLR 2022
FedSS convergence analysis with gradient noise
0.12
Fed5S gradient noise
0.10
0.08
0.06
0.04
0.02
0.00
2000
4000
6000
8000
10000
[
Figure 9: Empirical FedSS gradient noise analysis. As increase the sample size the difference between FedAvg (with FullSoftmax) and FedSS diminishes.
the scope of this work. Instead we provide an empirical gradient noise analysis for the proposed method. To do so we compute the expected difference between FedAvg (with FullSoftmax) and FedSS gradients, i.e. E((gFedAvg 9Fedss]). where gFedAug and @FedSs are client model changes aggregated by the server for FedAvg (with FullSoftmax) and FedSS methods, respectively. Given that FedSS is an estimate of FedAvg (with FullSoftmax) this difference essentially represents the noise in FedSS gradients.
To compute single instance of gradient noise we assume that the clients participating in the FL round has same DD with |D] 32. Please note that the clients will have different Ng. For given WW&| we compute the expectation of the gradient noise across multiple batches (D) of the SOP dataset. Figure 9 shows the FedSS gradient noise as a function of [Wt] For very small values of [&& the gradients can be noisy but as the [W&] increases the gradient noise drops exponentially.
A.4 IMAGENET-2IC EXPERIMENTS
Along with Landmarks- User- 160K (Hsu al., 2020) and the SOP (Song et al., 2016) datasets we also experiment with ImageNet-21k (Deng et al., 2009) dataset. It is super set of the widely used ImageNet- 1k (Russakovsky al., 2015) dataset. contains 14.2 million images distributed across 21 classes organized by the WordNet hierarchy. For every class we do random 80-20 split on its samples to generate the train and test splits, respectively. The train split used generate 25,691 clients, each containing approximately 400 images distributed across 20 class labels. ImageNet-211 requires large number of FL rounds given abundant training images, hence we set training budget of 25,000 FL rounds to make experiments manageable. Although the performance report on ImageNet-211
not comparable with the (converged) oate-off-heeart we emphasize that the setup is sufficient evaluate our FedSS method and demonstrate its effectiveness.
15 % of n
120
220
820
(0.3%)
(0.5%
(1.0%)
(1.9%)
(3.7%)
FedSS (Ours)
9.1 +0.4
9.2 +0.1
9.9 +0.3
10.0 40.5
40.5
NegOnly
3.9 +0.1
4.2 40.1
4.3 +0.2
4.4 +0.1
4.7 40.2
PosOnly
5.1 +0.4
FedAwS (Yu et al., 2020)
5.1 +0.1
FullSoftmax
11.3
Centralized
15.4
Table Top-1 accuracy (%) on ImageNet-21kk at the end of 25k FL rounds. PosOnly and FedAwS have ~0.1 % of class representations on the clients, whereas, FullSoftmax has all the ~21k class representations.
13

---

Under review as a conference paper at ICLR 2022
0.14
FullSoftmax
0.12
Top-1 Accuracy
0.10
0.08
0.06
FedSS (Ours)
0.04
NegOnly
0.02
PosOnly
FedAwS
0.00
1.3 7.7 14.1
26.9
52.5 1398 103]
# of params in the classification layer
Figure 10: ImageNet-21k: Top-1 accuracy vs number of parameters the classification layer trans- mitted and optimized by the clients.
Table 3 summarizes top-1 accuracy on the ImageNet-2lk test split. We experiment with five different choices of Syl. The FullSoftmax method reaches (best) top-1 accuracy of 11.30% by the end of 25,000 FL. rounds, while our method achieves top-1 accuracy 10.02 0.5%, but with less than 2% of the classes the clients. Figure 10 summarizess performance of different methods with respect to number of parameters in the classification layer transmitted to and optimized by the clients. Our client-driven negative sampling with positive inclusion method (FedSS) requires very small fraction of parameters in the classification layer while performing reasonably similar to the full softmax training (FullSoftmax).
A.5 OVERFITTING IN THE SOP FULLSOFTMAX EXPERIMENTS
The class labels in the train and test splits of the SOP dataset do not overlap. In addition, it has, on average, only 5 images per class label. This makes the SOP dataset susceptible to overfitting (Table 4). In this case, using FedSS mitigates the overfitting as only subset of class representations is updated every FL round.
Method
Top-1 Accuracy (train)
MAP@ 10 (test)
FedSS (Ours)
97.6 +0.2
26.5 +0.03
FullSoftmax
99.9
25.7
Centralized
99.9
25.4
Table 4: Top-1 accuracy the train split and corresponding MAP@ 10 the test split for the SOP dataset at the end of 2k FL rounds. The FedSS shown here is trained on |Sa] = 100.
14

---

Under review as a conference paper at ICLR 2022
A.6 DERIVATIONS FROM EQ. 5 TO EQ. 6
Proof. Starting from Eq. 5, we have
(k) Fedss (z. y) -0g + log Σ
exp(o))
jESu
= log
exp( -0))
Σ
explo))
jES
log
Σ
explo) di)
jESe
log
exp(d; - o)) + Σ
explo -000--)) oi)
jES&/(()
= log
1+
explo, o²)
(E))
This gives Eq.
15


Reviews:
[{"review_id": "pgSyR9HIb6", "reviewer": "Reviewer_UuaD", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Pros,\n\nThis paper is well-written and easy to follow.\n\nThe idea is simple.\n\nThe proposed method can reduce the communication cost\n\nCons,\n\nThe most significant weak point is that the proposed algorithm is not free from the privacy issue. Since every sampled set of each client has to include the classes that the client has, the central server can infer the classes the client has. In many federated learning applications, it is strictly prohibited that the server can reveal any information of clients.\n\nIt would be much better to provide how much the algorithm can save the communication cost. Since the main contribution of this work is not improved performance but reduced communication cost, the authors should provide more discussions about the communication cost.\n\nThe authors should discuss with other communication efficient algorithms since the main contribution of this work is saving the communication cost. For instance, there are many works for efficient distributed learning, e.g., TernGrad. \n\nExperiments should include more results considering other data sets and neural net architectures. Since this is not a theory paper, the proposed algorithm can be justified only by experiments. Therefore, ICLR papers are expected to have comprehensive experimental results. \n", "summary_of_the_paper": "This paper considers the federated learning problem, especially for the case that there are many classes and each client has a small set of classes. When there are many classes, the neural network has to have many parameters to define the last layer classifier, which makes a huge communication cost. To resolve this issue, the author proposes FedSS, by which each FL client samples a set of classes and defines the loss function only with the sampled classes.  Since unsampled classes are not involved in the client's training process, the client sends only the parameters corresponding to the sampled classes. The authors empirically show the proposed approach can reduce the communication cost without sacrificing performance. \n", "summary_of_the_review": "This paper should be very careful to define a class subsampling for federated learning applications since it can leak clients' private information. Also, this paper requires more comprehensive experiments with considering other communication-efficient algorithms.\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "fnD4GR-s73", "reviewer": "Reviewer_24vu", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "===== Strength ======\n1. The paper points out an interesting fact and problem: most of the parameters in a neural net are in the last fully connected (FC) layer, while in a large-number-of-class and non-IID setting, each client may only have data from a small portion of classes. The authors thus accordingly proposed the federated sampled softmax to save the communication cost by local training with only a small number of \"negative\" classes in addition to the \"positive\" classes.\n\n2. The authors conduct experiments on three large-scale datasets. The authors conduct the ablation study on how many negative classes are needed. \n\n==== Weakness =====\n1. The technical contribution is not sufficient. There have been many methods like (Bengio & Sen´ecal,2008) and [a-c] that subsample the negative classes in the softmax objective. Specifically, in [c], the proposed method did include the positive classes in a minibatch, together with a set of subsampled classes, to compute the softmax. While the authors compare different combinations of positive and negative classes in 3.4, to me, it seems pretty obvious that we should combine the clients' positive classes with subsampled negative classes. Besides, while the authors discussed several advanced methods in 3.3, they ended up choosing uniform sampling. I was wondering if any of these advanced approaches can be approximately implemented in a federated setting to further improve the performance.\n\n[a] Joulin et al., Learning Visual Features from Large Weakly Supervised Data, ECCV 2016\n\n[b] Mikolov et al., Distributed representations of words and phrases and their compositionality, NeurIPS 2013\n\n[c] Hu et al., Learning answer embeddings for visual question answering, CVPR 2018\n\n2. One potential direction to improve the novelty or technical contributions of the paper is to develop a series of methods for different kinds of representation learning in a federated setting, including learning a fully connected layer and metric learning. If we do not consider the computational cost at the client end, will metric learning (without learning the FC layer) outperform the proposed method? In [a], the authors proposed to learn multiple one-vs-all classifiers rather than a softmax classifier. Will it be effective in a federated setting?\n\n3. The experimental setup can be improved. There is no clear description of how many classes, images, local epochs for each dataset. I also have concerns about using the pre-trained features for 4.2 and 4.3, which makes 4.2 and 4.3 like downstream tasks rather than representation learning. I would suggest that the authors include a fixed feature baseline in Table 1 and Table 2 to demonstrate that fine-tuning the feature extractor in a federated setting could outperform the pre-trained features. \n\n4. The authors only apply the proposed method to the FedAvg baseline. Will the proposed method be applicable/effective to more advanced federated learning algorithms like FedProx [d], Scaffold [e], and FedDyn [f]?\n\n[d] Li et al., Federated optimization in heterogeneous networks. In MLSys, 2020\n\n[e] Karimireddy et al., Scaffold: Stochastic controlled averaging for federated learning. In ICML, 2020\n\n[f] Acar et al., Federated learning based on dynamic regularization. In ICLR, 2021\n\n==== Other comments ====\n1. I would suggest that the authors add \"supervised\" into their title, to contrast to many recent works on \"unsupervised\" representation learning.\n\n2. I would suggest that the authors cite some more papers on federated learning.", "summary_of_the_paper": "This paper works on \"supervised\" representation learning in a federated learning setting. The main goal is to save the communication cost: by preventing sending the entire fully-connected layer between the server and the clients if the clients only have data from parts of the classes. The authors proposed federated sampled softmax, which is to compute the softmax only over the \"positive\" classes of which a client has data and a small portion of the other negative classes. The authors empirically show that, by doing so, even with a very small set of the negative classes (so small communication cost), the resulting feature network can achieve comparable performance to learning with the conventional softmax. ", "summary_of_the_review": "Overall, I enjoy reading the paper as it pointed out an interesting fact and problem. However, the proposed method seems to be too straightforward with insufficient novelty --- negative class subsampling has been widely used and including positive classes of the client seems to be quite intuitive. I have also listed several potential ways to strengthen the paper (please see the main review). I think the current version of the paper is below the acceptance bar, and thus I give a score of \"3\".", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "Q4_3zIH2SMm", "reviewer": "Reviewer_nvFP", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Pros:\n1. The paper is generally well written and easy to follow.\n2. FedSS provides a potential solution to reduce the communication cost of the classification layer in federated learning settings.\n\nCons:\n1. Firstly, the proposed method only has a significant saving when the number of classes is huge (Figure 8, >>1000). However, the datasets used in this paper mostly have a limited number of classes (SOP and Landmarks). In this case, the saving is very small (smaller than 15%). On real tasks with many classes (ImageNet-21k), the proposed method failed to match the accuracy compared to the full Softmax (Table 3). And the increase of $|S_k|$ does not seem to close the gap. Therefore, it is questionable if the proposed algorithm has real-life benefits.\n2. The proposed method seems to be a direct application of the sampled Softmax algorithm under a federated learning setting, which restricts the novelty.\n3. The proposed method degrades the accuracy (even in Table 1, the results do not really match). How it behaves compared to just using a smaller model with the overall same model size? Since the communication savings on SOP and Landmarks are small, the smaller model baseline should also have roughly the same performance.", "summary_of_the_paper": "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.", "summary_of_the_review": "The paper is generally well written. However, the experimental results and technical novelty are less convincing. I would lean towards objection for now.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "PvIyXw0G8lt", "reviewer": "Reviewer_q79o", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "\n# Significance and impact\n- The problem studied by the authors has not received a lot of attention, save for the previous work cited by the authors, which was restricted to a narrower problem (1 class per client).\n- This paper brings a good practical solution for this problem, building on previous work, which makes it quite significant.\n\n# Writing and clarity\n\n- The paper is very well written and easy to follow, with helpful figures.\n- It could be helpful to provide in appendix a short derivation of the validity of the sampled softmax approximation.\n\n# Quality\n\nThe experiments are well conducted, with adapted baselines, and support well the claims of the paper. I have a few questions and suggestions:\n- in table 2, the proposed method reaches a better performance than fullSoftmax. How do the authors explain this result?\n- The paper mainly focuses on the choice of the classes to include in the sampled softmax, and bring a very good experimental support for the method. It would also have been interesting to study:\n  - different NN architectures (only mobilenet v3 is used)\n  - what is the effect of the proposed method when the number of local epochs increases? This number remains fixed to 1, leading to a very large number of rounds (2k or 5k);\n  - What is the effect of label class heterogeneity on the method? e.g. for the SOP dataset one could study splits with varying heterogeneity.", "summary_of_the_paper": "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.", "summary_of_the_review": "This paper introduces a novel and simple approach to a practical cross-device problem. The experiments are well conducted and support the claims of the paper. I think this paper should be accepted. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper revisits representation learning for extreme settings (large number of class categories) in a federated learning setup. The authors show how each client can 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. The authors investigate the interest of the approach for image classification and image retrieval. 

The reviewers appreciated the interest of the approach to reduce communication and the experimental evaluation on several datasets. The reviewers also expressed concerns about privacy, a central concern in federated learning. One reviewer noted for instance that ‘since every sampled set of each client has to include the classes that the client has, the central server can infer the classes the client has’. The reviewers would also have liked to see a more comprehensive evaluation, in the absence of the theoretical guarantees. Finally, the reviewers expressed regarding accuracy/efficiency trade-offs, one reviewer commenting that “the proposed method degrades the accuracy”. 

The authors submitted responses to the reviewers' comments. The authors discussed the challenges related to privacy. The authors also commented on other gradient sparsification communication-reducing competing approaches (FedAwS) and the choice of datasets. After reading the response, updating the reviews, and discussion, the reviewers found that ‘the current good results are only obtained on smaller-scale datasets with fewer classes [while] in machine learning, the phenomenon could be quite different at different scales’ and that ‘it is not clear if the proposed method can outperform TernGrad at the same amount of transferred data [and] TernGrad also has a better convergence proof compared to the proposed method’. 

We encourage the paper to pursue their approach further taking into account the reviewers' comments, encouragements, and suggestions. Recent progress in privacy protection theoretical frameworks in FL (secure multi party computation, etc.), see the recent survey by Kairouz et al. in FnT in ML, should help the authors develop guarantees for their approach. Moreover the reviewers suggested a clear path towards further improvements of the experimental evaluation. 

The revision of the paper will generate a stronger submission to a future venue.

Reject.
Decision: Reject