Original: Differentiable Top- k with Optimal Transport
Revised: Differentiable Top- k with Optimal Transport

Original: Anonymous Author(s)

Original: AfﬁliationAddress email

Original: Abstract
Revised: Abstract

Original: The top- k operation, i.e., ﬁnding the k largest or smallest elements from a collectionof scores, is an important model component, which is widely used in informationretrieval, machine learning, and data mining.
Revised: The top- k operation, i.e., ﬁnding the k largest or smallest elements from a collection of scores, is an important model component, which is widely used in information retrieval, machine learning, and data mining.

Original: However, if the top- k operation is implemented in an algorithmic way, e.g., using bubble algorithm, the resulting modelcannot be trained in an end-to-end way using prevalent gradient descent algorithms.
Revised: However, if the top- k operation is implemented in an algorithmic way, e.g., using bubble algorithm, the resulting model cannot be trained in an end-to-end way using prevalent gradient descent algorithms.

Original: This is because these implementations typically involve swapping indices, whosegradient cannot be computed.
Revised: This is because these implementations typically involve swapping indices, whose gradient cannot be computed.

Original: Moreover, the corresponding mapping from the inputscores to the indicator vector of whether this element belongs to the top- k set isessentially discontinuous.
Revised: Moreover, the corresponding mapping from the input scores to the indicator vector of whether this element belongs to the top- k set is essentially discontinuous.

Original: To address the issue, we propose a smoothed approximation, namely the SOFT (Scalable Optimal transport-based diFferenTiable) top-koperator.
Revised: To address the issue, we propose a smoothed approximation, namely the SOFT (Scalable Optimal transport-based diFferenTiable) top-

Original:  Speciﬁcally, our SOFT top- k operator approximates the output of thetop- k operation as the solution of an Entropic Optimal Transport (EOT) problem.
Revised: k operator. Speciﬁcally, our SOFT top- k operator approximates the output of the top- k operation as the solution of an Entropic Optimal Transport (EOT) problem.

Original: The gradient of the SOFT operator can then be efﬁciently approximated based onthe optimality conditions of EOT problem.
Revised: The gradient of the SOFT operator can then be efﬁciently approximated based on the optimality conditions of EOT problem.

Original: We apply the proposed operator tothe k -nearest neighbors and beam search algorithms, and demonstrate improvedperformance.
Revised: We apply the proposed operator to the k -nearest neighbors and beam search algorithms, and demonstrate improved performance.

Original:  Introduction
Revised:  Introduction

Original: The top- k operation, i e., ﬁnding the k largest or smallest elements from a set, is widely used forpredictive modeling in information retrieval, machine learning, and data mining.
Revised: The top- k operation, i e., ﬁnding the k largest or smallest elements from a set, is widely used for predictive modeling in information retrieval, machine learning, and data mining.

Original: For example, inimage retrieval (Babenko et al., 2014; Radenovi´c et al., 2016; Gordo et al., 2016), one needs to querythe k nearest neighbors of an input image under certain metrics; in the beam search (Reddy et al. ,1977; Wiseman and Rush, 2016) algorithm for neural machine translation, one needs to ﬁnd the ksequences of largest likelihoods in each decoding step.
Revised: For example, in image retrieval (Babenko et al., 2014; Radenovi´c et al., 2016; Gordo et al., 2016), one needs to query the k nearest neighbors of an input image under certain metrics; in the beam search (Reddy et al., 1977; Wiseman and Rush, 2016) algorithm for neural machine translation, one needs to ﬁnd the k sequences of largest likelihoods in each decoding step.

Original: Although the ubiquity of top- k operation continues to grow, the operation itself is difﬁcult to beintegrated into the training procedure of a predictive model.
Revised: Although the ubiquity of top- k operation continues to grow, the operation itself is difﬁcult to be integrated into the training procedure of a predictive model.

Original: For example, we consider a neuralnetwork-based k -nearest neighbor classiﬁer.
Revised: For example, we consider a neural network-based k -nearest neighbor classiﬁer.

Original: Given an input, we use the neural network to extractfeatures from the input.
Revised: Given an input, we use the neural network to extract features from the input.

Original: Next, the extracted features are fed into the top- k operation for identifying thek nearest neighbors under some distance metric.
Revised: Next, the extracted features are fed into the top- k operation for identifying the k nearest neighbors under some distance metric.

Original: We then obtain a prediction based on the k nearestneighbors of the input.
Revised: We then obtain a prediction based on the k nearest neighbors of the input.

Original: In order to train such a model, we choose a proper loss function, and minimizethe average loss across training samples using (stochastic) ﬁrst-order methods.
Revised: In order to train such a model, we choose a proper loss function, and minimize the average loss across training samples using (stochastic) ﬁrst-order methods.

Original: This naturally requiresthe loss function being differentiable with respect to the input at each update step.
Revised: This naturally requires the loss function being differentiable with respect to the input at each update step.

Original: Nonetheless, thetop- k operation does not exhibit an explicit mathematical formulation: most implementations of thetop- k operation, e.g., bubble algorithm and Q UICKSELECT (Hoare, 1961), involve operations on indicessuch as indices swapping.
Revised: Nonetheless, the top- k operation does not exhibit an explicit mathematical formulation: most implementations of the top- k operation, e.g., bubble algorithm and Q UICKSELECT (Hoare, 1961), involve operations on indices such as indices swapping.

Original: Consequently, the training objective is difﬁcult to formulate explicitly.
Revised: Consequently, the training objective is difﬁcult to formulate explicitly.

Revised: [Figure]

Original: Alternative perspective — taking the top- k operation as an operator — still cannot resolve thedifferentibility issue.
Revised: Alternative perspective — taking the top- k operation as an operator — still cannot resolve the differentibility issue.

Original: Speciﬁcally, the top- k operator 1 maps a set of inputs x 1 , . . .
Revised: Speciﬁcally, the top- k operatormaps a set of inputs x 1 , . . .

Original: , x n to an indexvector { 0 , 1 } n .
Revised: , x n to an index vector { 0 , 1 } n .

Original: Whereas the Jacobian matrix of such a mapping is not well deﬁned.
Revised: Whereas the Jacobian matrix of such a mapping is not well deﬁned.

Original: As a simple example, consider two scalars x 1 , x 2 .
Revised: As a simple example, consider two scalars x 1 , x 2 .

Original: The top- 1 operation as in Figure 1 returns a vector [ A 1 , A 2 ] (cid:62) ,with each entry denoting whether the scalar is the larger one ( 1 for true, 0 for false).
Revised: The top- 1 operation as in Figure 1 returns a vector [ A 1 , A 2 ] (cid:62) , with each entry denoting whether the scalar is the larger one ( 1 for true, 0 for false).

Original: Denote A 1 = f ( x 1 , x 2 ) .
Revised: Denote A 1 = f ( x 1 , x 2 ) .

Original: For a ﬁxed x 2 , A 1 jumps from 0 to 1 at x 1 = x 2 .
Revised: For a ﬁxed x 2 , A 1 jumps from 0 to 1 at x 1 = x 2 .

Original: It is clear that f is notdifferentiable at x 1 = x 2 , and the derivative is identically zero otherwise.[Figure]
Revised: It is clear that f is not differentiable at x 1 = x 2 , and the derivative is identically zero otherwise.

Original: Due to the aforementioned difﬁculty, existing worksresort to two-stage training for models with the topk operation.
Revised: Due to the aforementioned difﬁculty, existing works resort to two-stage training for models with the top- k operation.

Original: We consider the neural network-basedk -nearest neighbor classiﬁer again.
Revised: We consider the neural network-based k -nearest neighbor classiﬁer again.

Original: As proposedin Papernot and McDaniel (2018), one ﬁrst trainsthe neural network using some surrogate loss on theextracted features, e.g., using softmax activation inthe output layer and the cross-entropy loss.
Revised: As proposed in Papernot and McDaniel (2018), one ﬁrst trains the neural network using some surrogate loss on the extracted features, e.g., using softmax activation in the output layer and the cross-entropy loss.

Original: Next,one uses the k -nearest neighbor for prediction basedon the features extracted by the well-trained neural network.
Revised: Next, one uses the k -nearest neighbor for prediction based on the features extracted by the well-trained neural network.

Original: This training procedure, althoughcircumventing the top- k operation, makes the training and prediction misaligned; and the actualperformance suffers.
Revised: This training procedure, although circumventing the top- k operation, makes the training and prediction misaligned; and the actual performance suffers.

Original: In this work, we propose the SOFT (Scalable Optimal transport-based diFferenTiable) top-
Revised: In this work, we propose the SOFT (Scalable Optimal transport-based diFferenTiable) top-

Original: k operationas a differentiable approximation of the standard top- k operation in Section.
Revised: k operation as a differentiable approximation of the standard top- k operation in Section.

Original: Speciﬁcally, motivatedby the implicit differentiation (Duchi et al., 2008; Griewank and Walther, 2008; Amos and Kolter,2017; Luise et al., 2018) techniques, we ﬁrst parameterize the top- k operation in terms of the optimalsolution of an Optimal Transport (OT) problem.
Revised: Speciﬁcally, motivated by the implicit differentiation (Duchi et al., 2008; Griewank and Walther, 2008; Amos and Kolter, 2017; Luise et al., 2018) techniques, we ﬁrst parameterize the top- k operation in terms of the optimal solution of an Optimal Transport (OT) problem.

Original: Such a re-parameterization is still not differentiablewith respect to the input.
Revised: Such a re-parameterization is still not differentiable with respect to the input.

Original: To rule out the discontinuity, we impose entropy regularization to theoptimal transport problem, and show that the optimal solution to the Entropic OT (EOT) problemyields a differentiable approximation to the top- k operation.
Revised: To rule out the discontinuity, we impose entropy regularization to the optimal transport problem, and show that the optimal solution to the Entropic OT (EOT) problem yields a differentiable approximation to the top- k operation.

Original: Moreover, we prove that under mildassumptions, the approximation error can be properly controlled.
Revised: Moreover, we prove that under mild assumptions, the approximation error can be properly controlled.

Original: We then develop an efﬁcient implementation of the SOFT top- k operation in Section.
Revised: We then develop an efﬁcient implementation of the SOFT top- k operation in Section.

Original: Speciﬁcally,we solve the EOT problem via the Sinkhorn algorithm (Cuturi, 2013).
Revised: Speciﬁcally, we solve the EOT problem via the Sinkhorn algorithm (Cuturi, 2013).

Original: Given the optimal solution, wecan explicitly formulate the gradient of SOFT top- k operation using the KKT (Karush-Kuhn-Tucker)condition.
Revised: Given the optimal solution, we can explicitly formulate the gradient of SOFT top- k operation using the KKT (Karush-Kuhn-Tucker) condition.

Original: As a result, the gradient at each update step can be efﬁciently computed with complexity
Revised: As a result, the gradient at each update step can be efﬁciently computed with complexity

Original: O ( n ) , where n is the number of elements in the input set to the top- k operation.
Revised: O ( n ) , where n is the number of elements in the input set to the top- k operation.

Original: Our proposed SOFT top- k operation allows end-to-end training, and we apply SOFT top- k operationto k NN for classiﬁcation in Section C and beam search in Section D.
Revised: Our proposed SOFT top- k operation allows end-to-end training, and we apply SOFT top- k operation to k NN for classiﬁcation in Section C and beam search in Section D.

Original: The experimental results demonstrate signiﬁcant performance gain over competing methods, as an end-to-end training procedureresolves the misalignment between training and prediction.
Revised: The experimental results demonstrate signiﬁcant performance gain over competing methods, as an end-to-end training procedure resolves the misalignment between training and prediction.

Original:  SOFT Top- k Operator
Revised:  SOFT Top- k Operator

Original: We adopt the following deﬁnition of the (augment of) top- k operator.
Revised: We adopt the following deﬁnition of the (augment of) top- k operator.

Original: Given a set of scalars X
Revised: Given a set of scalars X

Original: =
Revised: =

Original: { x i } ni
Revised: { x i } ni

Original: =1 , the standard top- k operator returns a vector A =
Revised: =1 , the standard top- k operator returns a vector A =

Original: [ A 1 , . .
Revised: [ A 1 , . .

Original: , A n ] (cid:62) , such that
Revised: , A n ] (cid:62) , such that

Original: A i = (cid:26) 1 , if x i is a top- k element in X , 0 , otherwise .
Revised: [Equation] , if x i is a top- k element in X , 0 , otherwise .

Original: Note that the deﬁnition is essentially an "arg-top- k " operation since it marks the top- k indices as 1 ,instead of returning the top- k values .
Revised: Note that the deﬁnition is essentially an "arg-top- k " operation since it marks the top- k indices as 1 , instead of returning the top- k values .

Original: This allows more ﬂexibility since we can obtain the top- k valuesby multiplying A to X .
Revised: This allows more ﬂexibility since we can obtain the top- k values by multiplying A to X .

Original: The goal is to design a smooth relaxation of the standard top- k operator.
Revised: The goal is to design a smooth relaxation of the standard top- k operator.

Original: Without loss of generality, we refer to top- k elements as the smallest k elements.
Revised: Without loss of generality, we refer to top- k elements as the smallest k elements.

Original: 2.1 Parameterizing Top- k Operator as OT Problem
Revised: 2.1 Parameterizing Top- k Operator as OT Problem

Original: We ﬁrst show that the standard top- k operator can be parameterized in terms of the solution ofan Optimal Transport (OT) problem (Monge, 1781; Kantorovich, 1960).
Revised: We ﬁrst show that the standard top- k operator can be parameterized in terms of the solution of an Optimal Transport (OT) problem (Monge, 1781; Kantorovich, 1960).

Original: We brieﬂy introduce OTproblems for self-containedness.
Revised: We brieﬂy introduce OT problems for self-containedness.

Original: An OT problem ﬁnds a transport plan between two distributions,while the expected cost of the transportation is minimized.
Revised: An OT problem ﬁnds a transport plan between two distributions, while the expected cost of the transportation is minimized.

Original: We consider two discrete distributionsdeﬁned on supports A = { a i } ni =1 and B = { b j } mj =1 , respectively.
Revised: We consider two discrete distributions deﬁned on supports A = { a i } ni =1 and B = { b j } mj =1 , respectively.

Original: Denote P ( { a i } )
Revised: Denote P ( { a i } )

Original: = µ i and
Revised: = µ i and

Original: P ( { b j } ) =
Revised: P ( { b j } ) =

Original: ν j , and let µ = [ µ 1 , . .
Revised: ν j , and let µ = [ µ 1 , . .

Original: . , µ n ] (cid:62) and ν = [ ν 1 , . . .
Revised: . , µ n ] (cid:62) and ν = [ ν 1 , . . .

Original: , ν m ] (cid:62) .
Revised: , ν m ] (cid:62) .

Original: We further denote C ∈ R n × m asthe cost matrix with C ij being the cost of transporting mass from a i to b j .
Revised: We further denote C ∈ R n × m as the cost matrix with C ij being the cost of transporting mass from a i to b j .

Original: An OT problem can beformulated as
Revised: An OT problem can be formulated as

Original: [Equation]
Revised: [Equation]

Original: where 1 denotes a vector of ones.
Revised: where 1 denotes a vector of ones.

Original: The optimal Γ ∗ is referred to as the optimal transport plan .
Revised: The optimal Γ ∗ is referred to as the optimal transport plan .

Original: [Figure]

Original: In order to parameterize the top- k operator using the optimal transport plan Γ ∗ , we set the support
Revised: In order to parameterize the top- k operator using the optimal transport plan Γ ∗ , we set the support

Original: A = X and B = { 0 , 1 } in (1), with µ, ν deﬁned as
Revised: A = X and B = { 0 , 1 } in (1), with µ, ν deﬁned as

Original: [Equation]
Revised: [Equation]

Original: We take the cost to be the squared Euclidean distance, i.e., C i 1 = x 2 i and C i 2 =
Revised: We take the cost to be the squared Euclidean distance, i.e., C i 1 = x 2 i and C i 2 =

Original: ( x i − 1) 2 fori = 1 , . . . , n .
Revised: ( x i − 1) 2 for i = 1 , . . . , n .

Original: We then establish the relationship between the output A of the top- k operator and Γ ∗ .
Revised: We then establish the relationship between the output A of the top- k operator and Γ ∗ .

Original: Proposition 1.
Revised: Proposition 1.

Original: Consider the setup in the previous paragraph.
Revised: Consider the setup in the previous paragraph.

Original: Without loss of generality, we assume
Revised: Without loss of generality, we assume

Original: X has no duplicates.
Revised: X has no duplicates.

Original: Then the optimal transport plan Γ ∗ of (1) isΓ ∗ σ i , 1 = (cid:26) 1 /n, if i ≤ k, 0 , if k + 1 ≤ i ≤ n. , Γ ∗ σ i , 2 = (cid:26) 0 , if i ≤ k, 1 /n, if k + 1 ≤ i ≤ n, (2) with σ being the sorting permutation, i.e., x σ 1 < x σ 2 < · · · < x σ n .
Revised: Then the optimal transport plan Γ ∗ of (1) is [Equation] with σ being the sorting permutation, i.e., x σ 1 < x σ 2 < · · · < x σ n .

Original: Moreover, we have
Revised: Moreover, we have

Original: A = n Γ ∗ · [1 , 0] (cid:62) .
Revised: [Equation]

Original: The proof can be found in Appendix G. Figure 3(a) illustrates the corresponding optimal transportplan for parameterizing the top- 5 operator applied to a set of 7 elements.
Revised: The proof can be found in Appendix G. Figure 2(a) illustrates the corresponding optimal transport plan for parameterizing the top- 5 operator applied to a set of 7 elements.

Original: As can be seen, the massfrom the 5 closest points is transported to 0 , and meanwhile the mass from the 2 remaining points istransported to 1 .
Revised: As can be seen, the mass from the 5 closest points is transported to 0 , and meanwhile the mass from the 2 remaining points is transported to 1 .

Original: Therefore, the optimal transport plan exactly indicates the top- 5 elements.
Revised: Therefore, the optimal transport plan exactly indicates the top- 5 elements.

Original: 2.2 Smoothing by Entropy Regularization
Revised: 2.2 Smoothing by Entropy Regularization

Original: We next rule out the discontinuity of (1) to obtain a smoothed approximation to the top- k operator.
Revised: We next rule out the discontinuity of (1) to obtain a smoothed approximation to the top- k operator.

Original: Speciﬁcally, we employ entropy regularization to the OT problem (1):
Revised: Speciﬁcally, we employ entropy regularization to the OT problem (1):

Original: [Equation]
Revised: [Equation]

Original: where h (Γ) =
Revised: where h (Γ) =

Original: (cid:80) i,j Γ ij log Γ ij is the entropy regularizer.
Revised: (cid:80) i,j Γ ij log Γ ij is the entropy regularizer.

Original: We deﬁne A (cid:15) =
Revised: We deﬁne A (cid:15) =

Original: n Γ ∗ ,(cid:15) · [0 , 1] (cid:62) as asmoothed counterpart of output A in the standard top- k operator.
Revised: n Γ ∗ ,(cid:15) · [0 , 1] (cid:62) as a smoothed counterpart of output A in the standard top- k operator.

Original: Accordingly, SOFT top-
Revised: Accordingly, SOFT top-

Original: k operatoris deﬁned as the mapping from X to A (cid:15) .
Revised: k operator is deﬁned as the mapping from X to A (cid:15) .

Original: We show that the Jacobian matrix of SOFT top- k operatorexists and is nonzero in the following theorem.
Revised: We show that the Jacobian matrix of SOFT top- k operator exists and is nonzero in the following theorem.

Original: Theorem 1.
Revised: Theorem 1.

Original: For any (cid:15) > 0 , SOFT top- k operator: X (cid:55)→ A (cid:15) is differentiable, as long as the cost C ijis differentiable with respect to x i for any i, j .
Revised: For any (cid:15) > 0 , SOFT top- k operator: X (cid:55)→ A (cid:15) is differentiable, as long as the cost C ij is differentiable with respect to x i for any i, j .

Original: Moreover, the Jacobian matrix of SOFT top- k operatoralways has a nonzero entry for any X ∈ R n .
Revised: Moreover, the Jacobian matrix of SOFT top- k operator always has a nonzero entry for any X ∈ R n .

Original: The proof can be found in Appendix G. We remark that the entropic OT (4) is computationally morefriendly, since it allows the usage of ﬁrst-order algorithms (Cuturi, 2013).
Revised: The proof can be found in Appendix G. We remark that the entropic OT (4) is computationally more friendly, since it allows the usage of ﬁrst-order algorithms (Cuturi, 2013).

Original: The Entropic OT introduces bias to the SOFT top- k operator.
Revised: The Entropic OT introduces bias to the SOFT top- k operator.

Original: The following theorem shows that sucha bias can be effectively controlled.
Revised: The following theorem shows that such a bias can be effectively controlled.

Original: Theorem 2.
Revised: Theorem 2.

Original: Given a distinct sequence X and its sorting permutation σ , with Euclidean square costfunction, for the proposed top- k solver we have
Revised: Given a distinct sequence X and its sorting permutation σ , with Euclidean square cost function, for the proposed top- k solver we have

Original: [Equation]
Revised: [Equation]

Original: Therefore, with a small enough (cid:15) , the output vector A (cid:15) can well approximate A , especially whenthere is a large gap between x σ k and x σ k +1 .
Revised: Therefore, with a small enough (cid:15) , the output vector A (cid:15) can well approximate A , especially when there is a large gap between x σ k and x σ k +1 .

Original: Besides, Theorem 2 suggests a trade-off between thebias and regularization of SOFT top- k operator.
Revised: Besides, Theorem 2 suggests a trade-off between the bias and regularization of SOFT top- k operator.

Original: See Section F for a detailed discussion.
Revised: See Section F for a detailed discussion.

Original:  Efﬁcient Implementation
Revised:  Efﬁcient Implementation

Original: We now present our implementation of SOFT top- k operator, which consists of 1) computing A (cid:15)from X and 2) computing the Jacobian matrix of A (cid:15) with respect to X .
Revised: We now present our implementation of SOFT top- k operator, which consists of 1) computing A (cid:15) from X and 2) computing the Jacobian matrix of A (cid:15) with respect to X .

Original: We refer to 1) as the forwardpass and 2) as the backward pass.
Revised: We refer to 1) as the forward pass and 2) as the backward pass.

Revised: [Equation]

Original: Forward Pass The forward pass from X to A (cid:15) can be efﬁciently computed using Sinkhorn algorithm.
Revised: Forward Pass The forward pass from X to A (cid:15) can be efﬁciently computed using Sinkhorn algorithm.

Original: Speciﬁcally, we run iterative Bregman projections (Benamou et al., 2015), where at the (cid:96) -th iteration,we update
Revised: Speciﬁcally, we run iterative Bregman projections (Benamou et al., 2015), where at the (cid:96) -th iteration, we update

Original: [Equation]
Revised: [Equation]

Original: Here, the division is entrywise, q (0) = 1 2 / 2 , and G ∈ R n × m with G ij = e − Cij(cid:15) .
Revised: Here, the division is entrywise, q (0) = 1 2 / 2 , and G ∈ R n × m with G ij = e − Cij(cid:15) .

Original: Denote p ∗ and q ∗as the stationary point of the Bregman projections.
Revised: Denote p ∗ and q ∗ as the stationary point of the Bregman projections.

Original: The optimal transport plan Γ ∗ ,(cid:15) can be obtainedby Γ ∗ ,(cid:15)ij = p ∗ i G ij q ∗ j .
Revised: The optimal transport plan Γ ∗ ,(cid:15) can be obtained by Γ ∗ ,(cid:15)ij = p ∗ i G ij q ∗ j .

Original: The algorithm is summarized in Algorithm 1.
Revised: The algorithm is summarized in Algorithm 1.

Original: Backward Pass.
Revised: Backward Pass.

Original: Given A (cid:15) , we compute the Jacobian matrix dA (cid:15)
Revised: Given A (cid:15) , we compute the Jacobian matrix dA (cid:15)

Original: d X using implicit differentiation anddifferentiable programming techinques.
Revised: d X using implicit differentiation and differentiable programming techinques.

Original: Speciﬁcally, the Lagrangian function of Problem (4) is
Revised: Speciﬁcally, the Lagrangian function of Problem (4) is

Original: [Equation]
Revised: [Equation]

Original: where ξ and ζ are dual variables.
Revised: where ξ and ζ are dual variables.

Original: The KKT condition implies that Γ ∗ ,(cid:15) can be formulated using theoptimal dual variables ξ ∗ and ζ ∗ as (Sinkhorn’s scaling theorem, Sinkhorn and Knopp (1967)),
Revised: The KKT condition implies that Γ ∗ ,(cid:15) can be formulated using the optimal dual variables ξ ∗ and ζ ∗ as (Sinkhorn’s scaling theorem, Sinkhorn and Knopp (1967)),

Original: [Equation]
Revised: [Equation]

Original: Substituting (5) into the Lagrangian function, we obtain
Revised: Substituting (5) into the Lagrangian function, we obtain

Original: [Equation]
Revised: [Equation]

Original: − (cid:15) n,m (cid:88) i,j = e − Cij − ξ ∗ i − ζ ∗ j (cid:15) .
Revised: i,j = e − C ij − ξ ∗ i − ζ ∗ j (cid:15) .

Original: We now compute the gradient of ξ ∗ and ζ ∗ with respect to C , such that we can obtain d Γ ∗ ,(cid:15) /dCby the chain rule applied to (5).
Revised: We now compute the gradient of ξ ∗ and ζ ∗ with respect to C , such that we can obtain d Γ ∗ ,(cid:15) /dC by the chain rule applied to (5).

Original: Denote ω ∗ =
Revised: Denote ω ∗ =

Original: [( ξ ∗ ) (cid:62) , ( ζ ∗ ) (cid:62) ] (cid:62) , and φ ( ω ∗ ; C ) =
Revised: [( ξ ∗ ) (cid:62) , ( ζ ∗ ) (cid:62) ] (cid:62) , and φ ( ω ∗ ; C ) =

Original: ∂
Revised: ∂

Original: L ( ω ∗ ; C ) /∂ω ∗ .
Revised: L ( ω ∗ ; C ) /∂ω ∗ .

Revised: Algorithm

Original: Require:
Revised: Require:

Original: X =
Revised: X =

Original: [ x i ] ni =1 , k, (cid:15), L
Revised: [ x i ] ni =1 , k, (cid:15), L

Original: Y =
Revised: Y =

Original: [ y 1 , y 2 ] (cid:62) =
Revised: [ y 1 , y 2 ] (cid:62) =

Original: [0 , 1] (cid:62) µ = 1 n /n, ν =
Revised: [0 , 1] (cid:62) µ = 1 n /n, ν =

Original: [ k/n, ( n − K )
Revised: [ k/n, ( n − K )

Original: /n ] (cid:62)
Revised: /n ] (cid:62)

Original: C ij =
Revised: C ij =

Original: | x i − y j | 2 , G ij = e − Cij(cid:15) , q = 1 2 /for l = 1 , · · · , L do p = µ/
Revised: | x i − y j | 2 , G ij = e − Cij(cid:15) , q = 1 2 /for l = 1 , · · · , L do p = µ/

Original: ( Gq ) , q = ν/ ( G (cid:62) p ) end for Γ = diag( p ) (cid:12) G (cid:12) diag( q )
Revised: ( Gq ) , q = ν/ ( G (cid:62) p ) end for Γ = diag( p ) (cid:12) G (cid:12) diag( q )

Original: A (cid:15) =
Revised: A (cid:15) =

Original: n Γ · [0 , 1] (cid:62)
Revised: n Γ · [0 , 1] (cid:62)

Original: At the optimal dual variable ω ∗ , the KKT conditionimmediately yieldsφ ( ω ∗ ; C ) ≡ 0 .
Revised: At the optimal dual variable ω ∗ , the KKT condition immediately yields φ ( ω ∗ ; C ) ≡ 0 .

Original: By the chain rule, we havedφ
Revised: By the chain rule, we have dφ ( ω ∗ ; C ) dC = ∂φ (

Original: [Equation]
Revised: [Equation]

Original: Rearranging terms, we obtaindω
Revised: ∗ ; C ) ∂C + ∂φ ( ω ∗ ; C ) ∂ω

Original: [Equation]
Revised: [Equation]

Original:  (cid:19) [Equation] ; C ) ∂C .
Revised: dω ∗ dC = 0 . Rearranging terms, we obtain dω ∗ dC = − (cid:18) ∂φ ( ω ∗ ; C ) ∂ω ∗ (cid:19) − 1 ∂φ ( ω ∗ ; C ) ∂C .

Original: Combining (5), (6), C ij = ( x i − y j ) 2 , and A (cid:15) =
Revised: Combining (5), (6), C ij = ( x i − y j ) 2 , and A (cid:15) =

Original: n Γ ∗ ,(cid:15) · [1 , 0] (cid:62) , the Jacobian matrix dA (cid:15) /d
Revised: n Γ ∗ ,(cid:15) · [1 , 0] (cid:62) , the Jacobian matrix dA (cid:15) /d

Original: X canthen be derived using the chain rule again.
Revised: X can then be derived using the chain rule again.

Original: The detailed derivation and the corresponding algorithm for computing the Jacobian matrix canbe found in Appendix H. The time and space complexity of the derived algorithm is O ( n ) and
Revised: The detailed derivation and the corresponding algorithm for computing the Jacobian matrix can be found in Appendix H. The time and space complexity of the derived algorithm is O ( n ) and

Original: O ( kn ) for top- k and sorted top- k operators, respectively.
Revised: O ( kn ) for top- k and sorted top- k operators, respectively.

Original: We also include a Pytorch Paszke et al.(2017) implementation of the forward and backward pass in Appendix H by extending the autogradautomatic differentiation package.
Revised: We also include a Pytorch Paszke et al. (2017) implementation of the forward and backward pass in Appendix H by extending the autograd automatic differentiation package.

Original:  Experiments
Revised:  Experiments

Original: Our proposed SOFT top- k operation allows end-to-end training, and we apply SOFT top- k operationto deep k NN for  classiﬁcation and beam search for the machine translation task, and show the resultsin Table 1 and Table 2, respectively.
Revised: We apply SOFT top- k operation to deep k NN for image classiﬁcation and beam search for the machine translation task, and show the results in Table 1 and Table 2, respectively.

Original: The experimental results demonstrate signiﬁcant performancegain over competing methods.
Revised: The experimental results demonstrate signiﬁcant performance gain over competing methods.

Original: [Table]
Revised: [Table]

Original: [Table]

Original: [Table]
Revised: [Table]

Original: A Notations
Revised: A Notations

Original: We denote (cid:107) · (cid:107) 2 as the (cid:96) 2 norm of vectors, (cid:107) · (cid:107) F as the Frobenius norm of matrices.
Revised: We denote (cid:107) · (cid:107) 2 as the (cid:96) 2 norm of vectors, (cid:107) · (cid:107) F as the Frobenius norm of matrices.

Original: Given twomatrices B, D ∈ R n × m , we denote (cid:104) B, D (cid:105) as the inner product, i.e., (cid:104) B, D (cid:105) =
Revised: Given two matrices B, D ∈ R n × m , we denote (cid:104) B, D (cid:105) as the inner product, i e., (cid:104) B, D (cid:105) =

Original: (cid:80) n,mi =1 ,j =1 B ij D ij .
Revised: (cid:80) n,mi =1 ,j =1 B ij D ij .

Original: We denote B (cid:12) D as the element-wise multiplication of B and D .
Revised: We denote B (cid:12) D as the element-wise multiplication of B and D .

Original: We denote 1 ( · ) as the indicatorfunction, i.e., the output of 1 ( · ) is 1 if the input condition is satisﬁed, and is 0 otherwise.
Revised: We denote 1 ( · ) as the indicator function, i e., the output of 1 ( · ) is 1 if the input condition is satisﬁed, and is 0 otherwise.

Original: For matrix
Revised: For matrix

Original: B ∈ R n × m , we denote B i, : as the i -th row of the matrix.
Revised: B ∈ R n × m , we denote B i, : as the i -th row of the matrix.

Original: The softmax function for matrix B isdeﬁned as softmax i ( B ij )
Revised: The softmax function for matrix B is deﬁned as softmax i ( B ij )

Original: =
Revised: =

Original: e B ij / (cid:80) n(cid:96)
Revised: e B ij / (cid:80) n(cid:96)

Original: =1 e B lj .
Revised: =1 e B lj .

Original: For a vector b ∈ R n , we denote diag( b ) as the matrix where the i -th diagonal entries is b i .
Revised: For a vector b ∈ R n , we denote diag( b ) as the matrix where the i -th diagonal entries is b i .

Original: B Sorted SOFT Top- k Operator
Revised: B Sorted SOFT Top- k Operator

Original: In some applications, we not only need to distinguish the top- k elements, but also sort the top- kelements.
Revised: In some applications, we not only need to distinguish the top- k elements, but also sort the top- k elements.

Original: For example, in image retrieval (Gordo et al., 2016), the retrieved k images are expected tobe sorted.
Revised: For example, in image retrieval (Gordo et al., 2016), the retrieved k images are expected to be sorted.

Original: Our SOFT top- k operator can be extended to the sorted SOFT top- k operator.
Revised: Our SOFT top- k operator can be extended to the sorted SOFT top- k operator.

Original: Analogous to the derivation of the SOFT top- k operator, we ﬁrst parameterize the sorted top- koperator in terms of an OT problem.
Revised: Analogous to the derivation of the SOFT top- k operator, we ﬁrst parameterize the sorted top- k operator in terms of an OT problem.

Original: Speciﬁcally, we keep A = X and µ = 1 n /n and set
Revised: Speciﬁcally, we keep A = X and µ = 1 n /n and set

Original: B = [0 , 1 , 2 , · · · , k ] (cid:62) , and ν = [1 /n, · · · , 1 /n, ( n − k ) /n ] (cid:62) .
Revised: B = [0 , 1 , 2 , · · · , k ] (cid:62) , and ν = [1 /n, · · · , 1 /n, ( n − k ) /n ] (cid:62) .

Original: One can check that the optimal transport plan of the above OT problem transports the smallest elementin A to 0 in B , the second smallest element to 1 , and so on so forth.
Revised: One can check that the optimal transport plan of the above OT problem transports the smallest element in A to 0 in B , the second smallest element to 1 , and so on so forth.

Original: This in turn yields the sortedtop- k elements.
Revised: This in turn yields the sorted top- k elements.

Original: Figure 3(b) illustrates the sorted top- 2 operator and its optimal transport plan.
Revised: Figure 2(b) illustrates the sorted top- 2 operator and its optimal transport plan.

Original: The sorted SOFT top- k operator is obtained similarly to SOFT top- k operator by solving the entropyregularized OT problem.
Revised: The sorted SOFT top- k operator is obtained similarly to SOFT top- k operator by solving the entropy regularized OT problem.

Original: We can show that the sorted SOFT top- k operator is differentiable and thebias can be properly controlled.
Revised: We can show that the sorted SOFT top- k operator is differentiable and the bias can be properly controlled.

Original: C k -NN for Image Classiﬁcation
Revised: C k -NN for Image Classiﬁcation

Original: The proposed SOFT top- k operator enables us to train an end-to-end neural network-based k NNclassiﬁer.
Revised: The proposed SOFT top- k operator enables us to train an end-to-end neural network-based k NN classiﬁer.

Original: Speciﬁcally, we receive training samples { Z i , y i } Ni =1 with Z i being the input data andy i ∈ { 1 , . . .
Revised: Speciﬁcally, we receive training samples { Z i , y i } Ni =1 with Z i being the input data and y i ∈ { 1 , . . .

Original: , M } the label from M classes.
Revised: , M } the label from M classes.

Original: During the training, for an input data Z j (also known asthe query sample), we associate a loss as follows.
Revised: During the training, for an input data Z j (also known as the query sample), we associate a loss as follows.

Original: Denote Z \ j as all the input data excluding Z j (alsoknown as the template samples).
Revised: Denote Z \ j as all the input data excluding Z j (also known as the template samples).

Original: We use a neural network f θ parameterized by θ to extract featuresfrom all the input data, and measure the pairwise Euclidean distances between the extracted featuresof Z \ j and that of Z j .
Revised: We use a neural network f θ parameterized by θ to extract features from all the input data, and measure the pairwise Euclidean distances between the extracted features of Z \ j and that of Z j .

Original: Denote X \ j,θ as the collection of these pairwise distances, i.e., 
Revised: Denote X \ j,θ as the collection of these pairwise distances, i.e., X

Original: [Equation]
Revised: [Equation]

Original:  where the subscript of X emphasizes its dependence on θ .
Revised: , where the subscript of X emphasizes its dependence on θ .

Original: Next, we apply SOFT top- k operator to X \ j,ω , and the returned vector is denoted by A (cid:15) \ j,θ .
Revised: Next, we apply SOFT top- k operator to X \ j,ω , and the returned vector is denoted by A (cid:15) \ j,θ .

Original: Let
Revised: Let

Original: Y \ j ∈ R M × ( N − 1) be the matrix by concatenating the one-hot encoding of labels
Revised: Y \ j ∈ R M × ( N − 1) be the matrix by concatenating the one-hot encoding of labels

Original: y i for i (cid:54) =
Revised: y i for i (cid:54) =

Original: j ascolumns, and Y j ∈ R M the one-hot encoding of the label y j .
Revised: j as columns, and Y j ∈ R M the one-hot encoding of the label y j .

Original: The loss of Z j is deﬁned as
Revised: The loss of Z j is deﬁned as

Original: [Equation]
Revised: [Equation]

Original: Consequently, the training loss is L ( { Z j , y j } Nj =1 ) = 1 N (cid:80) Nj =1 (cid:96) ( Z j , y j ) .
Revised: Consequently, the training loss is L ( { Z j , y j } Nj =1 ) = 1 N (cid:80) Nj =1 (cid:96) ( Z j , y j ) .

Original: Recall that the Jacobianmatrix of A (cid:15) \ j,θ exists and has no zero entries.
Revised: Recall that the Jacobian matrix of A (cid:15) \ j,θ exists and has no zero entries.

Original: This allows us to utilize stochastic gradient descent algorithms to update θ in the neural network.
Revised: This allows us to utilize stochastic gradient descent algorithms to update θ in the neural network.

Original: Moreover, since N is often large, to ease the computation,we randomly sample a batch of samples to compute the stochastic gradient at each iteration.
Revised: Moreover, since N is often large, to ease the computation, we randomly sample a batch of samples to compute the stochastic gradient at each iteration.

Original: In the prediction stage, we use all the trainingsamples to obtain a predicted label of a querysample.
Revised: In the prediction stage, we use all the training samples to obtain a predicted label of a query sample.

Original: Speciﬁcally, we feed the query sampleinto the neural network to extract its features,and compute pairwise Euclidean distances to allthe training samples.
Revised: Speciﬁcally, we feed the query sample into the neural network to extract its features, and compute pairwise Euclidean distances to all the training samples.

Original: We then run the standardk NN algorithm (Hastie et al., 2009) to obtainthe predicted label.
Revised: We then run the standard k NN algorithm (Hastie et al., 2009) to obtain the predicted label.

Original: C.1 Experiment
Revised: C.1 Experiment

Original: We evaluate the performance of the proposedneural network-based k NN classiﬁer on twobenchmark datasets: MNIST dataset of handwritten digits (LeCun et al., 1998) and the CIFAR10 dataset of natural images (Krizhevsky et al., 2009) with the canonical splits for training andtesting without data augmentation.
Revised: We evaluate the performance of the proposed neural network-based k NN classiﬁer on two benchmark datasets: MNIST dataset of handwritten digits (LeCun et al., 1998) and the CIFAR10 dataset of natural images (Krizhevsky et al., 2009) with the canonical splits for training and testing without data augmentation.

Original: We adopt the coefﬁcient of entropy regularizer (cid:15) =
Revised: We adopt the coefﬁcient of entropy regularizer (cid:15) =

Original: 10 − 3 for
Revised: 10 − 3 for

Original: MNIST dataset and (cid:15) =
Revised: MNIST dataset and (cid:15) =

Original: 10 − 5 for CIFAR-10 dataset.
Revised: 10 − 5 for CIFAR-10 dataset.

Original: Further implementation details can be found in
Revised: Further implementation details can be found in

Original: Appendix I.
Revised: Appendix I.

Original: Baselines.
Revised: Baselines.

Original: We consider several baselines:
Revised: We consider several baselines:

Original: Standard k NN method.2.
Revised: Standard k NN method.

Original: Two-stage training methods: we ﬁrst extract the features of the images, and then perform k NN onthe features.
Revised: Two-stage training methods: we ﬁrst extract the features of the images, and then perform k NN on the features.

Original: The feature is extracted using Principle Component Analysis (PCA, top- 50 principlecomponents is adopted), autoencoder (AE), or a pretrained Convolutional Neural Network (CNN)using the Cross-Entropy (CE) loss.3.
Revised: The feature is extracted using Principle Component Analysis (PCA, top- 50 principle components is adopted), autoencoder (AE), or a pretrained Convolutional Neural Network (CNN) using the Cross-Entropy (CE) loss.

Original: Differentiable ranking + k NN: This includes NeuralSort (Grover et al., 2019) and Cuturi et al.(2019).
Revised: Differentiable ranking + k NN: This includes NeuralSort (Grover et al., 2019) and Cuturi et al.

Original: Cuturi et al.
Revised: Cuturi et al.

Original: (2019) is not directly applicable, which requires adaptations (see Appendix I).
Revised: (2019) is not directly applicable, which requires adaptations (see Appendix I).

Original: Stochastic k NN with Gumbel top-
Revised: Stochastic k NN with Gumbel top-

Original: k relaxation (Xie and Ermon, 2019):
Revised: k relaxation (Xie and Ermon, 2019):

Original: The model is referred asRelaxSubSample.5.
Revised: The model is referred as RelaxSubSample.

Original: Softmax Augmentation for smoothed top- k operation: A combination of k softmax operation isused to replace the top- k operator.
Revised: Softmax Augmentation for smoothed top- k operation: A combination of k softmax operation is used to replace the top- k operator.

Original: Speciﬁcally, we recursively perform softmax on X for k times(Similar idea appears in Plötz and Roth (2018)).
Revised: Speciﬁcally, we recursively perform softmax on X for k times (Similar idea appears in Plötz and Roth (2018)).

Original: At the k -th iteration, we mask the top- ( k − 1)entries with negative inﬁnity.6.
Revised: At the k -th iteration, we mask the top- ( k − 1) entries with negative inﬁnity.

Original: CNNs trained with CE without any top- k component 2 .
Revised: CNNs trained with CE without any top- k component 4 .

Original: For the pretrained CNN and CNN trained with CE, we adopt identical neural networks as our method.
Revised: For the pretrained CNN and CNN trained with CE, we adopt identical neural networks as our method.

Original: Results.
Revised: Results.

Original: We report the classiﬁcation accuracies on the standard test sets in Table 1.
Revised: We report the classiﬁcation accuracies on the standard test sets in Table 1.

Original: On both datasets,the SOFT k NN classiﬁer achieves comparable or better accuracies.
Revised: On both datasets, the SOFT k NN classiﬁer achieves comparable or better accuracies.

Original: D Beam Search for Machine Translation
Revised: D Beam Search for Machine Translation

Original: Beam search is a popular method for the inference of Neural Language Generation (NLG) models,e.g., machine translation models.
Revised: Beam search is a popular method for the inference of Neural Language Generation (NLG) models, e.g., machine translation models.

Original: Here, we propose to incorporate beam search into the trainingprocedure based on SOFT top- k operator.
Revised: Here, we propose to incorporate beam search into the training procedure based on SOFT top- k operator.

Original: D.1 Misalignment between Training and Inference
Revised: D.1 Misalignment between Training and Inference

Original: Denote the predicted sequence as y =
Revised: Denote the predicted sequence as y =

Original: [ y (1) , · · · , y ( T ) ] , and the vocabularies as { z 1 , · · · , z V } .
Revised: [ y (1) , · · · , y ( T ) ] , and the vocabularies as { z 1 , · · · , z V } .

Original: Consider a recurrent network based NLG model.
Revised: Consider a recurrent network based NLG model.

Original: The output of the model at the t -th decoding stepis a probability simplex
Revised: The output of the model at the t -th decoding step is a probability simplex

Original: [ P ( y ( t ) =
Revised: [ P ( y ( t ) =

Original: z i | h ( t ) ]
Revised: z i | h ( t ) ]

Original: Vi =1 , where h ( t ) is the hidden state associated with thesequence y (1: t ) =
Revised: Vi =1 , where h ( t ) is the hidden state associated with the sequence y (1: t ) =

Original: [ y (1) , ..., y ( t ) ] .
Revised: [ y (1) , ..., y ( t ) ] .

Original: Beam search recursively keeps the sequences with the k largest likelihoods, and discards the rest.
Revised: Beam search recursively keeps the sequences with the k largest likelihoods, and discards the rest.

Original: Speciﬁcally, at the ( t + 1) -th decoding step, we have k sequences (cid:101) y (1: t ) ,i ’s obtained at the t -th step, where i = 1 , ..., k indexes the sequences.
Revised: Speciﬁcally, at the ( t + 1) -th decoding step, we have k sequences (cid:101) y (1: t ) ,i ’s obtained at the t -th step, where i = 1 , ..., k indexes the sequences.

Original: The likelihood of (cid:101) y (1: t ) ,i is denoted by L s ( (cid:101) y (1: t ) ,i ) .
Revised: The likelihood of (cid:101) y (1: t ) ,i is denoted by L s ( (cid:101) y (1: t ) ,i ) .

Original: Wethen select the next k sequences by varying i = 1 , . . .
Revised: We then select the next k sequences by varying i = 1 , . . .

Original: , k and j = 1 , . .
Revised: , k and j = 1 , . .

Original: , V :
Revised: , V :

Original: [Equation]
Revised: [Equation]

Original: where L s ([ (cid:101) y (1: t ) ,i , z j ]) is the likelihood of the sequence appending z j to (cid:101)
Revised: where L s ([ (cid:101) y (1: t ) ,i , z j ]) is the likelihood of the sequence appending z j to (cid:101)

Original: y
Revised: y

Original: (1: t ) ,i deﬁned as
Revised: (1: t ) ,i deﬁned as

Original: [Equation]
Revised: [Equation]

Original: and h ( t +1) ,i is the hidden state generated from (cid:101) y
Revised: and h ( t +1) ,i is the hidden state generated from (cid:101) y

Original: (1: t ) ,i .
Revised: (1: t ) ,i .

Original: Note that z j ’s and (cid:101) y
Revised: Note that z j ’s and (cid:101) y

Original: (1: t ) ,i ’s together yield V kchoices.
Revised: (1: t ) ,i ’s together yield V k choices.

Original: Here we abuse the notation: (cid:101) y (1: t +1) ,(cid:96) denotes the (cid:96) -th selected sequence at the ( t + 1) -thdecoding step, and is not necessarily related to (cid:101) y (1: t ) ,i at the t -th decoding step, even if i = (cid:96) .
Revised: Here we abuse the notation: (cid:101) y (1: t +1) ,(cid:96) denotes the (cid:96) -th selected sequence at the ( t + 1) -th decoding step, and is not necessarily related to (cid:101) y (1: t ) ,i at the t -th decoding step, even if i = (cid:96) .

Original: For t = 1 , we set (cid:101) y (1) =
Revised: For t = 1 , we set (cid:101) y (1) =

Original: z s as the start token, L s ( y (1) )
Revised: z s as the start token, L s ( y (1) )

Original: = 1 , and h (1) =
Revised: = 1 , and h (1) =

Original: h e as the output of theencoder.
Revised: h e as the output of the encoder.

Original: We repeat the above procedure, until the end token is selected or the pre-speciﬁed maxlength is reached.
Revised: We repeat the above procedure, until the end token is selected or the pre-speciﬁed max length is reached.

Original: At last, we select the sequence y (1: T ) , ∗ with the largest likelihood as the prediction.
Revised: At last, we select the sequence y (1: T ) , ∗ with the largest likelihood as the prediction.

Original: Moreover, the most popular training procedure for NLG models directly uses the so-called“ teacher forcing ” framework.
Revised: Moreover, the most popular training procedure for NLG models directly uses the so-called “ teacher forcing ” framework.

Original: As the ground truth of the target sequence (i.e., gold sequence)¯ y =
Revised: As the ground truth of the target sequence (i.e., gold sequence) ¯ y =

Original: [¯ y (1) , · · · , ¯ y ( T ) ] is provided at the training stage, we can directly maximize the likelihood
Revised: [¯ y (1) , · · · , ¯ y ( T ) ] is provided at the training stage, we can directly maximize the likelihood

Original: [Equation]
Revised: [Equation]

Original: As can be seen, such a training framework only involve the gold sequence, and cannot take theuncertainty of the recursive exploration of the beam search into consideration.
Revised: As can be seen, such a training framework only involve the gold sequence, and cannot take the uncertainty of the recursive exploration of the beam search into consideration.

Original: Therefore, it yields amisalignment between model training and inference (Bengio et al., 2015), which is also referred asexposure bias (Wiseman and Rush, 2016).
Revised: Therefore, it yields a misalignment between model training and inference (Bengio et al., 2015), which is also referred as exposure bias (Wiseman and Rush, 2016).

Original: D.2 Differential Beam Search with Sorted SOFT Top- k
Revised: D.2 Differential Beam Search with Sorted SOFT Top- k

Original: To mitigate the aforementioned misalignment, we propose to integrate beam search into the trainingprocedure, where the top- k operator in the beam search algorithm is replaced with our proposedsorted SOFT top- k operator proposed in Section B.
Revised: To mitigate the aforementioned misalignment, we propose to integrate beam search into the training procedure, where the top- k operator in the beam search algorithm is replaced with our proposed sorted SOFT top- k operator proposed in Section B.

Original: Speciﬁcally, at the ( t + 1) -th decoding step, we have k sequences denoted by E (1: t ) ,i , wherei = 1 , ..., k indexes the sequences.
Revised: Speciﬁcally, at the ( t + 1) -th decoding step, we have k sequences denoted by E (1: t ) ,i , where i = 1 , ..., k indexes the sequences.

Original: Here E (1: t ) ,i consists of a sequence of D -dimensional vectors,where D is the embedding dimension.
Revised: Here E (1: t ) ,i consists of a sequence of D -dimensional vectors, where D is the embedding dimension.

Original: We are not using the tokens, and the reason behind will beexplained later.
Revised: We are not using the tokens, and the reason behind will be explained later.

Original: Let (cid:101) h ( t ) ,i denote the hidden state generated from E (1: t ) ,i .
Revised: Let (cid:101) h ( t ) ,i denote the hidden state generated from E (1: t ) ,i .

Original: We then consider
Revised: We then consider

Original: [Equation]
Revised: [Equation]

Original: where L s ( · ) is deﬁned analogously to (7), and w j ∈ R D is the embedding of token z j .
Revised: where L s ( · ) is deﬁned analogously to (7), and w j ∈ R D is the embedding of token z j .

Original: Recall that (cid:15) is the smoothing parameter.
Revised: Recall that (cid:15) is the smoothing parameter.

Original: We then apply the sorted SOFT top- k operator to X ( t ) toobtain { E (1: t +1) ,(cid:96) } k(cid:96)
Revised: We then apply the sorted SOFT top- k operator to X ( t ) to obtain { E (1: t +1) ,(cid:96) } k(cid:96)

Original: =1 , which are k sequences with the largest likelihoods.
Revised: =1 , which are k sequences with the largest likelihoods.

Original: More precisely, the sorted
Revised: More precisely, the sorted

Original: SOFT top- k operator yields an output tensor A ( t ) ,
Revised: SOFT top- k operator yields an output tensor A ( t ) ,

Original: (cid:15) ∈ R V × k × k , where A ( t ) ,(cid:15) ji,(cid:96) denotes the smoothedindicator of whether [ E (1: t ) ,i , w j ] has a rank (cid:96) .
Revised: (cid:15) ∈ R V × k × k , where A ( t ) ,(cid:15) ji,(cid:96) denotes the smoothed indicator of whether [ E (1: t ) ,i , w j ] has a rank (cid:96) .

Original: We then obtain
Revised: We then obtain

Original: [Equation]
Revised: [Equation]

Original: where r denotes the index i (for E (1: t ) ,i ’s) associated with the index (cid:96) (for E (1: t +1) ,(cid:96) ’s).
Revised: where r denotes the index i (for E (1: t ) ,i ’s) associated with the index (cid:96) (for E (1: t +1) ,(cid:96) ’s).

Original: This is whywe use vector representations instead of tokens: this allows us to compute E ( t +1) ,(cid:96) as a weightedsum of all the word embeddings [ w j
Revised: This is why we use vector representations instead of tokens: this allows us to compute E ( t +1) ,(cid:96) as a weighted sum

Original: ] Vj =1 , instead of discarding the un-selected words.

Original: Accordingly, we generate the k hidden states for the ( t + 1) -th decoding step:
Revised: Accordingly, we generate the k hidden states for the ( t + 1) -th decoding step:

Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: where h ( t ) ,i is the hidden state generated by the decoder based on E (1: t ) ,i .
Revised: where h ( t ) ,i is the hidden state generated by the decoder based on E (1: t ) ,i .

Original: After decoding, we select the sequence with largest likelihood E (1: T ) , ∗ , and maximize the likelihoodas follows,
Revised: After decoding, we select the sequence with largest likelihood E (1: T ) , ∗ , and maximize the likelihood as follows,

Original: [Equation]
Revised: [Equation]

Original: We provide the sketch of training procedure in Algorithm 2, where we denote logit ( t ) ,i as[log P ( y ( t +1) =
Revised: We provide the sketch of training procedure in Algorithm 2, where we denote logit ( t ) ,i as [log P ( y ( t +1) =

Original: ω
Revised: ω

Original: j | (cid:101) h ( t ) ,i ( E (1: t ) ,i ))]
Revised: j | (cid:101) h ( t ) ,i ( E (1: t ) ,i ))]

Original: Vj =1 , which is part of the output of the decoder.
Revised: Vj =1 , which is part of the output of the decoder.

Original: More technical details (e.g., backtracking algorithm for ﬁnding the index r in (9)) are provided in Appendix
Revised: More technical details (e.g., backtracking algorithm for ﬁnding the index r in (9)) are provided in Appendix

Original: I.
Revised: I.

Original: Require:
Revised: Require:

Original: Input sequence s , target sequence ¯ y ;
Revised: Input sequence s , target sequence ¯ y ;

Original: embedding matrix W ∈ R V × D ; max length T ; k ; regularization coefﬁcient (cid:15) ; number of Sinkhorn iteration L (cid:101) h (1) i = h e = Encoder ( s ) , E (1) ,i = w s for t = 1 , · · · , T − 1 do for i = 1 , · · · , k do logit ( t ) ,i , h ( t ) ,i =
Revised: embedding matrix W ∈ R V × D ; max length T ; k ; regularization coefﬁcient (cid:15) ; number of Sinkhorn iteration L (cid:101) h (1) i = h e = Encoder ( s ) , E (1) ,i = w s for t = 1 , · · · , T − 1 do for i = 1 , · · · , k do logit ( t ) ,i , h ( t ) ,i =

Original: Decoder ( E ( t ) ,i , (cid:101) h ( t ) ,i ) log L s ([ E (1: t ) ,i , w j ])
Revised: Decoder ( E ( t ) ,i , (cid:101) h ( t ) ,i ) log L s ([ E (1: t ) ,i , w j ])

Original: =
Revised: =

Original: log L s ( E (1: t ) ,i )+ logit( t ) ,i j
Revised: log L s ( E (1: t ) ,i )+ logit

Original: X ( t ) =
Revised: [Table]

Original: {− log L s ([ E (1: t ) ,i , w j ]) | j = 1 , · · · , V } end for A ( t ) ,(cid:15) = Sorted-SOFT-Top- k
Revised: the beam

Original: ( X ( t ) , k, (cid:15), L ) Compute E ( t +1) ,(cid:96) , (cid:101) h ( t +1) ,(cid:96) as in (9) and (10) end for Compute ∇L SOFT and update the model
Revised: [Table]

Original: Note that integrating the beam search intotraining essentially yields a very largesearch space for the model, which is notnecessarily affordable sometimes.
Revised: into training essentially yields a very large search space for the model, which is not necessarily affordable sometimes.

Original: To alleviate this issue, we further propose ahybrid approach by combining the teacherforcing training with beam search-typetraining.
Revised: To alleviate this issue, we further propose a hybrid approach by combining the teacher forcing training with beam search-type training.

Original: Speciﬁcally, we maximize theweighted likelihood deﬁned as follows,
Revised: Speciﬁcally, we maximize the weighted likelihood deﬁned as follows,

Original: [Equation]
Revised: [Equation]

Original: where ρ ∈ (0 , 1) is referred to as the“teaching forcing ratio”.
Revised: where ρ ∈ (0 , 1) is referred to as the “teaching forcing ratio”.

Original: The teaching forcing loss L tf can help reduce the searchspace and improve the overall performance.
Revised: The teaching forcing loss L tf can help reduce the search space and improve the overall performance.

Original: D.3 Experiment
Revised: D.3 Experiment

Original: We evaluate our proposed beam search+ sorted SOFT top- k training procedureusing WMT2014 English-French dataset.
Revised: We evaluate our proposed beam search + sorted SOFT top- k training procedure using WMT2014 English-French dataset.

Original: We adopt beam size 5 , teacher forcing ratio ρ = 0 . 8 , and (cid:15) =
Revised: We adopt beam size 5 , teacher forcing ratio ρ = 0 . 8 , and (cid:15) =

Original: 10 − 1 .
Revised: 10 − 1 .

Original: For detailed settings of thetraining procedure, please refer to Appendix I.
Revised: For detailed settings of the training procedure, please refer to Appendix I.

Original: We reproduce the experiment in Bahdanau et al.
Revised: We reproduce the experiment in Bahdanau et al.

Original: (2014), and run our proposed training procedurewith the identical data pre-processing procedure and the LSTM-based sequence-to-sequence model.
Revised: (2014), and run our proposed training procedure with the identical data pre-processing procedure and the LSTM-based sequence-to-sequence model.

Original: Different from Bahdanau et al.
Revised: Different from Bahdanau et al.

Original: (2014), here we also preprocess the data with byte pair encoding(Sennrich et al., 2015).
Revised: (2014), here we also preprocess the data with byte pair encoding (Sennrich et al., 2015).

Original: Results.
Revised: Results.

Original: As shown in Table 2, the proposed SOFT beam search training procedure achieves animprovement in BLEU score of approximately 0 .
Revised: As shown in Table 2, the proposed SOFT beam search training procedure achieves an improvement in BLEU score of approximately 0 .

Original: We also include other LSTM-based models forbaseline comparison.
Revised: We also include other LSTM-based models for baseline comparison.

Original: E Related Work
Revised: E Related Work

Original: We parameterize the top- k operator as an optimal transport problem, which shares the same spirit as
Revised: We parameterize the top- k operator as an optimal transport problem, which shares the same spirit as

Original: Cuturi et al.
Revised: Cuturi et al.

Original: Speciﬁcally, Cuturi et al.
Revised: Speciﬁcally, Cuturi et al.

Original: (2019) formulate the ranking and sorting problems as
Revised: (2019) formulate the ranking and sorting problems as

Original: OT problems.
Revised: OT problems.

Original: Ranking is more complicated than identifying the top- k elements, since one needs toalign different ranks to corresponding elements.
Revised: Ranking is more complicated than identifying the top- k elements, since one needs to align different ranks to corresponding elements.

Original: Therefore, the algorithm complexity per iterationfor ranking whole n elements is O ( n 2 ) .
Revised: Therefore, the algorithm complexity per iteration for ranking whole n elements is O ( n 2 ) .

Original: Cuturi et al.
Revised: Cuturi et al.

Original: (2019) also propose an OT problem for ﬁndingthe τ -quantile in a set of n elements and the algorithm complexity reduces to O ( n ) .
Revised: (2019) also propose an OT problem for ﬁnding the τ -quantile in a set of n elements and the algorithm complexity reduces to O ( n ) .

Original: Top-
Revised: Top-

Original: k operatoressentially ﬁnds all the elements more extreme than the ( n − k ) /n
Revised: k operator essentially ﬁnds all the elements more extreme than the ( n − k ) /n

Original: -quantile, and our proposedalgorithm achieves the same complexity O ( n ) per iteration.
Revised: -quantile, and our proposed algorithm achieves the same complexity O ( n ) per iteration.

Original: The difference is that top- k operatorreturns the top- k elements in a given input set, while ﬁnding a quantile only yields a certain threshold.
Revised: The difference is that top- k operator returns the top- k elements in a given input set, while ﬁnding a quantile only yields a certain threshold.

Original: Gumbel-Softmax trick (Jang et al., 2016) can also be utilized to derive a continuous relaxation ofthe top- k operator.
Revised: Gumbel-Softmax trick (Jang et al., 2016) can also be utilized to derive a continuous relaxation of the top- k operator.

Original: Speciﬁcally, Kool et al.
Revised: Speciﬁcally, Kool et al.

Original: (2019) adapted such a trick to sample k elements fromn choices, and Xie and Ermon (2019) further applied the trick to stochastic k NN, where neuralnetworks are used to approximating the sorting operator.
Revised: (2019) adapted such a trick to sample k elements from n choices, and Xie and Ermon (2019) further applied the trick to stochastic k NN, where neural networks are used to approximating the sorting operator.

Original: However, as shown in our experiments (see
Revised: However, as shown in our experiments (see

Original: Table 1), the performance of stochastic k NN is not as good as deterministic k NN.
Revised: Table 1), the performance of stochastic k NN is not as good as deterministic k NN.

Original: Our SOFT beam search training procedure is inspired by several works that incorporate some ofthe characteristics of beam search into the training procedure (Wiseman and Rush, 2016; Goyalet al., 2018; Bengio et al., 2015).
Revised: Our SOFT beam search training procedure is inspired by several works that incorporate some of the characteristics of beam search into the training procedure (Wiseman and Rush, 2016; Goyal et al., 2018; Bengio et al., 2015).

Original: Speciﬁcally, Wiseman and Rush (2016) and Goyal et al.
Revised: Speciﬁcally, Wiseman and Rush (2016) and Goyal et al.

Original: (2018)both address the exposure bias issue in beam search.
Revised: both address the exposure bias issue in beam search.

Original: Wiseman and Rush (2016) propose a newloss function in terms of the error made during beam search.
Revised: Wiseman and Rush (2016) propose a new loss function in terms of the error made during beam search.

Original: This mitigates the misalignment of training and testing in beam search.
Revised: This mitigates the misalignment of training and testing in beam search.

Original: Later, Goyal et al.
Revised: Later, Goyal et al.

Original: (2018) approximates the top- k operator usingk softmax operations (This method is described and compared to our proposed method in C).
Revised: (2018) approximates the top- k operator using k softmax operations (This method is described and compared to our proposed method in C).

Original: Suchan approximation allows an end-to-end training of beam search.
Revised: Such an approximation allows an end-to-end training of beam search.

Original: Besides, our proposed training loss
Revised: Besides, our proposed training loss

Original: L ﬁnal is inspired by Bengio et al.
Revised: L ﬁnal is inspired by Bengio et al.

Original: (2015), which combines teacher forcing training procedure and greedy decoding, i.e., beam search with beam size 1 .
Revised: (2015), which combines teacher forcing training procedure and greedy decoding, i.e., beam search with beam size 1 .

Original: F Discussion
Revised: F Discussion

Original: Relation to automatic differentiation.
Revised: Relation to automatic differentiation.

Original: We compute the Jacobianmatrix of SOFT top- k operator directly in the backward pass.
Revised: We compute the Jacobian matrix of SOFT top- k operator directly in the backward pass.

Original: The
Revised: The

Original: OT plan can be obtained by the Sinkhorn algorithm (Algorithm 1),which is iterative and each iteration only involves multiplication andaddition.
Revised: OT plan can be obtained by the Sinkhorn algorithm (Algorithm 1), which is iterative and each iteration only involves multiplication and addition.

Original: Therefore, we can also apply automatic differentiation(auto-diff) to compute the Jacobian matrix.
Revised: Therefore, we can also apply automatic differentiation (auto-diff) to compute the Jacobian matrix.

Original: Speciﬁcally, we denoteΓ (cid:96) as the transport plan at the t -th iteration of Sinkhorn algorithm.
Revised: Speciﬁcally, we denote Γ (cid:96) as the transport plan at the t -th iteration of Sinkhorn algorithm.

Original: The update of Γ (cid:96) can be written as Γ (cid:96) +

Revised: The update of Γ (cid:96) can be written as Γ (cid:96) +

Original: = T (Γ (cid:96) ) , where T denotesthe update of the Sinkhorn algorithm.
Revised: = T (Γ (cid:96) ) , where T denotes the update of the Sinkhorn algorithm.

Original: In order to apply auto-diff,we need to store all the intermediate states, e.g., p, q, G in eachiteration, as deﬁned in Algorithm 1 at each iteration.
Revised: In order to apply auto-diff, we need to store all the intermediate states, e.g., p, q, G in each iteration, as deﬁned in Algorithm 1 at each iteration.

Original: This requires a huge memory size proportionalto the number of iterations of the algorithm.
Revised: This requires a huge memory size proportional to the number of iterations of the algorithm.

Original: In contrast, our backward pass allows us to save memory.
Revised: In contrast, our backward pass allows us to save memory.

Original: Bias and regularization trade-off.
Revised: Bias and regularization trade-off.

Original: Theorem 2 suggests a trade-off between the regularization andbias of SOFT top- k operator.
Revised: Theorem 2 suggests a trade-off between the regularization and bias of SOFT top- k operator.

Original: Speciﬁcally, a large (cid:15) has a strong smoothing effect on the entropic OTproblem, and the corresponding entries of the Jacobian matrix are neither too large nor too small.
Revised: Speciﬁcally, a large (cid:15) has a strong smoothing effect on the entropic OT problem, and the corresponding entries of the Jacobian matrix are neither too large nor too small.

Original: This eases the end-to-end training process.
Revised: This eases the end-to-end training process.

Original: However, the bias of SOFT top- k operator is large, whichcan deteriorate the model performance.
Revised: However, the bias of SOFT top- k operator is large, which can deteriorate the model performance.

Original: On the contrary, a smaller (cid:15) ensures a smaller bias.
Revised: On the contrary, a smaller (cid:15) ensures a smaller bias.

Original: Yet
Revised: Yet

Original: the
Revised: the

Original: SOFT top-
Revised: SOFT top-

Original: k operator is less smooth, which in turn makes the end-to-end training less efﬁcient.
Revised: k operator is less smooth, which in turn makes the end-to-end training less efﬁcient.

Original: On the other hand, the bias of SOFT top- k operator also depends on the gap between x σ k +1 andx σ k .
Revised: On the other hand, the bias of SOFT top- k operator also depends on the gap between x σ k +1 and x σ k .

Original: In fact, such a gap can be viewed as the signal strength of the problem.
Revised: In fact, such a gap can be viewed as the signal strength of the problem.

Original: A large gap impliesthat the top- k elements are clearly distinguished from the rest of the elements.
Revised: A large gap implies that the top- k elements are clearly distinguished from the rest of the elements.

Original: Therefore, the biasis expected to be small since the problem is relatively easy.
Revised: Therefore, the bias is expected to be small since the problem is relatively easy.

Original: Moreover, in real applications such asneural network-based k NN classiﬁcation, the end-to-end training process promotes neural networksto extract features that exhibit a large gap (as illustrated in Figure 5).
Revised: Moreover, in real applications such as neural network-based k NN classiﬁcation, the end-to-end training process promotes neural networks to extract features that exhibit a large gap (as illustrated in Figure 5).

Original: Hence, the bias of SOFT top- koperator can be well controlled in practice.
Revised: Hence, the bias of SOFT top- k operator can be well controlled in practice.

Original: G Theoretical Guarantees
Revised: G Theoretical Guarantees

Original: Proposition 1.
Revised: Proposition 1.

Original: Consider the setup in the previous paragraph.
Revised: Consider the setup in the previous paragraph.

Original: Without loss of generality, we assume
Revised: Without loss of generality, we assume

Original: X has no duplicates.
Revised: X has no duplicates.

Original: Then the optimal transport plan Γ ∗ of (1) is
Revised: Then the optimal transport plan Γ ∗ of (1) is

Original: [Equation]
Revised: [Equation]

Original: (11) with σ being the sorting permutation, i.e., x σ 1 < x σ 2 < · · · < x σ n .
Revised: with σ being the sorting permutation, i.e., x σ 1 < x σ 2 < · · · < x σ n .

Original: Moreover, we have A = n Γ ∗ · [1 , 0] (cid:62) . 
Revised: Moreover, we have A = n Γ ∗ · [1 , 0] (cid:62) . [Equation]

Original: Proof.
Revised: Proof.

Original: We expand the objective function of (1) as
Revised: We expand the objective function of (1) as

Original: [Equation]
Revised: [Equation]

Original: n

Original: [Equation]

Original: n − k n −n (cid:88) i

Original: = x i Γ i, 2 .

Original: Therefore, to minimize (cid:104) C, Γ (cid:105) , it sufﬁces to maximize (cid:80) ni
Revised: Therefore, to minimize (cid:104) C, Γ (cid:105) , it sufﬁces to maximize (cid:80) ni

Original: =1 x i Γ i, 2 .
Revised: =1 x i Γ i, 2 .

Original: It is straightforward to check n(cid:88) i
Revised: It is straightforward to check n

Original: = Γ i, 2 = n − k n and Γ i, 2 ≤ 1 n for any i = 1 , . . . , n .
Revised: [Equation] for any i = 1 , . . . , n .

Original: Hence, maximizing (cid:80) n i =1 x i Γ i, 2 is essentially selecting the largest n − Kelements from X , and the maximum is attained atΓ ∗ σ i , 2 = (cid:26) 0 , if i ≤ k, 1 /n, if k + 1 ≤ i ≤ n.
Revised: Hence, maximizing (cid:80) ni =1 x i Γ i, 2 is essentially selecting the largest n − K elements from X , and the maximum is attained at [Equation]

Original: The constraint Γ 1 m = µ further implies that Γ ∗ i, 1 satisﬁes (11).
Revised: The constraint Γ 1 m = µ further implies that Γ ∗ i, 1 satisﬁes (11).

Original: Thus, A can be parameterized as
Revised: Thus, A can be parameterized as

Original: A = n Γ ∗ · [1 , 0] (cid:62) .
Revised: A = [Equation] [Equation] [Equation]

Revised: [Equation]

Original: We then show that after adding entropy regularization the problem is differentiable.
Revised: We then show that after adding entropy regularization the problem is differentiable.

Original: Theorem 1.
Revised: Theorem 1.

Original: For any (cid:15) > 0 , SOFT top- k operator: X (cid:55)→ A (cid:15) is differentiable, as long as the cost C ijis differentiable with respect to x i for any i, j .
Revised: For any (cid:15) > 0 , SOFT top- k operator: X (cid:55)→ A (cid:15) is differentiable, as long as the cost C ij is differentiable with respect to x i for any i, j .

Original: Moreover, the Jacobian matrix of SOFT top- k operatoralways has a nonzero entry for any X ∈ R n .
Revised: Moreover, the Jacobian matrix of SOFT top- k operator always has a nonzero entry for any X ∈ R n .

Original: Proof.
Revised: Proof.

Original: We ﬁrst prove the differentiability.
Revised: We ﬁrst prove the differentiability.

Original: This part of proof mirrors the proof in Luise et al.
Revised: This part of proof mirrors the proof in Luise et al.

Original: By Sinkhorn’s scaling theorem,
Revised: By Sinkhorn’s scaling theorem,

Original: [Equation]
Revised: [Equation]

Original: Therefore, since C ij is differentiable, Γ ∗ ,(cid:15) is differentiable if ( ξ ∗ , ζ ∗ ) is differentiable as a functionof input scores X .
Revised: Therefore, since C ij is differentiable, Γ [Equation] ,(cid:15) is differentiable if ( ξ ∗ , ζ ∗ ) is differentiable as a function of input scores X .

Original: Let us set
Revised: Let us set

Original: [Equation]
Revised: [Equation]

Original: and recall that ( ξ ∗ , ζ ∗ ) = argmax ξ,ζ L ( ξ, ζ ; µ, ν, C ) .
Revised: and recall that ( ξ ∗ , ζ ∗ ) = argmax ξ,ζ L ( ξ, ζ ; µ, ν, C ) .

Original: The differentiability of ( ξ ∗ , ζ ∗ ) is proved usingthe Implicit Function theorem and follows from the differentiability and strong convexity in ( ξ ∗ , ζ ∗ )of the function L .
Revised: The differentiability of ( ξ ∗ , ζ ∗ ) is proved using the Implicit Function theorem and follows from the differentiability and strong convexity in ( ξ ∗ , ζ ∗ ) of the function L .

Original: Now we prove that dA (cid:15) /dx
Revised: Now we prove that dA (cid:15) /dx

Original: (cid:96) always has a nonzero entry for l = 1 , · · · , n .
Revised: (cid:96) always has a nonzero entry for l = 1 , · · · , n .

Original: First, we prove thatfor any (cid:96) ∈ { 1 , · · · , n } , d Γ ∗ ,(cid:15) /dx
Revised: First, we prove that for any (cid:96) ∈ { 1 , · · · , n } , d Γ ∗ ,(cid:15) /dx

Original: (cid:96) always has a nonzero entry.
Revised: (cid:96) always has a nonzero entry.

Original: We will prove it by contradiction.
Revised: We will prove it by contradiction.

Original: Speciﬁcally, the KKT conditions for the stationarity are as follows
Revised: Speciﬁcally, the KKT conditions for the stationarity are as follows

Original: [Equation]
Revised: [Equation]

Original: If we view the above formula as a linear equation set of the dual variables, it has nm equations andm + n variables.
Revised: If we view the above formula as a linear equation set of the dual variables, it has nm equations and m + n variables.

Original: Therefore, there are nm − m − n redundant equations.
Revised: Therefore, there are nm − m − n redundant equations.

Original: Suppose one of the scoresx (cid:96) , has an inﬁnitesimal change δx (cid:96) .
Revised: Suppose one of the scores x (cid:96) , has an inﬁnitesimal change δx (cid:96) .

Original: Assuming Γ ∗ ,(cid:15) does not change, we have a new set of linearequations,
Revised: Assuming Γ ∗ ,(cid:15) does not change, we have a new set of linear equations,

Original: [Equation]
Revised: [Equation]

Original: Easy to verify that this set of linear equations has no solution.
Revised: Easy to verify that this set of linear equations has no solution.

Original: Therefore, there must be at least oneentry in Γ ∗ ,(cid:15) has changed.
Revised: Therefore, there must be at least one entry in Γ ∗ ,(cid:15) has changed.

Original: As a result, d Γ ∗ ,(cid:15) /dx
Revised: As a result, d Γ ∗ ,(cid:15) /dx

Original: (cid:96) always has a nonzero entry.
Revised: (cid:96) always has a nonzero entry.

Original: We denote this entry asΓ ∗ ,(cid:15)i (cid:48) j (cid:48) .
Revised: We denote this entry as Γ ∗ ,(cid:15)i (cid:48) j (cid:48) .

Original: Since Γ ∗ ,(cid:15)i (cid:48)
Revised: Since Γ ∗ ,(cid:15)i (cid:48)

Original: j (cid:48)
Revised: j (cid:48)

Original: + Γ ∗ ,(cid:15)i (cid:48) , 3 − j (cid:48)
Revised: + Γ ∗ ,(cid:15)i (cid:48) , 3 − j (cid:48)

Original: = µ i (cid:48) , we have
Revised: = µ i (cid:48) , we have

Original: [Equation]
Revised: [Equation]

Original: Therefore, there must be a nonzero entry in the ﬁrst column of d Γ ∗ ,(cid:15) /dx
Revised: Therefore, there must be a nonzero entry in the ﬁrst column of d Γ ∗ ,(cid:15) /dx

Original: (cid:96) .
Revised: (cid:96) .

Original: Recall A (cid:15) is the ﬁrstcolumn of Γ ∗ ,(cid:15) .
Revised: Recall A (cid:15) is the ﬁrst column of Γ ∗ ,(cid:15) .

Original: As a result, there must be a nonzero entry in dA (cid:15) /dx
Revised: As a result, there must be a nonzero entry in dA (cid:15) /dx

Original: (cid:96) for any (cid:96) ∈ { 1 , · · · , n } .
Revised: (cid:96) for any (cid:96) ∈ { 1 , · · · , n } .

Original: Second, we would like to know after smoothness relaxation, how much bias is introduced to A (cid:15) .
Revised: Second, we would like to know after smoothness relaxation, how much bias is introduced to A (cid:15) .

Original: Lemma 1.
Revised: Lemma 1.

Original: Denote the feasible set of optimal transport problem as ∆ =
Revised: Denote the feasible set of optimal transport problem as ∆ =

Original: { Γ : Γ ∈ [0 , 1] n × m , Γ 1 m =µ, Γ 1 n = ν } .
Revised: { Γ : Γ ∈ [0 , 1] n × m , Γ 1 m = µ, Γ 1 n = ν } .

Original: Assume the optimal transport plan is unique.
Revised: Assume the optimal transport plan is unique.

Original: Denote Γ ∗ as the optimal transport plan,
Revised: Denote Γ ∗ as the optimal transport plan,

Original: [Equation]
Revised: [Equation]

Original: and Γ ∗ ,(cid:15) as the entropy regularized transport plan,
Revised: and Γ ∗ ,(cid:15) as the entropy regularized transport plan,

Original: [Equation]
Revised: [Equation]

Revised: =

Revised: [Equation]

Original: We can bound the difference between Γ ∗ and Γ ∗ ,(cid:15) to be
Revised: We can bound the difference between Γ ∗ and Γ ∗ ,(cid:15) to be

Original: [Equation]
Revised: [Equation]

Original: , where (cid:107) · (cid:107) F is the Frobenius norm, and B is a positive constant irrelevant to (cid:15) .
Revised: where (cid:107) · (cid:107) F is the Frobenius norm, and B is a positive constant irrelevant to (cid:15) .

Original: Proof.
Revised: Proof.

Original: Note that H (Γ) is the entropy function.
Revised: Note that H (Γ) is the entropy function.

Original: Since 0 ≤ Γ ij ≤ 1 and (cid:80) ij Γ ij = 1 for any Γ ∈ ∆ ,we can view ∆ as the subset of a simplex.
Revised: Since 0 ≤ Γ ij ≤ 1 and (cid:80) ij Γ ij = 1 for any Γ ∈ ∆ , we can view ∆ as the subset of a simplex.

Original: Therefore,
Revised: Therefore,

Original: H (Γ) is non-negative.
Revised: H (Γ) is non-negative.

Original: The maximum of H (Γ) in the simplex can be obtained at Γ ij ≡ 1 nm .
Revised: The maximum of H (Γ) in the simplex can be obtained at Γ ij ≡ 1 nm .

Original: Therefore themaximum value is (ln n + ln m ) .
Revised: Therefore the maximum value is (ln n + ln m ) .

Original: Therefore, 0 ≤ H (Γ) ≤ (ln n + ln m ) for any Γ ∈ ∆ .
Revised: Therefore, 0 ≤ H (Γ) ≤ (ln n + ln m ) for any Γ ∈ ∆ .

Original: Since H (Γ) ≥ 0 , we have f (cid:15) (Γ) ≤ f (Γ) for any Γ ∈ ∆ .
Revised: Since H (Γ) ≥ 0 , we have f (cid:15) (Γ) ≤ f (Γ) for any Γ ∈ ∆ .

Original: As a result, we have f (cid:15) (Γ ∗ ,(cid:15) )
Revised: As a result, we have f (cid:15) (Γ ∗ ,(cid:15) )

Original: ≤ f (Γ ∗ ) .
Revised: ≤ f (Γ ∗ ) .

Original: Inother words, we have
Revised: In other words, we have

Original: [Equation]
Revised: [Equation]

Original: Therefore,
Revised: Therefore,

Original: [Equation]
Revised: [Equation]

Original: Since the optimal transport problem is a linear optimization problem, Γ ∗ is one of the vertices of ∆ .
Revised: Since the optimal transport problem is a linear optimization problem, Γ ∗ is one of the vertices of ∆ .

Original: Denote e 0 , e 1 , · · · , e J as the vertices of ∆ , and without loss of generality we assume e 0{+
+} = Γ ∗ .
Revised: Denote e 0 , e 1 , · · · , e J as the vertices of ∆ , and without loss of generality we assume e {+
+} = Γ ∗ .

Original: SinceΓ ∗ ,(cid:15) ∈ ∆ , we can denote Γ ∗ ,(cid:15) =
Revised: Since Γ ∗ ,(cid:15) ∈ ∆ , we can denote Γ ∗ ,(cid:15) =

Original: (cid:80) Jj =

Revised: (cid:80) Jj =

Original: λ
Revised: λ

Original: j e j , where λ j ≥ 0 , and (cid:80) j λ j = 1 .
Revised: j e j , where λ j ≥ 0 , and (cid:80) j λ j = 1 .

Original: Since Γ ∗ is unique, wehave(cid:104) C, e j − e 0 (cid:105) > 0 , ∀ j = 1 , · · · , J.
Revised: Since Γ ∗ is unique, we have [Equation]

Original: Denote B j = (cid:104) C,
Revised: Denote B j = (cid:104) C,

Original: e j − e 0 (cid:105) .
Revised: e j − e 0 (cid:105) .

Original: Since the space we are considering is Euclidean space (if we reshape thematrices into vectors), we can write the inner product as 
Revised: Since the space we are considering is Euclidean space (if we reshape the matrices into vectors), we can write the inner product as B

Original: [Equation]
Revised: [Equation]

Original: So we have cos θ ( C,e j − e 0 ) > 0 .
Revised: So we have cos θ ( C,e j − e 0 ) > 0 .

Original: In other words, the angle between C and e j − e 0 is always smaller thanπ 2 .
Revised: In other words, the angle between C and e j − e 0 is always smaller than π 2 .

Original: Therefore, the angle between C and the afﬁne combination of e j − e 0 , namely (cid:80) Jj =

Revised: Therefore, the angle between C and the afﬁne combination of e j − e 0 , namely (cid:80) Jj =

Original: λ j
Revised: λ j

Original: ( e j − e 0 ) ,is also smaller than π 2 .
Revised: ( e j − e 0 ) , is also smaller than π 2 .

Original: More speciﬁcally, we have
Revised: More speciﬁcally, we have

Original: [Equation]
Revised: [Equation]

Revised: F .

Original: Therefore, we have
Revised: Therefore, we have

Original: [Equation]
Revised: [Equation]

Original: Denote B = min
Revised: Denote B = min

Original: j
Revised: j

Original: [Equation]

Original: j

Original: [Equation]
Revised: [Equation]

Original: F , and we have the conclusion.
Revised: and we have the conclusion.

Original: Remark 1.
Revised: Remark 1.

Original: In Theorem 1 we restricted the optimal solution to be unique, only for clarity purpose.
Revised: In Theorem 1 we restricted the optimal solution to be unique, only for clarity purpose.

Original: Ifit is not unique, similar conclusion holds, except that the proof is more tedious – instead of divide thevertices into e 0 and others, we need to divide it into the vertices that are optimal solutions and theothers.
Revised: If it is not unique, similar conclusion holds, except that the proof is more tedious – instead of divide the vertices into e 0 and others, we need to divide it into the vertices that are optimal solutions and the others.

Original: Lemma 2.
Revised: Lemma 2.

Original: At each of the vertices of ∆ , the entries of Γ are either 0 or 1 /n for Γ ∈ ∆ .
Revised: At each of the vertices of ∆ , the entries of Γ are either 0 or 1 /n for Γ ∈ ∆ .

Original: Proof.
Revised: Proof.

Original: The key idea is to prove by contradiction: If there exist i, j such that Γ ij ∈ (0 , 1 /n ) , then Γcannot be a vertex.
Revised: The key idea is to prove by contradiction: If there exist i, j such that Γ ij ∈ (0 , 1 /n ) , then Γ cannot be a vertex.

Original: To ease the discussion, we denote Z = n Γ .
Revised: To ease the discussion, we denote Z = n Γ .

Original: We will ﬁrst prove that the entries of Z are either 0 or 1at the vertices.
Revised: We will ﬁrst prove that the entries of Z are either 0 orat the vertices.

Original: Notice that
Revised: Notice that

Original: [Equation]
Revised: [Equation]

Original: If there exists an entry Z i (cid:48) ,j (cid:48) ∈ (0 , 1) , then
Revised: If there exists an entry Z i (cid:48) ,j (cid:48) ∈ (0 , 1) , then

Original: Z i (cid:48) , 3 − j (cid:48) ∈ (0 , 1) .2.
Revised: Z i (cid:48) , 3 − j (cid:48) ∈ (0 , 1) .

Original: there must exist i (cid:48)(cid:48) (cid:54) =
Revised: there must exist i (cid:48)(cid:48) (cid:54) =

Original: i (cid:48) , such that Z i (cid:48)(cid:48) ,j (cid:48) ∈ (0 , 1) .
Revised: i (cid:48) , such that Z i (cid:48)(cid:48) ,j (cid:48) ∈ (0 , 1) .

Original: This is because (cid:80) ni =1 Z i,j is an integer,and Z i (cid:48) ,j (cid:48) is not.
Revised: This is because (cid:80) ni =1 Z i,j is an integer, and Z i (cid:48) ,j (cid:48) is not.

Original: As a result, Z i (cid:48)(cid:48) , 3 − j (cid:48) ∈ (0 , 1) . 
Revised: As a result, Z i (cid:48)(cid:48) , 3 − j (cid:48) ∈ (0 , 1) . Therefore, consider δ ∈ ( − min { 1 − Z i (cid:48) .j (cid:48) , Z i (cid:48) ,j (cid:48) } , min { 1 − Z i (cid:48) .j (cid:48) , Z i (cid:48) ,j (cid:48) } ) and denote (cid:101)

Original: Therefore, consider δ ∈ ( − min { 1 − Z 
Revised: Z (1) ij =  

Revised: Z i (cid:48) ,j (cid:48)

Original:  i 
Revised: + δ, if i = i (cid:48) , j = j (cid:48) , Z i (cid:48) [Equation] j (cid:48)

Revised: [Equation]

Revised: δ, if i = i (cid:48) , j =

Revised: [Equation]

Revised: j (cid:48) ,

Revised: Z i (cid:48)(cid:48) ,j (cid:48) − δ, if i =

Revised: i (cid:48)(cid:48) , j = j (cid:48) ,

Revised: Z i (cid:48)(cid:48)

Original: [Equation]
Revised: [Equation]

Original: .j (cid:48),
Revised: − j (cid:48)

Original: Z i 
Revised: + δ, if i = i (cid:48)(cid:48) , j = 3 − j (cid:48) , Z i,j , otherwise . (cid:101)

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: } , min { 1 − Zi
Revised: =      Z

Original: [Equation]
Revised: [Equation]

Original: .j
Revised: (cid:48) ,

Original: [Equation]
Revised: [Equation]

Original: , Z i 
Revised: (cid:48) − δ, if i =

Original: [Equation]
Revised: [Equation]

Original:  j 
Revised: (cid:48) , j =

Original: [Equation]
Revised: [Equation]

Original: } ) and denote
Revised: (cid:48) , Z

Original: [Equation]
Revised: [Equation]

Original: We can easily verify that (cid:101)
Revised: We can easily verify that (cid:101)

Original: Z (1) /n, (cid:101)
Revised: Z (1) /n, (cid:101)

Original: Z (2) /n ∈ ∆ , and also Z =
Revised: Z (2) /n ∈ ∆ , and also Z =

Original: ( (cid:101) Z (1) + (cid:101) Z (2) ) / 2 .
Revised: ( (cid:101) Z (1) + (cid:101) Z (2) ) / 2 .

Original: Therefore, Z cannotbe a vertex.
Revised: Therefore, Z cannot be a vertex.

Original: Lemma 3.
Revised: Lemma 3.

Original: Given a set of scalar { x 1 , · · · , x n } , we sort it to be { x σ 1 , · · · , x σ n } .
Revised: Given a set of scalar { x 1 , · · · , x n } , we sort it to be { x σ 1 , · · · , x σ n } .

Original: If Euclidean squarecost is adopted, Γ ∗ has the following form,
Revised: If Euclidean square cost is adopted, Γ ∗ has the following form,

Original: [Equation]
Revised: [Equation]

Original: And min j
Revised: And min j

Revised: B j

Original: [Equation]
Revised: [Equation]

Original:  F is attained at at a vertex Γ ∗∗ , where Γ ∗∗ ij = Γ ∗ ij except that the σ k -th row and theσ k +1 -th row are swapped.
Revised: − e [Equation] F is attained at at a vertex Γ ∗∗ , where Γ ∗∗ ij = Γ ∗ ij except that the σ k -th row and the σ k +1 -th row are swapped.

Original: As a result, we have
Revised: As a result, we have

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: From Lemma 2, in each vertex the entries of Γ is either 0 or 1 /n .
Revised: From Lemma 2, in each vertex the entries of Γ is either 0 or 1 /n .

Original: Also, Γ ∗ ∈ ∆ =
Revised: Also, Γ ∗ ∈ ∆ =

Original: { Γ :Γ ∈ [0 , 1] n × m , Γ 1 m = 1 n /n, Γ 1 n = [ k/n, ( n − k ) /n ] (cid:62) } .
Revised: { Γ : Γ ∈ [0 , 1] n × m , Γ 1 m = 1 n /n, Γ 1 n = [ k/n, ( n − k ) /n ] (cid:62) } .

Original: Therefore, for the j -th vertex, there arek entries with value 1 /n in the ﬁrst row of Γ .
Revised: Therefore, for the j -th vertex, there are k entries with value 1 /n in the ﬁrst row of Γ .

Original: Denote the row indices of these k entries as I j , andΩ = { 1 , · · · , n } .
Revised: Denote the row indices of these k entries as I j , and Ω = { 1 , · · · , n } .

Original: Then for each vertex we have 
Revised: Then for each vertex we have Γ i, 1 = 1 /n, ∀ i ∈ I

Original: [Equation]
Revised: j Γ i, 1 = 0 , ∀ i ∈ Ω \I j Γ i, 2 = 1 /n, ∀ i ∈ Ω \I j Γ i, 2 = 0 , ∀ i ∈ I j .

Original: Denote I ∗ = { σ 1 , · · · , σ k } .
Revised: Denote I ∗ = { σ 1 , · · · , σ k } .

Original: We now prove that I ∗ corresponds to the optimal solution Γ ∗ .
Revised: We now prove that I ∗ corresponds to the optimal solution Γ ∗ .

Original: This isbecause for any j ∈ { 1 , · · · , J } 
Revised: This is because for any j ∈ { 1 , · · · , J } Γ( I j ) − Γ( I ∗ ) =  (cid:88) i ∈I j x 2 i + (cid:88) i ∈ Ω \I j ( x i − 1)  −  (cid:88) i ∈I ∗ x 2 i + (cid:88) i ∈

Revised: Ω \I ∗ ( x i − 1)  =   (cid:88) i ∈

Original: {+Ω x 
+}
Revised: {+Ω x 
+}

Original: [Equation]
Revised: i − (cid:88) i ∈ Ω \I j 2 x i + ( n − k )   −   (cid:88) i ∈ Ω x 2 i − (cid:88)

Revised: i ∈

Revised: Ω \I ∗ 2 x i + ( n − k ) 

Revised:  =  (cid:88) i ∈

Revised: Ω \I ∗ x i − (cid:88) i ∈

Original:  where the last step is because the elements with indices Ω \I j is the largest n − k elements.
Revised: Ω \I j x i   ≥ 0 , where the last step is because the elements with indices Ω \I j is the largest n − k elements.

Original: Thereforewe have Γ( I ∗ ) =
Revised: Therefore we have Γ( I ∗ ) =

Original: Γ ∗ .
Revised: Γ ∗ .

Original: Now let’s compute min j (cid:54) =
Revised: Now let’s compute min j (cid:54) =

Original: B j / (cid:107) e
Revised: B j / (cid:107) e

Original: j − e 0 (cid:107) .
Revised: j − e 0 (cid:107) .

Original: Denote set subtraction A − B as the set if elements thatbelongs to A but do not belong to B , and |A| as the number of elements in A . 
Revised: Denote set subtraction A − B as the set if elements that belongs to A but do not belong to B , and |A| as the number of elements in A . B

Original: [Equation]
Revised: [Equation]

Original: where the second line can be obtained by substituting the deﬁnition of B j .
Revised: where the second line can be obtained by substituting the deﬁnition of B j .

Original: Notice that I j −I ∗ ∈
Revised: Notice that I j −I ∗ ∈

Original: Ω \I ∗and I ∗ − I j ∈
Revised: Ω \I ∗ and I ∗ − I j ∈

Original: I ∗ .
Revised: I ∗ .

Original: Any element with index in Ω \I ∗ is larger than any element in I ∗ by at leastx σ K +1 − x σ K .
Revised: Any element with index in Ω \I ∗ is larger than any element in I ∗ by at least x σ K +1 − x σ K .

Original: Then we have
Revised: Then we have

Original: [Equation]
Revised: [Equation]

Original: where the last step is because for j (cid:54) = 0 , |I ∗ − I j | is at least 1 .
Revised: where the last step is because for j (cid:54) = 0 , |I ∗ − I j | is at least 1 .

Original: Also notice that the value n ( x σ k +1 − x σ k ) can be attained at I j ∗ = { σ 1 , · · · , σ k − 1 , σ k +1 } .
Revised: Also notice that the value n ( x σ k +1 − x σ k ) can be attained at I j ∗ = { σ 1 , · · · , σ k − 1 , σ k +1 } .

Original: Thereforewe have
Revised: Therefore we have

Original: [Equation]
Revised: [Equation]

Original: Theorem 2.
Revised: Theorem 2.

Original: Given a distinct sequence X and its sorting permutation σ , with Euclidean square costfunction, for the proposed top- k solver we have
Revised: Given a distinct sequence X and its sorting permutation σ , with Euclidean square cost function [Equation] for the proposed top- k solver we have

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: This is a direct conclusion with Lemma 1 and Lemma 3.
Revised: This is a direct conclusion with Lemma 1 and Lemma 3.

Original: H The Expression of the Gradient of A (cid:15)
Revised: H The Expression of the Gradient of A (cid:15)

Original: In this section we will derive the expression of dA (cid:15) /dx
Revised: In this section we will derive the expression of dA (cid:15) /dx

Original: i .
Revised: i .

Original: We ﬁrst list a few reminders that will beused later:
Revised: We ﬁrst list a few reminders that will be used later:

Original: • { x i } ni =1 is a scalar set to be solved for top- k . {
Revised: • { x i } ni =1 is a scalar set to be solved for top- k . {

Original: y j } mj =1 is taken to be { 0 , 1 } .
Revised: y j } mj =1 is taken to be { 0 , 1 } .

Original: • C ∈ R n × m is the cost matrix, usually deﬁned as C ij =
Revised: • C ∈ R n × m is the cost matrix, usually deﬁned as C ij =

Original: ( x i − y j ) 2 .
Revised: ( x i − y j ) 2 .

Original: • The loss function of entropic optimal transport is
Revised: • The loss function of entropic optimal transport is

Original: [Equation]
Revised: [Equation]

Original: where ∆ =
Revised: where ∆ =

Original: { Γ : Γ ∈ [0 , 1] n × m , Γ 1 m = µ, Γ 1 n = ν } .
Revised: { Γ : Γ ∈ [0 , 1] n × m , Γ 1 m = µ, Γ 1 n = ν } .

Original: • The dual problem of the above optimization problem is
Revised: • The dual problem of the above optimization problem is

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original: , ζ ; C )

Original: , where
Revised: where

Original: [Equation]
Revised: [Equation]

Original: And it is connected to the prime form by
Revised: And it is connected to the prime form by

Original: [Equation]
Revised: [Equation]

Original:  Algorithm 
Revised: The converged p, q in Algorithm 1 is actually e ξ ∗ (cid:15) and e ζ ∗ (cid:15) .

Original: If we obtain the expression for dξ ∗ dC and dζ ∗ dC , we can obtain the expression for dA (cid:15) dx i .
Revised: If we obtain the expression for dξ ∗ dC and dζ ∗ dC , we can obtain the expression for dA (cid:15) dx i .

Original: In this section only, we denote Γ =
Revised: In this section only, we denote Γ =

Original: Γ ∗ ,(cid:15) , to shorten the notation.
Revised: Γ ∗ ,(cid:15) , to shorten the notation.

Original: The multiplication of 3 rd-ordertensors mirrors the multiplication of matrices: we always use the last dimension of the ﬁrst input tomultiplies the ﬁrst dimension of the second input.
Revised: The multiplication of 3 rd-order tensors mirrors the multiplication of matrices: we always use the last dimension of the ﬁrst input to multiplies the ﬁrst dimension of the second input.

Original: We denote ¯ b = b : − 1 as b removing the last entry,¯ ν = ν : − 1 as ν removing the last entry, ¯Γ =
Revised: We denote ¯ b = b : − 1 as b removing the last entry, ¯ ν = ν : − 1 as ν removing the last entry, ¯Γ =

Original: Γ : , : − 1 as Γ removing the last column.
Revised: Γ : , : − 1 as Γ removing the last column.

Original: Theorem 3.
Revised: Theorem 3.

Original: dξ ∗ dC and dζ ∗ dC have the following expression,
Revised: dξ ∗ dC and dζ ∗ dC have the following expression,

Original: [Equation]
Revised: [Equation]

Original: where − H − 1 D ∈ R ( n + m − 1) × n × m , 0 ∈ R 1 × n × m , and
Revised: where − H − 1 D ∈ R ( n + m − 1) × n × m , 0 ∈ R 1 × n × m , and

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: Notice that there is one redundant dual variable, since µ 1 N = ν 1 M = 1 .
Revised: Notice that there is one redundant dual variable, since µ 1 N = ν 1 M = 1 .

Original: Therefore, we canrewrite L ( ξ, ζ ; C ) as
Revised: Therefore, we can rewrite L ( ξ, ζ ; C ) as

Original: [Equation]
Revised: [Equation]

Original: Denote
Revised: Denote

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Original: where
Revised: where

Original: [Equation]
Revised: [Equation]

Original: Since ( ξ ∗ , ¯ ζ ∗ ) is a maximum of L ( ξ, ¯ ζ ; C ), we haveφ
Revised: Since ( ξ ∗ , ¯ ζ ∗ ) is a maximum of L ( ξ, ¯ ζ ; C )

Original: [Equation]
Revised: [Equation]

Original: , ¯

Original: [Equation] [Equation] [Equation] 0 , ψ 
Revised: we have φ ( ξ ∗ , ¯ ζ ∗ , C ) = 0 , ψ ( ξ

Original: [Equation]
Revised: [Equation]

Original: , ¯ ζ ∗ , C )
Revised: , ¯ ζ ∗ , C )

Original: = 0 .
Revised: = 0 .

Original: Therefore,
Revised: Therefore,

Original: [Equation]
Revised: [Equation]

Original: Therefore,
Revised: Therefore,

Revised: [Equation]

Revised: ∂

Original: [Equation]
Revised: [Equation]

Original: D
Revised: D

Original: [Equation]
Revised: [Equation]

Original: (cid:44) − H −
Revised: (cid:44) − H −

Original: [Equation]
Revised: [Equation]

Original: D
Revised: D

Original: [Equation]
Revised: [Equation]

Original: Now let’s compute each of the terms
Revised: Now let’s compute each of the terms

Original: [Table]
Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Original: [Equation] 
Revised: [Equation] [Table] [Equation] [Table]

Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation][Table] [Equation] [Table]
Revised: [Equation]

Original: [Equation] [Table] [Equation] [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table][Equation] [Table] [Equation]
Revised: [Table]

Original: [Table]

Original: [Equation]
Revised: [Equation]

Original: To sum up, we have
Revised: To sum up, we have

Original: [Equation]
Revised: [Equation]

Original: Following the formula for inverse of block matrices,
Revised: Following the formula for inverse of block matrices,

Original: [Equation]
Revised: [Equation]

Original: denote
Revised: denote

Original: [Equation]
Revised: [Equation]

Original: Note that K is just a scalar for SOFT top- k operator, and is a ( k − 1) ×
Revised: Note that K is just a scalar for SOFT top- k operator, and is a ( k − 1) ×

Original: ( k − 1) matrix for sorted
Revised: ( k − 1) matrix for sorted

Original: SOFT top- k operator.
Revised: SOFT top- k operator.

Original: Therefore computing its inverse is not expensive.
Revised: Therefore computing its inverse is not expensive.

Original: Finally we have
Revised: Finally we have

Original: [Equation]
Revised: [Equation]

Original: And also
Revised: And also

Original: [Equation]
Revised: [Equation]

Original: The above derivation can actually be viewed as we explicitly force b m = 0 , i.e., no matter how Cchanges, b m does not change.
Revised: The above derivation can actually be viewed as we explicitly force b m = 0 , i.e., no matter how C changes, b m does not change.

Original: Therefore, we can treat db m dC = 0 n × m , and we get the equation in thetheorem.
Revised: Therefore, we can treat db m dC = 0 n × m , and we get the equation in the theorem.

Original: After we obtain dξ ∗ dC and dζ ∗ dC , we can now compute d Γ dC .
Revised: After we obtain dξ ∗ dC and dζ ∗ dC , we can now compute d Γ dC .

Original: [Equation]
Revised: [Equation]

Original: Finally, in the back-propagation step, we can compute the gradient of the loss L w.r.t.
Revised: Finally, in the back-propagation step, we can compute the gradient of the loss L w.r.t.

Original: C ,dL
Revised: C ,

Original: [Equation]
Revised: [Equation]

Original:  
Revised: +

Original: [Equation]
Revised: [Equation]

Original: 
Revised: dL

Original: [Equation]
Revised: [Equation]

Original: We summarize the above procedure for computing the gradient for sorted SOFT top- k operator in
Revised: We summarize the above procedure for computing the gradient for sorted SOFT top- k operator in

Original: Algorithm 3.
Revised: Algorithm 3.

Original: This naive implementation takes O ( n 2 k ) complexity, which is not efﬁcient.
Revised: This naive implementation takes O ( n 2 k ) complexity, which is not efﬁcient.

Original: Therefore,we modify the algorithm using the associative law of matrix multiplications, so that the complexity islowered to O ( nk ) .
Revised: Therefore, we modify the algorithm using the associative law of matrix multiplications, so that the complexity is lowered to O ( nk ) .

Original: We summarize the modiﬁed algorithm in Algorithm 4.
Revised: We summarize the modiﬁed algorithm in Algorithm 4.

Original: We also include the PyTorch implementation of the forward pass and backward pass as shownbelow.
Revised: We also include the PyTorch implementation of the forward pass and backward pass as shown below.

Original: The code is executed by creating an instance of TopK_custom , and the forward pass and thebackward pass is run similar to any other PyTorch model.
Revised: The code is executed by creating an instance of TopK_custom , and the forward pass and the backward pass is run similar to any other PyTorch model.

Original: Algorithm 3 Gradient for Sorted Top- k
Revised: Algorithm 3 Gradient for Sorted Top- k

Original: Require: C ∈ R n × ( k +1) , µ ∈ R n , ν ∈ R
Revised: Require: C ∈ R n × ( k +1) , µ ∈ R n , ν ∈ R

Original: k
Revised: k

Original: +1 , d L d Γ ∈ R n × ( k +1) , (cid:15)
Revised: +1 , d L d Γ ∈ R n × ( k +1) , (cid:15)

Original: Run forward pass to get Γ ¯ ν = ν
Revised: Run forward pass to get Γ ¯ ν = ν

Original: [: − 1] , ¯Γ =
Revised: [: − 1] , ¯Γ =

Original: Γ[: , : − 1]
Revised: Γ[: , : − 1]

Original: K ← diag(¯ ν ) − ¯Γ T (diag( µ )) − 1 ¯Γ # K ∈ R k × k
Revised: K ← diag(¯ ν ) − ¯Γ T (diag( µ )) − 1 ¯Γ # K ∈ R k × k

Original: H 1 ← (diag( µ )) − 1 + (diag( µ )) − 1 ¯Γ K − 1 ¯Γ T (diag( µ ))
Revised: H 1 ← (diag( µ )) − 1 + (diag( µ )) − 1 ¯Γ K − 1 ¯Γ T (diag( µ ))

Original: −# H 1 ∈ R n × n
Revised: −# H 1 ∈ R n × n

Original: H 2 ← − (diag( µ )) − 1 ¯Γ K −# H 2 ∈ R n × k
Revised: H 2 ← − (diag( µ )) − 1 ¯Γ K −# H 2 ∈ R n × k

Original: H 3 ← ( H 2) T # H 3 ∈ R k × n
Revised: H 3 ← ( H 2) T # H 3 ∈ R k × n

Original: H 4 ← K −# H 4 ∈ R k × k
Revised: H 4 ← K −# H 4 ∈ R k × k

Original: Pad H 2 to be [ n, k + 1] in the 
Revised: Pad H 2 to be [ n, k + 1] in the last column with value

Revised: Pad H 4 to be [ k, k + 1] in the last column with value

Original: [Table]
Revised: [Table]

Original: Require:
Revised: Require:

Original: C ∈ R N ×
Revised: C ∈ R N ×

Original: ( K +1) , µ ∈ R N , ν ∈ R K +1 ,
Revised: ( K +1) , µ ∈ R N , ν ∈ R K +1 ,

Original: d L d Γ ∈ R N × ( K +1) , (cid:15)
Revised: d L d Γ ∈ R N × ( K +1) , (cid:15)

Original: Run forward pass to get Γ ¯ ν = ν
Revised: Run forward pass to get Γ ¯ ν = ν

Original: [: − 1] , ¯Γ =
Revised: [: − 1] , ¯Γ =

Original: Γ[: , : − 1]
Revised: Γ[: , : − 1]

Original: K ← diag(¯ ν ) − ¯Γ T (diag( µ )) − 1 ¯Γ # K ∈ R K × K µ (cid:48) i
Revised: K ← diag(¯ ν ) − ¯Γ T (diag( µ )) − 1 ¯Γ # K ∈ R K × K µ (cid:48) i

Original: = µ − 

Revised: = µ −

Original: i
Revised: i

Original: L ← (diag( µ )) − 1 ¯Γ K −# L ∈ R N × K
Revised: L ← (diag( µ )) − 1 ¯Γ K −# L ∈ R N × K

Original: G 1 ← d L d Γ (cid:12) Γ # G 1 ∈ R N × K g 1 ← [ G 1] 1 K , g 2 ← [ G 1] T 1 N # g 1 ∈ R N , g 2 ∈ R K
Revised: G 1 ← d L d Γ (cid:12) Γ # G 1 ∈ R N × K g 1 ← [ G 1] 1 K , g 2 ← [ G 1] T 1 N # g 1 ∈ R N , g 2 ∈ R K

Original: G 21 ← ( g 1 (cid:12) µ (cid:48) )
Revised: G 21 ← ( g 1 (cid:12) µ (cid:48) )

Original: .expand _ dims (1) (cid:12) Γ # G 21 ∈ R N ×
Revised: .expand _ dims (1) (cid:12) Γ # G 21 ∈ R N ×

Original: ( K +1)
Revised: ( K +1)

Original: G 22 ← (( g 1) T L ¯Γ T (cid:12) µ (cid:48) )
Revised: G 22 ← (( g 1) T L ¯Γ T (cid:12) µ (cid:48) )

Original: .expand _ dims (1) (cid:12) Γ # G 22 ∈ R N ×
Revised: .expand _ dims (1) (cid:12) Γ # G 22 ∈ R N ×

Original: ( K +1)
Revised: ( K +1)

Original: G 23 ← − (( g 1) T L ) .pad
Revised: G 23 ← − (( g 1) T L ) .pad

Original: _ last _ entry (0) .expand _ dims (0) (cid:12) Γ # G 23 ∈ R N × ( K +1)
Revised: _ last _ entry (0) .expand _ dims (0) (cid:12) Γ # G 23 ∈ R N × ( K +1)

Original: G 2 =
Revised: G 2 =

Original: G 21 + G 22 + G# G 2 ∈ R N ×
Revised: G 21 + G 22 + G# G 2 ∈ R N ×

Original: ( K +1)
Revised: ( K +1)

Original: g 2 ← g 2[: − 1]
Revised: g 2 ← g 2[: − 1]

Original: G 31 ← − ( L ( g 2))
Revised: G 31 ← − ( L ( g 2))

Original: .expand _ dims (1) (cid:12) Γ # G 31 ∈ R N ×
Revised: .expand _ dims (1) (cid:12) Γ # G 31 ∈ R N ×

Original: ( K +1)
Revised: ( K +1)

Original: G 32 ← ( K − 1 ( g 2)) .pad
Revised: G 32 ← ( K − 1 ( g 2)) .pad

Original: _ last _ entry (0) .expand _ dims (0) (cid:12) Γ # G 32 ∈ R N ×
Revised: _ last _ entry (0) .expand _ dims (0) (cid:12) Γ # G 32 ∈ R N ×

Original: ( K +1)
Revised: ( K +1)

Original: G 3 =
Revised: G 3 =

Original: G 31 + G# G 3 ∈ R N ×
Revised: G 31 + G# G 3 ∈ R N ×

Original: ( K +1)
Revised: ( K +1)

Original: d L dC ← 1 (cid:15)( − G 1 + G 2 + G 3) def sinkhorn_forward (C, mu, nu , epsilon , max_iter ):bs , n, k_ = C. size ()v = torch . ones ([ bs , 1, k_ ])/( k_) G = torch . exp( − C/ epsilon )if torch .
Revised: d L dC ← 1 (cid:15)

Original: cuda . is_available ():v = v. cuda ()for i in range ( max_iter ):u = mu/(G ∗ v ).sum ( − 1, keepdim=True)v =
Revised: [Table]

Original: nu /(G ∗ u ).sum ( − 2, keepdim= True) Gamma = u ∗ G ∗ vreturn Gammadef sinkhorn_forward_stablized (C, mu, nu , epsilon , max_iter ):bs , n, k_ = C. size ()
Revised: [Table]

Original: k = k _ − 1f = torch .
Revised: _forward

Original: zeros ([ bs , n, 1])g = torch . zeros ([ bs , 1, k+1])if torch . cuda .
Revised: [Table]

Original: is_available ():f = f . cuda ()g = g. cuda ()epsilon_log_mu
Revised: stablized (C, mu, nu , epsilon ,

Original: = epsilon ∗ torch . log (mu)epsilon_log_nu = epsilon ∗ torch . log (nu)def min_epsilon_row (Z, epsilon ):return − epsilon ∗ torch .
Revised: [Table]

Original: logsumexp(( − Z)/ epsilon , − 1, keepdim= True) 619def min_epsilon_col (Z, epsilon ):return − epsilon ∗ torch . logsumexp(( − Z)/ epsilon , − 2, keepdim= True)for i in range ( max_iter ):f = min_epsilon_row (C − g, epsilon )+epsilon_log_mug = min_epsilon_col (C − f , epsilon )+epsilon_log_nu
Revised: ) else

Original: [Table]
Revised: [Table]

Original: [Figure] [Figure] [Figure] [Figure]
Revised: Gamma

Original: [Figure]

Original: [Figure]

Original: [Figure]
Revised: [Table]

Original: with
Revised: sinkhorn

Original: [Figure]
Revised: [Table]

Original: if
Revised: forward

Original: [Figure]
Revised: [Table]

Original: Gamma = sinkhorn_forward (C, mu, nu , epsilon , max_iter )if bool (torch . any(Gamma!=Gamma)):print ( ’Nan appeared in Gamma, re − computing ... ’)
Revised: stablized (

Original: Gamma =
Revised: [Table]

Original: sinkhorn_forward_stablized (C, mu, nu , epsilon, max_iter )else :
Revised: , mu, nu [Table] epsilon

Original: Gamma =
Revised: [Table]

Original: sinkhorn_forward_stablized (C, mu, nu , epsilon , max_iter )ctx .
Revised: ) ctx [Table]

Original: save_for_backward (mu, nu , Gamma)ctx . epsilon = epsilonreturn Gamma@staticmethoddef backward(ctx , grad_output_Gamma ):epsilon = ctx . epsilonmu, nu , Gamma =
Revised: [Table]

Original: ctx .
Revised: ctx .

Original: saved_tensors
Revised: [Table] = epsilon return Gamma @staticmethod [Table] backward(ctx ,

Original: # mu [1 , n, 1]
Revised: [Table] [Table] [Table] epsilon

Original: # nu [1 , 1, k+1]#Gamma [bs , n, k+1]with torch . no_grad ():grad_C = sinkhorn_backward (grad_output_Gamma , Gamma, mu, nu , epsilon )return grad_C , None , None , None , Noneclass TopK_custom( torch .nn. Module ):def __init__ ( self , k, epsilon =0.1 , max_iter = 200):super (TopK_custom1 , self ).
Revised: [Table]

Original: __init__ ()self .k = kself .

Original: epsilon = epsilonself .
Revised: epsilon [Table]

Original: anchors = torch . FloatTensor
Revised: Gamma

Original: [k − i for i in range (k+1)]).
Revised: [Table] # mu [ [Table]

Original: view ([1 ,1 , k+1])self . max_iter = max_iterif torch .
Revised: n

Original: cuda .
Revised: [Table]

Original: is_available ():self .
Revised: # nu [Table] [Table]

Original: anchors = self .
Revised: Gamma [Table]

Original: anchors . cuda ()def forward ( self , scores ):bs , n = scores . size ()scores = scores . view ([ bs , n, 1])#find the − inf value and replace it withthe minimum value except − infscores_ = scores .
Revised: with

Original: clone (). detach ()max_scores = torch .max( scores_ ). detach ()scores_ [ scores_==float(’ − inf ’)]
Revised: [Table]

Original: = float ( ’ inf ’)min_scores = torch .min( scores_ ). detach ()filled_value = min_scores − (max_scores − min_scores )mask =
Revised: no

Original: scores==float(’ − inf ’) scores = scores . masked_fill (mask , filled_value )
Revised: [Table]

Original: C = ( scores − self .
Revised: [Table] self .

Original: anchors)
Revised: anchors)

Original: ∗∗ [-
-] C = C / (C.max(). detach ())
Revised: ∗∗ [-
-]2 C = C / (C.max(). detach ())

Original: mu = torch .
Revised: mu = torch .

Original: ones ([1 , n, 1] , requires_grad=False )/nnu =
Revised: ones ([1 , n, 1] , requires_grad=False )/n nu =

Original: [1./ n for _ in range ( self .k)]nu.
Revised: [1./ n for _ in range ( self .k) nu.

Original: append ((n − self .k)/n)nu
Revised: append ((n − self .k)/n)

Original:  = torch .
Revised: nu = torch .

Original: FloatTensor (nu ).
Revised: FloatTensor (nu ).

Original: view ([1 , 1, self .k+1])if torch .
Revised: view ([1 , 1, self .k+1]) if torch .

Original: cuda .
Revised: cuda .

Original: is_available ():mu = mu.
Revised: is_available (): mu = mu.

Original: cuda ()nu = nu. cuda ()[-
-]
Revised: cuda () nu = nu. cuda ()[-
-]

Original: Gamma = TopKFunc.
Revised: Gamma = TopKFunc.

Original: apply (C, mu, nu , self .
Revised: apply (C, mu, nu , self .

Original: epsilon , self .
Revised: epsilon , self .

Original: max_iter )
Revised: max_iter )

Original: A =
Revised: A =

Original: Gamma[: ,: ,: self .k] ∗ nreturn A, None
Revised: Gamma[: ,: ,: self .k] ∗ n return A, None

Original: I Experiment Settings 
Revised: I Experiment Settings I.1 k NN

Original: I.1 k NN

Original: The settings of the neural networks, the training procedure, and the number of neighbors k , and thetuning procedures are similar to Grover et al.
Revised: The settings of the neural networks, the training procedure, and the number of neighbors k , and the tuning procedures are similar to Grover et al.

Original: The tuning o (cid:15) ranging from 10 − 6 to 10 − 2 .
Revised: The tuning o (cid:15) ranging from 10 − 6 to 10 − 2 .

Original: Other settings are shown in Table 3.
Revised: Other settings are shown in Table 3.

Original: [Table]
Revised: [Table]

Original: Note that f θ is a feature extraction neural network, so that model speciﬁed in the last row of Table 3does not contain the ﬁnal activation layer and the linear layer.
Revised: Note that f θ is a feature extraction neural network, so that model speciﬁed in the last row of Tabledoes not contain the ﬁnal activation layer and the linear layer.

Original: Baselines.
Revised: Baselines.

Original: In the baselines, the results of k NN, k NN+PCA, k NN+AE, k NN+NeuralSort is copiedfrom Grover et al.
Revised: In the baselines, the results of k NN, k NN+PCA, k NN+AE, k NN+NeuralSort is copied from Grover et al.

Original: The result of RelaxSubSample is copied from Xie and Ermon (2019).
Revised: The result of RelaxSubSample is copied from Xie and Ermon (2019).

Original: The implementation of k NN+Cuturi et al.
Revised: The implementation of k NN+Cuturi et al.

Original: (2019) is based on Grover et al.
Revised: (2019) is based on Grover et al.

Original: Speciﬁcally, theoutputs of the models in Cuturi et al.
Revised: Speciﬁcally, the outputs of the models in Cuturi et al.

Original: (2019) and Grover et al.
Revised: (2019) and Grover et al.

Original: (2019) are both doubly stochasticmatrices.
Revised: (2019) are both doubly stochastic matrices.

Original: So in the implementation of k NN+Cuturi et al.
Revised: So in the implementation of k NN+Cuturi et al.

Original: (2019), we adopt the algorithm in Groveret al.
Revised: (2019), we adopt the algorithm in Grover et al.

Original: (2019), except that we replace the module of computing the doubly stochastic matrix to be theone in Cuturi et al.
Revised: (2019), except that we replace the module of computing the doubly stochastic matrix to be the one in Cuturi et al.

Original: We extensively tuned k , (cid:15) and the learning rate, but cannot achieve abetter score for this experiment.
Revised: We extensively tuned k , (cid:15) and the learning rate, but cannot achieve a better score for this experiment.

Original: The baselines k NN+Softmax k times, k NN+pretrained CNN, and CE+CNN adopts the identicalneural networks as our model.
Revised: The baselines k NN+Softmax k times, k NN+pretrained CNN, and CE+CNN adopts the identical neural networks as our model.

Original: We remark that the scores reported in Grover et al.
Revised: We remark that the scores reported in Grover et al.

Original:  for CNN+CEare 99 .
Revised: (2019) for CNN+CE are 99 .

Original: 4% for MNIST and 95 .
Revised: 4% for MNIST and 95 .

Original: 1% for CIFAR-10.
Revised: 1% for CIFAR-10.

Original: However, our experiments using their code cannotreproduce the reported scores: and the scores are 99 . 0% and 90 .
Revised: However, our experiments using their code cannot reproduce the reported scores: and the scores are 99 . 0% and 90 .

Original: 9% , respectively.
Revised: 9% , respectively.

Original: Therefore, thereported score for MNIST is implemented by us, and the score for CIFAR-10 is copied from He et al.(2016).
Revised: Therefore, the reported score for MNIST is implemented by us, and the score for CIFAR-10 is copied from He et al.

Original: I.2 Beam Search
Revised: I.2 Beam Search

Original: Algorithm.
Revised: Algorithm.

Original: We now elaborate how to backtrack the predecessors E (1: t ) ,r for an embedding E ( t +1) ,(cid:96) ,and how to compute the likelihood L s ( E (1: t +1) ,(cid:96) ) , which we have omitted in Algorithm 2.
Revised: We now elaborate how to backtrack the predecessors E (1: t ) ,r for an embedding E ( t +1) ,(cid:96) , and how to compute the likelihood L s ( E (1: t +1) ,(cid:96) ) , which we have omitted in Algorithm 2.

Original: Speciﬁcally, in standard beam search algorithm, each selected token (cid:101) y
Revised: Speciﬁcally, in standard beam search algorithm, each selected token (cid:101) y

Original: ( t +1) ,(cid:96) is generated from a speciﬁcpredecessor, and thus the backtracking is straightforward.
Revised: ( t +1) ,(cid:96) is generated from a speciﬁc predecessor, and thus the backtracking is straightforward.

Original: In beam search with sorted SOFT top- koperator, however, each computed embedding E (1: t ) ,r is a weighted sum of the output from allpredecessors, so that it is not corresponding to one speciﬁc predecessor.
Revised: In beam search with sorted SOFT top- k operator, however, each computed embedding E (1: t ) ,r is a weighted sum of the output from all predecessors, so that it is not corresponding to one speciﬁc predecessor.

Original: To address this difﬁculty, weselect the predecessor for E ( t +1) ,(cid:96) with the largest weight, i.e.,
Revised: To address this difﬁculty, we select the predecessor for E ( t +1) ,(cid:96) with the largest weight, i.e.,

Original: [Equation]
Revised: [Equation]

Original: This is a good approximation because A ( t ) ,(cid:15) is a smoothed 0 - 1 tensor, i.e., for each (cid:96) , there is onlyone entry that is approximately 1 in A ( t ) ,(cid:15) : , : ,(cid:96) , while the others are approximately 0 .
Revised: This is a good approximation because A ( t ) ,(cid:15) is a smoothed 0 - 1 tensor, i.e., for each (cid:96) , there is only one entry that is approximately 1 in A ( t ) ,(cid:15) : , : ,(cid:96) , while the others are approximately 0 .

Original: The likelihood isthen computed as follows
Revised: The likelihood is then computed as follows

Original: [Equation]
Revised: [Equation]

Original: Implementation.
Revised: Implementation.

Original: The implemented model is identical to Bahdanauet al.
Revised: The implemented model is identical to Bahdanau

Original: Different from

Original: Bahdanau et al.

Original: (2014), here we also preprocess the data with byte pair encoding (Sennrich et al.,2015).

Original: We adopt beam size 5 , teacher forcing ratio ρ = 0 . 8 , and (cid:15) =
Revised: We adopt beam size 5 , teacher forcing ratio ρ = 0 . 8 , and (cid:15) =

Original: 10 − 1 .
Revised: 10 − 1 .

Original: The training procedure isas follows:
Revised: The training procedure is as follows:

Original: We ﬁrst pretrain the model with teacher forcing training procedure.
Revised: We ﬁrst pretrain the model with teacher forcing training procedure.

Original: The pretrainingprocedure has initial learning rate 1 , learning rate decay 0 .
Revised: The pretraining procedure has initial learning rate 1 , learning rate decay 0 .

Original: 1 starting from iteration 5 × 10 5 for every10 5 iterations.
Revised: 1 starting from iteration 5 × 10 5 for every 10 5 iterations.

Original: We pretrain it for 10 6 iterations in total.
Revised: We pretrain it for 10 6 iterations in total.

Original: We then train the model using the combinedtraining procedure for 10 5 iterations with learning rate 0 .
Revised: We then train the model using the combined training procedure for 10 5 iterations with learning rate 0 .

Original: I.3 Top- k Attention for Machine Translation
Revised: I.3 Top- k Attention for Machine Translation

Original: We apply SOFT top- k operator to yield sparse attention scores.
Revised: We apply SOFT top- k operator to yield sparse attention scores.

Original: Attention module is an integral partof various natural language processing tasks, allowing modeling of long-term and local dependencies.
Revised: Attention module is an integral part of various natural language processing tasks, allowing modeling of long-term and local dependencies.

Original: Speciﬁcally, given the vector representations of a source sequence s = [ s 1 , · · · , s N ] (cid:62) and targetsequence y =
Revised: Speciﬁcally, given the vector representations of a source sequence s = [ s 1 , · · · , s N ] (cid:62) and target sequence y =

Original: [ y 1 , · · · , y M ] (cid:62) , we compute the alignment score between s i and y j by a compatibilityfunction f
Revised: [ y 1 , · · · , y M ] (cid:62) , we compute the alignment score between s i and y j by a compatibility function f

Original: ( s i , y j ) , e.g., f ( s i , y j ) =
Revised: ( s i , y j ) , e.g., f ( s i , y j ) =

Original: s (cid:62) i y j , which measures the dependency between s i and y j .
Revised: s (cid:62) i y j , which measures the dependency between s i and y j .

Original: Asoftmax function then transforms the scores [ f ( s i , y j )]
Revised: A softmax function then transforms the scores [ f ( s i , y j )]

Original: Ni =1 to a sum-to-one weight vector w j for eachy j .
Revised: Ni =1 to a sum-to-one weight vector w j for each y j .

Original: The output o j of this attention module is a weighted sum of s i ’s, i.e., o j = w (cid:62) j s.
Revised: The output o j of this attention module is a weighted sum of s i ’s, i.e., o j = w (cid:62) j s.

Original: The attention module described above is called the soft attention, i.e., the attention scores w j of y j isnot sparse.
Revised: The attention module described above is called the soft attention, i.e., the attention scores w j of y j is not sparse.

Original: This may lead to redundancy of the attention (Zhu et al., 2018; Schlemper et al., 2019).
Revised: This may lead to redundancy of the attention (Zhu et al., 2018; Schlemper et al., 2019).

Original: Empirical results show that hard attention, i.e., enforcing sparsity structures in the score w j ’s, yieldsmore appealing performance (Shankar et al., 2018).
Revised: Empirical results show that hard attention, i.e., enforcing sparsity structures in the score w j ’s, yields more appealing performance (Shankar et al., 2018).

Original: Therefore, we propose to replace the softmaxoperation on [ f ( s i , y j )]
Revised: Therefore, we propose to replace the softmax operation on [ f ( s i , y j )]

Original: Ni =1 by the standard top- k operator to select the top- k elements.
Revised: Ni =1 by the standard top- k operator to select the top- k elements.

Original: In order for anend-to-end training, we further deploy SOFT top- k operator to substitute the standard top- k operator.
Revised: In order for an end-to-end training, we further deploy SOFT top- k operator to substitute the standard top- k operator.

Original: Given [ f ( s i , y j )]
Revised: Given [ f ( s i , y j )]

Original: Ni =1 , the output of SOFT top- k operator is denoted by A (cid:15)j , and the weight vector w jis now computed asw [Equation] = softmax([ f ( s 1 , y j ) , . . . , f ( s N ,
Revised: Ni =1 , the output of SOFT top- k operator is denoted by A (cid:15)j , and the weight vector w j is now computed as

Original: [Equation]

Original: j )] (cid:62)

Original: [Equation]
Revised: [Equation]

Original: Here log is the entrywise logarithm.
Revised: Here log is the entrywise logarithm.

Original: The output o j of the attention module is computed the sameo j = w (cid:62) j s .
Revised: The output o j of the attention module is computed the same o j = w (cid:62) j s .

Original: Such a SOFT top- k attention will promote the top- k elements in [ f ( s i , y j )]
Revised: Such a SOFT top- k attention will promote the top- k elements in [ f ( s i , y j )]

Original: Ni =1 to be evenlarger than the non-top- k elements, and eventually promote the attention of y j to focus on k tokens ins .
Revised: Ni =1 to be even larger than the non-top- k elements, and eventually promote the attention of y j to focus on k tokens in s .

Original: I.3.1 Experiment
Revised: I.3.1 Experiment

Original: We evaluate the proposed top- k attention on WMT2016 English-German dataset.
Revised: We evaluate the proposed top- k attention on WMT2016 English-German dataset.

Original: Our implementationand settings are based on Klein et al.
Revised: Our implementation and settings are based on Klein et al.

Original: For a fair comparison, we implement a standard softattention using the same settings as the baseline.
Revised: For a fair comparison, we implement a standard soft attention using the same settings as the baseline.

Original: The details are provided in Appendix I.
Revised: The details are provided in Appendix I.

Original: Results.
Revised: Results.

Original: As shown in Table 4, the proposed SOFT top- k attention training procedure achieves animprovement in BLEU score of approximately 0 .
Revised: As shown in Table 4, the proposed SOFT top- k attention training procedure achieves an improvement in BLEU score of approximately 0 .

Original: We visualize the top- k attention in Figure 6.
Revised: We visualize the top- k attention in Figure 6.

Original: Theattention matrix is sparse, and has a clear semantic meaning – “truck" corresponds to “Lastwagen",“blue" corresponds to “blauen", “standing" corresponds to “stehen", etc.
Revised: The attention matrix is sparse, and has a clear semantic meaning – “truck" corresponds to “Lastwagen", “blue" corresponds to “blauen", “standing" corresponds to “stehen", etc.

Original: [Table]
Revised: [Table]

Original: J Visualization of the Gradients
Revised: J Visualization of the Gradients

Original: [Figure] In this section we visualize the computed gradient using a toy[Figure] example mimicking the settings of k NN classiﬁcation.
Revised: In this section we visualize the computed gradient using a toy example mimicking the settings of k NN classiﬁcation.

Original: Speciﬁcally, we input 10 scores computed from 10 images, i.e., X ={ 0 , 1 , 2 , · · · , 9 } , into the SOFT top- k operator, and select the top3 elements.
Revised: Speciﬁcally, we input 10 scores computed from 10 images, i.e., X = { 0 , 1 , 2 , · · · , 9 } , into the SOFT top- k operator, and select the top3 elements.

Original: Denote the indices of the images with the same labelsas the query sample as I .
Revised: Denote the indices of the images with the same labels as the query sample as I .

Original: Similar to k NN classiﬁcation, we wantto maximize (cid:80) i ∈I
Revised: Similar to k NN classiﬁcation, we want to maximize (cid:80) i ∈I

Original: A (cid:15)i .
Revised: A (cid:15)i .

Original: (a) I = { 0 , 1 , 2 } .
Revised: (a) I = { 0 , 1 , 2 } .

Original: (b) I = { 2 , 3 , 4 } .
Revised: (b) I = { 2 , 3 , 4 } .

Original: We visualize the gradient on X with respect to this objectivefunction in Figure 7.
Revised: We visualize the gradient on X with respect to this objective function in Figure 7.

Original: In Figure 7(a), I is the same as the indicesof top- 3 scores.
Revised: In Figure 7(a), I is the same as the indices of top- 3 scores.

Original: In this case, the gradient will push the gap between the top- 3 scores and the restscores even further.
Revised: In this case, the gradient will push the gap between the top- 3 scores and the rest scores even further.

Original: In Figure 7(b), I is different from the indices of top- 3 scores.
Revised: In Figure 7(b), I is different from the indices of top- 3 scores.

Original: In this case, thescores corresponding to I are pushed to be smaller, while the others are pushed to be larger.
Revised: In this case, the scores corresponding to I are pushed to be smaller, while the others are pushed to be larger.