Learning distributed representations with efficient SoftMax normalization
Authors: Lorenzo Dall'Amico
Enrico Maria Belliardo
Venue: Accepted by TMLR ()

Abstract: Learning distributed representations, or embeddings, that encode the relational similarity patterns among objects is a relevant task in machine learning. A popular method to learn the embedding matrices $X, Y$ is optimizing a loss function of the term ${\rm SoftMax}(XY^T)$.  The complexity required to calculate this term, however, runs quadratically with the problem size, making it a computationally heavy solution. In this article, we propose a linear-time heuristic approximation to compute the normalization constants of ${\rm SoftMax}(XY^T)$ for embedding vectors with bounded norms. We show on some pre-trained embedding datasets that the proposed estimation method achieves higher or comparable accuracy with competing methods. From this result, we design an efficient and task-agnostic algorithm that learns the embeddings by optimizing the cross entropy between the softmax and a set of probability distributions given as inputs. The proposed algorithm is interpretable and easily adapted to arbitrary embedding problems. We consider a few use cases and observe similar or higher performances and a lower computational time than similar ``2Vec'' algorithms.

Published in Transactions on Machine Learning Research (05/2025)
Learning distributed representations with efficient SoftMax normalization
Lorenzo Dall' Amico ISI Foundation
lovenso.dallcmioo@iss...
Enrico Maria Belliardo ISI Foundation
enrico. purico.bellaardo@isst
Reviewed on OpenReview:
https: // openreview. net/ met/forwee forum? id- 9M4NKM2OPU
Abstract
Learning distributed representations, or embeddings, that encode the relational similarity patterns among objects is relevant task in machine learning. A popular method to learn the embedding matrices Y is optimizing a loss function of the term SoftMax(XY")). The complexity required to calculate this term, however, runs quadratically with the problem size, making it computationally heavy solution. In this article, we propose a lincar- time heuristic approximation to compute the normalization constants of (oltMaxXXY") for embedding vectors with bounded norms. We show on some pre-trained embedding datasots that the proposed estimation method achieves higher or comparable accuracy with competing methods. From this result, design an efficient and lask-agnostic algorithm that learns the embeddings by optimizing the cross entropy between the softmax and set of probability distributions given as inputs. The proposed algorithm is interpretable and easily adapted to arbitrary embedding problems. We consider few use cases and observe similar or higher performances and a lower computational time than similar "2Vec" algorithms.
1 Introduction
A foundational role of machine learning is providing complex objects with expressive vector representations that preserve some properties of the represented data. These vectors, also known as embeddings or distributed representations, enable otherwise ill-defined operations on the data, such as assessing their similarity (Bengio et al. 2013 Bellet el al., 2015). Relational patterns among the represented objects are a relevant example of a property that embedding vectors can efficiently encode. For instance, in language processing, distributed reprosentations capture the semantic similarity between words, leveraging the patterns in which they appear in text. One way of achieving this result is by training the softmax scores SoftMax(XYY)) where X, Y are two embedding matrices, possibly with X = Y. The softmax scores range between 0 and and can be used to promote similar representations for related objects. The optimization of the softmax scores is at the core of all "tyyce algorithms but is also key in attention- based transformers (Vaswani et al. 2017).
A known limitation of optimizing the softmax scores the computational complexity, which scales quadrat- ieally with the system size and is impractical for large datasets. As turnaround, Mikolov al. (2013)
considered two approximation strategies to train the embedding vectors, namely hierarchical softmax Good-
man
2001 Morin & Bengio 2005). and negative sampling. Thanks to its efficiency, negative sampling is
largely adopted and several works proposed analyses or extensions {Mu et al. 2019 Rawat et al. 2019 Shan
el al.
2018 Bamler & Mandi 2020), but also highlighted the weaknesses of this approach (Landgral &
Bellay 2017 Chen et al. 2018 Qin et al. 2016 Mimno Thompson 2017). Both efficient implementations of 2Vec-type algorithms, however, consider de facto modified loss functions with respect the originally proposed one, thus circumventing the computational bottlenock. A large body of literature, instead, focused and is still actively exploring methods to approximate and train softmax efficiently. These works include

---

Published in Transactions on Machine Learning Research (05/2025)
importance sampling methods (Bengio & Senécal
2003
Blanc & Rendle 2018
Rawat et al.
Baharav Choromanski et al.
2019
et al. 2024), random feature methods based on the kernel trick [Rahimi & Recht, 2007
2020
P Peng et al. 2021), low-rank approximations (Drineas et al.) 2005 Shim et al. 2017
Xiong et al. 2021 1
and local hashing techniques (Mussmann et al. 2017 Zahoor et al. 2020 Beltagy et al. 2020).
In this work, we provide a closed formula to approximate the softmax normalization constants in linear time, with theoretical and empirical results supporting our method. We evaluate the accuracy of the proposed approximation on pre-trained embeddings and compare it with several related methods. Building on this result, design an efficient embedding algorithm in the 2Vec spirit, training loss function containing the term log[SoftMax(XXY) Given their good scalability and simplicity of design, the 2Vec algorithms proved extremely useful for machine learning users and have been applied to various contexts. However, we recall that 2Vec algorithms do not optimize the proposed loss function in their most popular implementation. Our algorithm, instead, optimizes the original loss function efficiently and leverages the existing theoretical works that analyzed it, like (Jaffe et al. 2020). We note that our proposod approximation of the softmax score normalization is sufficient. to efficiently compute lgg(SoftMax(2 (XXXT) because this term is composed of two parts: the numerator XY that is low rank and can consequently be efficiently optimized; the denominator involving the softmax normalization constants that is the computational bottleneck. Our estimation formula allows us to compute it in ((n) operations, and thus to efficiently optimize the embedding cost function. We showcase a few practical applications evidencing that our algorithm is competitive or outperforms comparable methods in speed and performance in n, and do not optimize the softmax function as we do. We stress that we do not claim sampling approaches are bad per se, but rather, we investigate how to optimize cost function containing the softmax scores efficiently. A Python implementation of our algorithm is available at gtthabb.com//rrozzoaaaammi///o/
! We recall that all these benchmark algorithms have a linear complexity
2 Main result
This section describes how to efficiently approximate the normalization constants of SoftMax(XY3) where X E Raxd, Y E Rmnd are two embodding matrices, with a potentially different number of rows, but with the same number of columns. We denote with fæ.)i=,..... and {ya)a=I....m the sets of vectors contained in the rows of X and Y, respectively. The i-th softmax score normalization constant reads
2, vo
(1)
In the remainder, we let m = On). implying that computing all the Z,'s requires O(dnm) = Odn??) operations, which is prohibitively expensive for large datasets. This condition determines the relevance of our problem setting; m < n, one can efficiently compute the normalization constants. We further assume d < n, a necessary condition to estimate all Z; in linear time. We remark that other works, such as Baharav] et al. (2024), have considered a complementary view to ours, by estimating the softmax normalization in sub-linear time in d. This setting is relevant for high-dimensional embedding vectors with d = (n).
2.1 Linear-time softmax normalization
We consider vector x; Rd fixed (corresponding to the i-th of X) and treat (yI...., ym) as random independent vectors drawn from an unknown distribution. This allows us to study Z, as random variable and describe its concentration properties. Note that we are not assuming the embedding dimensions to be independent, thus preventing the representation of complex similarity patterns. Such relations are captured by the undorlying unknown probability distribution on which we make no assumption besides requiring embedding vectors to have bounded norms. We derive our approximation show the concentration of Z//m around its expectation for which we obtain an integral form, depending on the unknown distribution fi of the scalar product a! ya- In the sccond, we introducc variational approximation by substituting the unknown distribution of the scalar product fa with Gaussian distribution. In the third, we introduce a clustering-baset method applied to the rows of Y to improve the estimation accuracy.
three steps. In the first, we
1 All codes are run on Dell Inspiron laptop with 16 GB of RAM and with a processor 11th Gen Intel Core 17-11390H @ 3.40GHx 8.
2

---

Published in Transactions on Machine Learning Research (05/2025)
Concentration properties of the normalization constants
We succinctly describe the main concentration result of the normalization constants Z, which we formally emunciate and prove in Appendix A We denote with Si.m the unknown empirical distribution of the scalar
product of ya and suppose its support is in I-h, h| for h = (1)) Assuming the convergence in distribution of Jim to Sú, with high probability
Z == +00 m7
S"" di ee s:())
lim
(2)
Our result holds unchanged also the case X Y. We remark that we do not require any assumptions on the embedding vectors' distribution (besides having finite norms). nor that the embedding vectors are identically distributed. For more accurate enunciation, refer the reader to Theorem A.1 The constant h controls the concentration speed of Zi/m around its mean: smaller h implies faster convergence, The assumption h = (1) guarantees good concentration properties, and it can be easily enforced by considering embedding vectors with bounded norms. Equation 2 is obtained by combining Theorem A.1 with Theorem A.3 in which we show that the E/Zl//m converges to the integral form on the right hand-side. This is generally not true under the loose assumption of convergence in distribution we fornulated.
Gaussian approximation of the scalar product distribution
Equation 2 gives tractable expression of the random variable Z,, but it cannot be solved without additional hypotheses. We thus proceed by introducing variational approximation of fi. denoted with ft. Thanks to Theorem A.3 the goodness of this approximation only needs to hold in distribution sense, i... the cumulative density functions of fi and S; should be close. We choose to write Si as a Gaussian distribution. Recalling that f; is the distribution of scalar product i.e., a sum of random variables we expect that for d large enough, it is well approximated by the normal distribution. As shown in Section [.2] the empirical evidence confirms the goodness of this approximation real data. We denote with μ., Z the mean and covariance of y and with E[|] the expectation over We obtain Elx? y) x! μ and W(xr y] = xx Sax and write
lim 7: = dt e f R m-0000 m
-f"
[ c'N of u. (m.)) ) exp
x 1 + a. Sax }
Note that this approximation does not require the embedding vectors to be drawn from a multivariate Gaussian distribution. As a remark, we move from an integral on [-h, h] to one over the real axis, but the contribution from the tails is negligible since the integrand goes to zero at least as fast as e144, The considerable advantage of Equation 4 that, given l and Ω2 the normalization constant Z, computed Odd) operations, independently of m. This allows us estimate all Z, 's in O(n) operations. Since μe and Ω are both estimated in O(n) operations, this formula allows the lincar-tima estimation of the normalization constants. We note the quadratic complexity in d, but, since we are working under the assumption d < this does not hamper the computational efficiency of the proposed method.
Multivariate Gaussian approximation of the scalar product distribution
To obtain Equation (3,, we assumed all embedding vectors be well described by the same distribution. In practice, it is common that. the represented objects form clusters in the embedded space that mirror affinity groups. We can leverage these clusters to improve the estimation accuracy by clustering the embedding vectors to reduce the within-group embedding variance. We subdivide the set V of all elements in non- overlapping subsets Va V. The normalization constant then reads
m
7
DeV.
Z,a
3

---

Published in Transactions on Machine Learning Research (05/2025)
If Va (m) for all a, then Zia respects the same concentration properties of Equation 2
This
assumption is also necessary to keep a low computational complexity. We update Equation 3 as follows
Z
R
No ea No t0
(4)
m
where Ha. Sa are the mean vector and covariance matrix of the vectors y in class a, and Ta = Val/m.
VVl. m Na
To
VV
a;; Sh Val
Σ (x pa) pa)² IEV.
LEV.
We choose the partition of V to minimize the variance x Saax, within each class, We define the total variance V and show the following relation with the k-means objective (MacQueen 1967
Σ
Σ a=l LEV.
Σ
af (y. (M))- s |1
Σ 1l Ma a=
aEV.
Using k-means obtain the clusters Va=l....s we minimize the upper bound of the total variance, hence improving the estimation performance. The complexity of Lloyd algorithm (Lloyd 1982) to perform k-means clustering scales as (ndd). As such, with Equation 4 all Z,8 can be computed in O(nxd?) operations.
2.2 Empirical evaluation
We consider 6 datasets taken from the NLPL. word embeddings repository ? Kutuzov et al. 2017 repre- senting word embeddings obtained with different algorithms and trained different corpora:
0. British National Corpus; Continuous Skip-Gram, 163.473, 300;
English Wikipedia Dump 02/ 22/2017; Global Vectors, n = 273.930, 300;
16. Gigaword 5th Edition; fast Skipgram, 292.967, d = 300;
30. Ancient Greek CoNLL17 corpus; Word2Vec Continuous Skip-gram, n 45742, d 100;
187. Taiga corpus; Text Continuous Bag-of-Words, 192.415, 300;
224. Ukrainian CoNLL17 corpus; Continuous Bag-of-Word, 99.884, d = 200.
The number reported in the list above corresponds to the ID used in the repository. For each dataset, we (1) rescale the embedding vectors that their average norm equals 1; (2) sample 1000 random indices; (3) compute the corresponding exact and the estimated Z, values for different approximation orders k. We then repeat the same procedure by imposing that all embedding vectors have unitary norms. Figure shhows the results of this procedure. The first column compares the empirical distribution cumulative density functions (CDF) of f; and Ji obtained for = 5. Here, is a randomly selected node and = Y, and we use the rescaled version of the embedding matrix. The third column shows the same result for the normalized embedding matrices. The plots confirm that in all cases, the multivariate Gaussian approximation introduced achieves high accuracy in estimating fi. The second and fourth columns of Figure | evaluate the accuracy of the estimation method, by showing the cumulative density function of the error ween the exact and the estimated values of Z,. As expected, the precision of our method increases with k.
Performance comparison
We compare the accuracy of our estimation method with several competing approaches. For all methods, we set the hyper-parameters to have a comparable execution time with our method.
Sampling. simple way of estimating the normalization constant 2, efficiently is sampling g < m indices of the sum appcaring in Equation I Small values of improve the algorithm's speed but lead
2The datasets can be found http://eeetees.npp.mm/mmppponnn and shared under the CC BY 4.0 license.
4

---

Published in Transactions on Machine Learning Research (05/2025)
PBI
Nysrom Topk
hystrom
Emarical
Adgli solmma
muGasalen
Emp
Adapt sotmmx
Tip-c
1.0
1.0
1.0
1.0
CDL
0.5
0.5
0.0
0.00
0.25
0.50
0.75
1.00
0.0
-5
-3
-2
0.0
0.00
0.25
0.50
0.75
1.00
0.0
-5
-4
-3
-2
-1
1.0
1.0
1.0
1.0
CDF
0.5
0.5
0.5
0.0
0.0
0.0
-4
-3
-2
1.0
1.0
1.0
1.0
CDF
0.5
0.5
0.5
0.5
0.0
0.0
0.00
1.0
1.0
1.0
1.0
CD1
0.5
0.5
0.5
0.5
0.0
0.0
1.0
1.0
1.0
1.0
CDF
0.5
0.5
0.0
0.00
1.0
1.0
1.0
CDF
0.5
0.5
0.5
0.0
0.0
-2
0.0
1.0
-5
-3
ya
lognole)
x y
logiole)
Figure 1: Evaluation of the softmax normalization constants approximation. Each row corresponds to one of the six embeddings listed in Section 2.2 which appear in the same order. The first two columns consider the rescaled embedding matrices in which the average norms are equal to onc. In the last two, all embedding vectors are normalized to one. The first and third columns show the cumulative density function (CDF) of the scalar product distribution f. for an arbitrary index The black dashed-dotted curve is the estimated multivariate Gaussian approximation fç for 5. The second and fourth columns show the CDF of the error = 1Z Zeal/n n in estimating the softmax score normalization constants. The solid lines refer to our proposed method for three values of (color-coded). The other lines refer to the methods described in Section 2.2] "Adapt. softmax", pink dashed-dotted line (Blanc & Rendle 2018): "RF", red loosely-dashed linc RRahim & Rocht 2007); "PRF", orange dashed line (Choromanski et al.) 2020); "Nyström", gray dottod line (Drineas el al. 2005 "Top-k", green dashed double dotted line (Mussmann al. 2017
to higher variance estimates. In (Bengio & Senécal 2003). the authors showed that importance sampling reduces the estimation variance and proposed an unbiased estimator that, however, requires the exact softmax distribution. The authors of [lanc & Rendle 2018) proposed adaptive softmaz in which the sampling probability of index a proportional to + (xx ) Leveraging a tree structure and the kernel trick, they further showed how compute the softmax normalization constants in Od2nlog n) operations. This algorithm's performance is denoted by "Adapt. softmax" in Figure)
5

---

Published in Transactions on Machine Learning Research (05/2025)
Random features. The terms eal ye are closely related to the Gaussian kernel e e-Ilm:-valll and the random feature methods exploit the kernel trick to linearize the softmax score matrix. One defines a random matrix W E Rxx with entries distributed according to (O,,, La), and a mapping, ow R 4 R2D so that (w((ow(a)) @w =e الليي This approximation allows one to compute Z, efficiently:
m
o yo
10,1
m
I (a:)) @ow (w (ya))
E (o:()) mw
where mw = Ee li 12
ow (ya) needs to be computed only once. The complexity of these methods runs as (n.)). In Figure 1 the line denoted as "RF" considers the mapping & using random features as proposed in [Rahimi & Recht, 2007 while "PRF" the line obtained for the positive random features proposed in CChoromankii al. 2020) Both methods are implemented for D 1000.
La=1 2
Low-rank approximations. Letting Qia -llle.-vall" the softmax normalization constant. reads Z, (Q1m); If Q can be written as the product of two low-rank matrices A, B, then Z, AAB"1, 1,, can be computed efficiently by multiplying from right to left, without materializing Q. One way of performing such a low-rank approximation is to use the Nyström method (Drineas et al. 2005), which we consider for X = Y, leading to a symmetric Q. We sample g indices and compute the corresponding rows and columns of Q. Without loss generality, we let the sampled indices be the first and write
Q-(aum
Qss (oss Que
Que Quu
=gmm
(as Qus
One QuaQu OT
The matrix Quu E (((--))((-sp) represents the entries of the unsampled indices and is approximated usr where i denotes the Moore-Penross pesudo-inverse. The matrix Qt Rax9 small, while Qus (n Assuming g & n, this method estimates all Z₁ in O(ndg) operations. The curve "Nyström" in Figure reports its performance for g 50.
with Qt QTe
dominated by the large values of at Yas because of the exponential function. One can leverage the large literature on mazimum inner product search to find the k closest embedding vectors to æ; and use them to define an estimator of Z; with lower variance. For instance, following (Mussmann et al. 2017), we let S, be set with the indices corresponding to the k closest vectors to and T; set of randomly sampled indices from V S. The estimator of Z; reads
Top-k approximation. The sum in Equation
Zi R3
Σ aES;
71
SS |
" v
Several methods exist to efficiently estimate the set S;, and all the Z,8 can be computed in O(nd + g)) operations. The curve "Top-k" in Figure shows the result of this method for k 25.
The results show our method systematically outperforms the competing ones, with one exception which the Nyström-based approach provides better results, even if on other datasets it outperformed by other methods. Most importantly, we observe high performance for 1. Here, the clustering step can be omitted and is of particular interest, since the isator is obtained from simple matrix operations.
3 EDRep: an algorithm for efficient distributed representations
Building the results of Section 2 we describe an efficient algorithm which name EDRep to obtain distributed representations in the 2Vec spirit. These algorithms are still extremely popular among machine learning users thanks to their efficiency, scalability, and flexibility. They build on Word2Vec (Mikolov el al., 2013), an algorithm designed for word embeddings, and were adapted to a variety of domains, including
graphs (Perozzi et al. 22014; Grover & Leskovee, [016 Gao et al. 2019 Rozemberezki et al., 2019 Nickel &
Kiela 2017 Narayanan al.) (017). time (Kazemi al.] 2019).
temporal contact sequences (Goyal et al.
2020
Rahman al.] 2018 Nguyen et al. 2018 Sato et. al. 2021 Torricelli et. al. 2020 biological entities

---

Published in Transactions on Machine Learning Research (05/2025)
Ng 2017 tweets DDiinrra et al. 2016 and higher order interactions Billings et al, 2019) among others. Word2Vec trains the embedding in an unsupervised fashion by letting words commonly appearing in the same context have a similar representation. The generalizations of Word2Vec to other contexts roquire the "translation" of the input dataset into a text on which Word2Vec is then applied. For instance, Mode2Vec, a lext is created by performing random walks on a graph.
[uu et al. 2019
In our formulation, we take different perspective and consider a probability matrix P the input of our problem. This matrix encodes relational patterns between object pairs and is well-suited for datasets describing affinity measures among the embedded objects. We formulate general embedding problem easily customized to an arbitrary setting. The optimal design of the probability matrix P is problem-dependent and beyond the scope of this article. However, we show in a few applications that simple choices lead to comparable results with competing or better 2Vec algorithms but with lower computation time.
3.1 Problem formulation
We consider set V of n items to be embedded. The relational patterns among the objects in V are encoded by an affinity probability matrix P E R=xx Our goal is to learn an embedding matrix E Rd that preserves the relational patterns encoded the matrix P. do so, we adopt a variational approach and minimize the cross entropy between the rows of P and of SoftMax(XXX) over the embedding vectors. This allows to learn the best parametric approximation X to fit the imput matrix P. More formally, we define the embedding matrix X as the solution to the following optimization problem:
X = arg min Σ
Σ JEV
Poy log SoftMax(YY"))) 16)
+
y
Yj
(6)
Y Y'llaxd
n
iEV
JEV
cross - entropy
regularization
where Uaxd denotes the set of all matrices of size n x d having in their rows unitary vectors. The loss function includes a regularization term that promoles embedding matrices with a centered mean. We empir- ically observed that this term improves the embedding quality in practical applications. The computational bottleneck of this optimization problem lies in the calculation of the softmax score normalization constants. We denote with the number of non-zero entries of P, with 1, the all-ones vector of size n, and with tr(-) the trace operator. Then, the optimization problem of Equation @ can be reformulated follows:
X = arg min
in (r
+
log Z,) + 1
(7)
Yeulxdd
SEV
O(Ed)
((nd)
(ddr?)
where the values below the brackets indicate the computational cost. required for each element. The derivation
of this expression is reported in Appendix B In many relevant settings, P is sparse matrix, thus E < n² The calculation of the softmax score normalization constants, instead, requires Oddn² operations regardless of E and is the computational bottleneck. We can thus adopt the approximation introduced in Equation (4) to efficiently optimize this cost. function and to define an efficient embedding algorithm.
Remark 3.1. In Equation (6) and in the remainder, we focus on square matrices P, in which rows and columns are defined over the same set V. The optimization problem in Equation (6) can however be general- ized to an asymmelric scenario in which the entries Pia are defined for E and a W, thus inwoking the term SoftMax(XY")), for X Y. We provide Python codes to obtain the embedding also in this setting.
Let us now detail the main steps needed to translate the result of Equation into practical algorithm produce efficient distributed representations.
3.2 Optimization strategy
We obtain the embedding matrix X by optimizing the problem formulated in Equation 6 with stochastic gradient descent. We substitute the approximated values of Z; introduced in Equation and we let

---

Published in Transactions on Machine Learning Research (05/2025)
Algorithm EDRep
Input: P E R probability matrix encoding similaritios; d, embodding dimension; l E fl,.... x} node
label vector; No, learning rate; n_epochs, number of training epochs
Output: Xaxd, embedding matrix
X initialize the embedding matrix with random unitary vectors
") 7lu initial learning rate #a=1..... as per Equation for 5 S a_epochs do
5
Ma=l....* -o=1...... & update the parameters as per Equation 5
{g. + gradient matrix as Equation S
for iiin do
g - 9: [g] æ, remove the parallel component
9. g |lg6 normalize
2; - v 1 x.
ng": gradient descent step
end for
1) 1) epochs linear update of the learning rate end for
M E Rexd have the (a values rows. We further define 2 E Rax* as
Zia = Ta exp { x No + 2 Pati
With this notation Z:/n (21,); The iE V and fl,...,af d) gradient component reads
gig = - ig
1n1 +
ZM +
Zia (Xa())))
(8)
((Ed)
((nd)
(Ound2))
where the values underneath the brackets indicate the computational complexity required to compute each addend of the gradient matrix. The derivation of Equation is reported in Appendix C To keep the normalization, we first compute g' removing the component parallel to æ;; then we normalize and obtain g finally update the embedding follows for S 1) S x.t VI- - r7 x; ng.. that implies |æ:-o||| = 1. Note that Algorithm requires the labeling vector f as input but it is generally unknown. A workaround consists in running EDRep for 1 for which l = 1,; then run k-class clustering k-meuns and rerun EDRep algorithm for the so-obtained vector L.
3.3 Computational complexity
To determine the complexity of Algorithm let us focus on its computationally heaviest steps:
The calculation of is 1. This is obtained in O(nRd) operations with k-means algorithm.
The parameters updale as per Equation 5 This step is performed in O(nnd²) operations.
3. The gradient calculation per Equation
This is obtained in O(Ed + nxd?) operations
indicated by the brackets in Equation
The gradient calculation is thus the most expensive operation. Our approximation reduces the complexity required to compute the "Z part" of the gradient from O(dn?) to O(nnd), with K, & n. The most expensive term thus requires O(Ed) operations. This complexity is prohibitive for dense matrices, but in typical settings P is sparse and the product can be performed efficiently. Nonetheless, even for large values of F,, if P be written the product (or of products) of sparse matrices, PX can still be
8

---

Published in Transactions on Machine Learning Research (05/2025)
Exact
EDRep
comp. time (s)
comp. time (s)
-3.0
(A)
(B)
(C)
40
1o
-3.5
2-
20
-4.0
1
y yax²
y x
-4.5
5000
10000
5000
10000
50
100
150
epoch
n
Figure 2: Comparison with exact gradient calculation. Panels A, B: computation time of the opti- mization of Equation with gradient descent as function of the size n. Panel A refers the exact gradient, panel B is Algorithm The blue dots the mean obtained over 10 realizations, the shadow line has the width of one standard deviation. Panel C: logarithm of embedding error C. = + x|x,, -X,X;") || between the true and the estimated embedding matrices at training epoch t. In this experiment, n = 3000. In color code and marker style, we report the results for different values of k. The embedding algorithms are run with the same initial condition and parameters: Yoo = 0.7, d = 32, Repocius 25 (for the first two plots).
computed elficiently. In fact, let P Pm Pm- ... Pi for some positive m, then PX can be obtained without materializing P, taking the products from right to left:
PX (Pm Pm-1 Pm--1..... ..X))
thus speeding up the computational bottleneck our algorithm. In our implementation, we explicitly consider this representation of as an input. When a non-factorized dense matrix P provided, one could envision adopting method such the one presented in (Le Magoarou & Gribonval 2016) to approximate dense matrix P with the product of sparse matrices to speed up the algorithm.
3.4 Comparison with exact gradient computation
We compare the EDRep algorithm described in the previous section with its analogous counterpart in which the gradient of Equation 6 is computed analytically in O(?d)) operations. We considered P the row- normalized random matrix in which the (ij) entry is set to 1 with probability proportional to 0,0; and 0 follows a negative binomial distribution with parameters N = 3, generate heterogencity in the P structure while being capable of controlling the size n. Figure D shows the computational time corresponding to the exact gradient calculation, while panel B reports the same result for EDRep. Let X, be the EDRep embedding at epoch t and Y, be the corresponding one of the full gradient calculation, we define C. = - fI, quantifying the deviation between the two embedding methods. Figure _ shows the behavior of C, for different values, evidencing only slight disagreements between the exact and the approximated embeddings that, as expected, decrease with
0.3. This choice allows us to
4
Use cases
We consider few use cases of algorithm to test it and showcase its flexibility in practical settings. To perform these tests, we must specify the set P and adopt simple strategies to define it. We show that, even for our simple choices, the EDRep approach achieves competitive (sometimes superior) results in terms of performance with competing 2Vec algorithms, with a (much) lower computational time 3 We would like to underline that the sampling probabilities choice a hard and problem-dependent task and optimally addressing it is beyond the scope of this article. Our aim not to develop slate-of-the-art algorithms for
3It. should be noted that performances typically increase with training time The parameter choice of the 2Vec algorithms is such that the competing algorithms comparable on one two measures that the other be evenly compared.

---

Published in Transactions on Machine Learning Research (05/2025)
(A)
(C)
(E)
0.9
6
non interacting interacting
DyAME EDRipp
infected
200
NMI
100
EDRep
0.3
Node2Vec
2
0.0
2
5
-0.5
0.0
0.5
1.0
100
360
400
α
x/ X (EDRep)
time
(B)
(D)
(F)
Prediction EDRep
1.5
t (sec)
PDF
100
20-
2
5 0.0
10
100
150
200
α
x/x, (Gene2Vec)
Prediction DyANE
Figure 3: EDRep cases. First column: community detection. Panel A: normalized mutual information (NMI) as function of the problem hardness a (see Appendix D for DCSBM graph. We consider graphs with n = 30 000 nodes, expected average degree = 10, and g = 4 communities. The green circles refer to the EDRep algorithm with d = 32, x = 1 and w = 3, the brown diamonds are DeepWalk with d = 32, while the blue narrow diamonds are the spectral clustering algorithm of (Dall' Amico et al. 2021). Panel B: corresponding computational time in scconds. Panels A and B: averages are over 10 samples and the error bar width equals the standard deviation. The two panels share the same legend. Second column: gene embedding. Scalar product between gene embedding representation for non-interacting (blue histogram) and interacting (purple histogram) gene pairs. Panel C: embedding obtained with EDRep; Panel D embedding obtained with the Gene2Vec algorithm. Panels C and D share the same legend. Third column: dynamic aware node embeddding. Panel E: number of infected individuals of a SIR process on a proximity network (black solid line) and reconstructed values by DyANE (orange dashed line) and using the EDRep embedding (blue dashed-dotted line), with 200 for both methods. The shaded lines are the standard deviations over 50 random trials of the process. Panel F: with reference Panel E, this is the scatter plot between the predicted number of infected per time-stamp by the two strategies.
specific problems but to show that with simple choices we can adapt our algorithm to compete with the closest competing mothods in terms of speed and accuracy. For further implementation details regarding the next section, we refer the reader to Appendix D
4.1 Community detection
Graphs are mathematical objects that model complex relations between pairs of items. They are formed by set of n modes V and set of edges & connecting node pairs (Newman 2003). Graphs can be represented with the adjacency matriz A _ that Aú 1 if (ij) & and equals zero otherwise. A relevant problem in graph lcarning is community detection, the task of determining a non-overlapping node partition, unveiling more densely connected groups of nodes (Fortunato & Hrie 2016). common way of proceeding see e.g- (Vn Luxburg 2007) - is to create a node embedding encoding the community structure and then clustering the nodes in the embedded space. Following this strategy, we adopt the EDRep to produce node embedding with P being P w Σ (L,w)', where Lrw is the row-normalized adjacency matriz. The entry Pg is the limiting probability that random walker
goes from node to j in one fewer steps.
We evaluate our algorithm on synthetic graphs generated from the degree-correctedr stochastic block model (DCSBM) [Krrerr & Newman 2011., capable of creating graphs with a community structure and an arbitrary degree distribution R The inference accuracy is expressed with the normalized mutual information wetw
4The degree of node is its number connections.
10

---

Published in Transactions on Machine Learning Research (05/2025)
the inferred and the ground truth partition. This score ranges between (random assignment) and 1 (perfect assignment). Figure 3 shows the NMT for different values of a, function of the generative model parameters, controlling for the hardness of the reconstruction problem [ The results are compared against other two algorithms that were alternatively deployed to obtain the embedding: the spectral method of (Dall Amico al.) 2021) that was shown to be nearly Bayes-oplimal for this task and DeepWalk (Perozzi ct al. 2014 The communitics are obtained from the embeddings using k-class k-means clustering.
The results show that the EDRep-based algorithm performs almost well as the optimal algorithm of Dall' Amico et al., 2021 and slight mismatch is only observed for approaching 1. This is particularly challenging setting in which only a few algorithms can retrieve the community structure. Compared to the DeepWalk approach, our method generally yields better results for all a. Figure 3B further compares the computation times, giving the EDRep approach a decisive advantage with respect to Deepwalk. The main advantage with respect to the spectral algorithm, instead, the algorithm's 's computational complexity. For graph with q communitics, the considered spectral clustering algorithm runs in (nqq)) operations, while the complexity of EDRep is independent of q-
4.2 Gene embeddings
In (Du et al. 2019), the authors develop an algorithm to embed DNA genes from a list of pairs whose co-cxpression exceeds a threshold value. The dataset comprises 8832 genes and 263016 gene pairs and a list of gene pairs with a binary label indicating whether or not that corresponds to an interacting pair. The Gene2Vec algorithm of (Du al. 2019) builds on Word2Vec to obtain meaningful gene vector representations based on their co-expression and uses it predict pairs of interacting genes.
We obtain the EDRep embedding using the normmoral gene co-occurrence matrix The Gene2Vec embedding is generated with the code provided by the authors with default parameters. Both embeddings have dimension d = 200. The computation time of EDRep is approximately 6 seconds against the 12 seconds needed for Gene2Vec. We then train a logistic regression classifier on the embedding cosine similarities with the 70% of the labeled data and test on the remaining 30% of the data. Our model achieves an accuracy of 92% and outperforms Gene2Vec which has an accuracy of 84%. Figures B D) show the histogram of the cosine similarities between interacting and non-interacting groups that visually explains the performance gap.
choice of P.
4.3 Causality aware temporal graph embeddings
In (Sato et al. 2021 the authors describe a method to embed temporal networks while preserving the role of time defining causality. Temporal networks are represented as a sequence of temporal edges (i,j, I), denoting an interaction between and j at time t. The method relies on the embedding of a supra-adjaceney matrix, Asupra in DDD where D -t= [i]] and V. is the set of active nodes time Here, each node corresponds to pair "node-time" in the original temporal graph. The supra-adjacency matrix is the adjacency matrix of weighted directed acyelic graph, accounting for time-driven causality. In (Sato et. al. 2021) the authors use DeepWalk to obtain embedding from Asupra and use it to reconstruct the states of a partially observed dynamical process taking place on the temporal graph (such as an epidemic spreading) from few observations. Following the same procedure of (Sato et al.] 2021), we obtain Figure BEEFF), in which we compare the reconstruction of an epidemic spreading obtained using DeepWalk against EDRep, with P being the row-normalization of Asupra- The results are barely distinguishable, but EDRep more than 5 times faster than the competing approach.
Detection is theoretically kessible if and only if 1 Culikers al. 2018 FFor the spectral method used the authors' Python impimeentatinn availa ble lorenzodallamico. www.gb.tooooooooo htt.uubboo//oo/eee under the CC BY 4.0 license. For DeepWalk, used the C++ implementation of atthub.com/thiaane httph//aaaaaaaiiaee////oaaa honeeeeeedeatteeeeeeeeeeeeeee node2vec with default values, released under the Apache License 2.0.
2017
11

---

Published in Transactions on Machine Learning Research (05/2025)
5 Conclusions
This article introduces linear-time approximation of the softmax scores normalization for embedding vectors with bounded norms. Our result relies on a variational approximation of the unknown scalar product probability distribution between embedding vectors. We provide both theoretical and empirical validation for our estimation formula. By testing our approximation several empirical datasets, we showed that it can achieve high accuracy levels, outperforming the competing methods.
Based on this result, we describe a simple embedding algorithm in the 2Vec style with a loss computed in Onn²) operations. Because of its complexily, the competing methods use alternative loss functions, while our algorithm efficiently optimizes the original loss function. The proposed EDRep algorithm is general-purpos and takes a probability matrix P as input. To prove its efficiency, we tested it on a few use cases and made specific choices for the matrix P. The simulations showed that simple and intuitive definitions of such matrix could lead to higher or similar performances compared with competing algorithms. We also observed the EDRep algorithm to be systematically faster than its competitors.
Let us now consider some limitations of our work. Given the generality of the formulation, we did not provide a bound to the error on the Z, estimation. However, we extensively tested our method on several embeddings and matrices P, weighted and unweighted, symmetric and not, and with different sparsity levels. In all cases, the results confirmed the goodness of our proposed approach. The reasons justifying these results are two: (1) our approximation needs only to hold in "distribution" sense and we do not need more stringent point-wise accuracy; (2) the multivariate normal approximation is particularly powerful to approximate the distribution of the scalar product. Unsurprisingl, we also observed that the performance of EDRep compared to the 2Vec methods- - highly depends on the matrix P. In some cases, our method provides a neat advantage in terms of performance, while in others the results are essentially identical, with our method being faster. We lack a clear interpretation of the role of P in determining the embedding quality and the convergence speed and we believe this aspect deserves further investigation in the future. We highlight, however, that this analysis is desk-dep,,,,n,, and finding good P should specifically address precise research question.
Our approximation is simple yet efficient method to bypass the quadratic computational complexity re- quired by the softmax normalization. The direct application of this result to the proposed EDRep algorithm has several advantages. Despite its simplicity, the practicality of this algorithm makes particularly appeal- ing to machine learning users who can easily adapt the algorithm to an arbitrary embedding problem. For instance, the EDRep algorithm has already been adopted to define the distance between pairs of temporal graphs in (Dall' Amico et al. 2024). Moreover, observed in (Levy et al., 2015)) for word embeddings, tailored fine-tuning, and data preprocessing often have higher impact in determining the embedding quality than the architecture itself. As such, the EDRep constitutes simple, interpretable, minimal unit to elficiently cre- ate distributed representations. Moreover, given the generality of our framework, one can effortlessly adapt EDRep to other similar cost functions. The most immodiate change is to consider a contrastive learning setting which also the term SoflMax(-XY")) appears. With minor modifications to Algorithm ] one can also account for non-normalized (but bounded) embedding vectors, or choose as an arbitrary non-negative matrix. On the other hand, while efficiently dealing with the computation of the softmax scores is crucial in transformer architectures, our results do not directly allow the design of an efficient transformer. How- ever, we envision that ours can be significant contribution to the design of efficient methods to generate distributed representations.
Acknowledgments
LD acknowledges support from the the Lagrange Project of the ISI Foundation funded by CRT Foundation, from the European Union's Horizon 2020 research and innovation programme under grant agreement No. 101016233 (PERISCOPE) and from Fondation Botnar (EPFL COVID-19 Real Time Epidemiology I-DAIR Pathfinder). EMB acknowledges support from CRT Lagrange Fellowships Data Science for Social Impact of the ISI Foundation. The authors thank Ciro Cattuto for fruitful discussions and Cosme Louart for his valuable feedback on the proofs of Theorems A.1
A.3
12

---

Published in Transactions on Machine Learning Research (05/2025)
References
Tavor Baharav, Ryan Kang, Colin Sullivan, Mo Tiwari, Eric Sager Luxenberg, David Tse, and Mert Pilanci. Adaptive sampling for efficient softmax approximation. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
R. Bamler and S Mandt. Extreme classification via adversarial softmax approximation. In International Conference on Learning Represenlations, 2020.
Albert-László Barabási and Réka Albert. Emergence of scaling in random networks. Science, 286(5439): 509-512, 1999.
Aurélien Bellet, Amaury Habrard, and Mare Sebban. Metric learning. Synthesis lectures on artificial intel- ligence and machine learning, 9(1):1-151, 2015.
Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arvi preprint arXic:2004.05150 2020.
Yoshua Bengio and Jebn-Sébastien Senécal. Quick training of probabilistie neural nets by importance sam- pling- In International Workshop Artificial Intelligence and Statistics, pp. 17-24. PMT.R 2003.
Yoshua Bengio, Aaron Courville, and Pascal Vincent. Representation learning: A review and new perspec- tives. IEEE transactions on pottern analysis and machine intelligence, 35(8):1798 1828, 2013.
Jacob Charles Wright Billings, Mirko Hu, Giulia Lerda, Alexey Medvedev, Francesco Moltes, Adrian On- icas, Andrea Santoro, and Giovanni Petri. Simplex2vec embeddings for community detection in simplicial complexes. arXiv preprint Xrxiv:1906. 09068, 2019.
Guy Blanc and Steffen Rendle. Adaptive sampled softmax with kernel based sampling. Tn International conference on machine learning, pp. 590 599. PMLR, 2018.
Ciro Cattuto, Wouter Van den Broeck, Alain Barrat, Vittoria Colizza, Jean-François Pinton, and Alessandro Vespignani. Dynamics of person-to-person intoractions from distributed rfid sensor networks. PLOS ONE, 5(7)xe11596, 07 2010. doi: ho0/1//jormal..pne..000 2F 2Fjournal. pone. www.:poon......... 00115996
URL http://dx. o0oooorrr/11.11111
Long Chen, Fajie Yuan, Joemon M Jose, and Weinan Zhang, Improving negative sampling for word repre- sentation using self-embedded features. In Proceedings of the Eleventh ACM International Conference on Web Search and Data Mining, pp. 99-107, 2018.
Krzysatof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreca Gane, Tamas Sarlos, Peter Hawkins, Jared Quiney Davis, Afroz Mohiuddin, Lukasz Kaiser, al. Rethinking attention with performers. In International Conference on Learning Representations, 2020.
Lorenzo Dall' Amico, Romain Couileet, and Nicolas Tremblay. Revisiting the betho-hessian: improved com- munity detection in sparse heterogeneous graphs. In Advances in Neural Information Processing Systems, pp. 4039 4049, 2019.
Lorenzo Dall' Amico, Romain Couillet, and Nicolas Tremblay. A unified framework for spectral clustering in sparse graphs. Journal of Machine Learning Research, 22(217):1-56, 2021.
Lorenzo Dall' Amico, Alain Barrat, and Ciro Cattuto. An embedding-bbase distance for temporal graphs. Nature Communications,, 15(1), November 2024, ISSN 2041-1723. doi: (0.1088/4411667-224-52222 URL
http://dd..da..orr/00..033
1038// 00010///211477-2444228880
Bhuwan Dhingra, Zhong Zhou, Dylan Fitzpatrick, Michael Muchl, and William W. Cohen. Tweet2vec: Character-based distributed representations for social media. 54th Annual Meeling of the Association for Computational Linguistics, ACL 2016 Short Popers, pp. 269 274, 2016. doi: 10. /v v1/p16111pvepp 2044.
Petros Drincas, Michael W Mahoney, and Nello Cristianini. On the nyström method for approximating a gram matrix for improved kernel-based learning. Journal of Machine Learning Research, 6(12), 2005.
13

---

Published in Transactions on Machine Learning Research (05/2025)
Jingcheng Du, Peilin Jia, Yulin Dai, Cui Tao, Zhongming Zhao, and Degui Zhi. Gene2vec: distributed representation of genes based on co-expression. BMC genomics, 20(1):7-15, 2019.
Santo Fortunato and Darko Hric. Community detection in networks: user guide. Physics reports, 659: 1 44, 2016.
Zheng Gao, Gang Fu, Chumping Ouyang, Satoshi Tsutsui, Xiaozhong Liu, Jeremy Yang, Christopher Gessner Brian Foote, David Wild, Ying Ding, et al. edge2voc: Representation learning using edge semantics for biomedical knowledge discovery. BMC bioinformatics, 20(1):1 15, 2019.
Joshua Goodman. Classes for fast maximum entropy training. In 2001 IEEE International Conference on Acoustics, Speech, and Signal Processing. Proceedings (Cal. No. 01 01CH72221, volume 1, pp. 561-564. LEEE, 2001.
Palash Goyal, Sujit Rokka Chhetri, and Arquimedes Canedo. dyngraph2vec: Capturing network dynamics using dynamic graph representation learning. Knowledge-Based Systems, 187:104816, 2020.
Aditya Grover and Jure Leskovec. node2voc: Scalable feature learning for networks. In Proceedings of the 22nd ACM SIGKDD international conference on Knowledge discovery and data mining, pp. 855 864, 2016.
Lennart Gulikers, Mare Lelarge, and Laurent Massoulié. A spectral method for community detection in moderately sparse degree-corrected stochastic block models. Advances in Applied Probabilily, 49(3):686- 721, 2017.
Lennart Gulikers, Mare Lelarge, and Laurent Massoulié. An impossibility result for reconstruction in the degroo-corrocted stochastic block model. The Annals of Applied Probability, 2018.
Ariel Jaffe, Yuval Kluger, Ofir Lindenbaum, Jonathan Patsenker, Erez Peterfreund, and Stefan Steinerberger. The spectral underpinning of word2vec. Frontiers in applied mathematics and statistics, 6:593406, 2020.
Brian Karrer and Mark EJ Newman. Stochastic blockmodels and community structure in networks. Physical review E, 83(1)0016107, 2011.
Seyed Mehran Kazemi, Rishab Goel, Sepehr Eghbali, Janahan Ramanan, Jaspreet Sahota, Sanjay Thakur, Stella Wu, Cathal Smyth, Pascal Poupart, and Marcus Brubaker. Time2ve:: Learning a vector represen- tation of time. arXiv preprint ar arXix:1907.055 2019.
Matt J Keeling and Pejman Rohani. Modeling infectious discases in humans and animals. Princeton uni- versity press, 2011.
Andrei Kutuzov, Murhaf Fares, Stephan Oepen, and Erik Velldal. Word vectors, reuse, and replicability: To- wards a community repository of large-text resources. In Proceedings of the 58th Conference on Simulation and Modelling, pp. 271-276. Linköping University Electronic Press, 2017.
Andrew J Landgraf and Jeremy Bellay. Word2vec skip-gran with negative sampling weighted logistic pca. XrXiu preprint arXio:1705.0077
2017.
Lue Le Magoarou and Rémi Gribonval. Flexible multilayer sparse approximations of matrices and applica- tions. IEEE Journal of Selected Topics in Signal Processing, 10(4):688 700, 2016.
Michel Ledoux. The concentration of measure phenomenon. American Mathematical Soc., 2001.
Omer Levy, Yoav Goldberg, and Ido Dagan. Improving distributional similarity with lessons learned from word emboddings. Transactions of the association for computational linguistics, 3:211-225, 2015.
Stuart Lloyd. Least squares quantization in pem. IEEE transactions on information theory, 28(2):129-13) 1982.
WAJ Luxemburg. Arzela's dominated convergence theorem for the riemann integral. The American Mothe- matical Monthly, 78(9):970 979, 1971.
14

---

Published in Transactions on Machine Learning Research (05/2025)
James MacQueen. Some methods for classification and analysis of multivariate observations. In Proceedings of the Fifth Berkeley Symposium on Mathemalical Statistics and Probabilily, Volume 1: Stafistics, volume 5, pp. 281-282222 298. University of California press, 1967.
Tomas Mikolov, Nya Sutskever, Kai Chen, Greg S Corrado, and Jeff Dean. Distributed representations of words and phrases and their compositionality. Advances in neural information processing systems, 26, 2013.
David Mimno and Laure Thompson. The strange geometry of skip-gram with negative sampling. In Empirical Methods in Natural Language Processing, 2017.
Frederic Morin and Yoshua Bengio. Hierarchical probabilistic neural network language model. In Interna- tional workshop on artificial intelligence and statistics, pp. 246 252. PMLR, 2005.
Cun Mu, Guang Yang, and Yan Zheng. Revisiting skip-gram negative sampling model with rectification. In Intelligent Computing: Proceedings of the 2019 Computing Conference, Volume 1, pp. 485-497. Springer, 2019.
Stephen Mussmann, Daniel Levy, and Stefano Ermon. Fast amortized inference and learning in log- lincar models with randomly perturbed nearest neighbor search. arXiv preprint arXtt:1707.083773 2017.
Annamalai Narayanan, Mahinthan Chandramohan, Rajasekar Venkatesan, Lihui Chen, Yang Liu, and Shan- tanu Jaiswal. graph2vec: Learning distributed representations of graphs. arXiv preprint 1rx: 05005, 2017.
Mark EJ Newman. The structure and function of complex networks. STAM review, 45(2):167-25) 2003.
Patrick Ng. dna2vec: Consistent vector representations of variable-length k-mers. arXiv:1701 701.0629177 2017.
or preprint
Giang Hoang Nguyen, John Boaz Lee, Ryan A Rossi, Nesreen K Ahmed, Eunyee Koh, and Sungchul Kim. Continuous-time dynamic network embeddings. In Componion proceedings of the the web conference 2018, pp. 969-976, 2018.
Maximillian Nickel and Douwe Kiela. Poincaré embeddings for learning hicrarchical representations. Ad- vances in neural information processing systems, 30, 2017.
Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah A Smith, and Lingpeng Kong. Random feature attention. In 9th International Conference on Learning Representations, ICLR 2021, 2021.
Bryan Perozzi, Rami Al-Rfou, and Steven Skiena. Deepwalk: Online learning of social representations. In Proceedings of the 20th ACM SIGKDD international conference on Knowledge discovery and data mining, pp. 701-710, 2014.
Pengda Qin, Weiran Xu, and Jun Guo. A novel negative sampling basod on tfidf for learning word repre- sentation. Nearocompuling, 177:257 265, 2016.
Ali Rahimi and Benjamin Rocht. Random features for large-scale kernel machines. Advances in neural information processing systems, 20, 2007.
Mahmudur Rahman, Tanay Kunar Saha, Mohammad Al Hasan, Kevin S Xu, and Chandan K Reddy, Dylink2vec: Effective feature representation for link prediction in dynamic networks. arXio preprint arXiv:1804.0572 2018.
Ankit Singh Rawat, Jiecao Chen, Felix Xinnan X Yu, Ananda Theertha Suresh, and Sanjiv Kumar. Sampled softmax with random fourier features. Advances in Neural Information Processing Systems, 32, 2019.
Benodek Rozemberezki, Ryan Davies, Rik Sarkar, and Charles Sutton. Gemsoe: Graph embedding with self clustering. In Proceedings of the 2019 IEEE/ACM international conference on advances in social networks analysis and mining, pp. 65 72, 2019.
15

---

Published in Transactions on Machine Learning Research (05/2025)
Koya. Sato, Mizuki Oka, Alain Barrat, and Ciro Cattuto. Predicting partially observed processes on temporal networks by dynamics-aware node embeddings (dyane). EPJ Data Science, 10(1), 2021. doi: 10.1140/ epids//113688-021-002777
Yingchun Shan, Chenyang Bu, Xiaojian Liu, Shengwei Ji, and Lei Li. Confidence-aware negative sampling method for noisy knowledge graph embedding. In 2018 IEEE International Conference on Big Knowledge (ICBK), pp. 33-40. IEEE, 2018.
Kyuhong Shim, Minjae Lee, Iksoo Choi, Yoonho Boo, and Wonyong Sung. Svd-softmax: Fast softmax approximation on large vocabulary neural networks. Advances in neural information processing systems. 30, 2017.
Michel Talagrand. Concentration of measure and isoperimetric inequalities in product spaces. Publications Mathêmatiques de l'Intiiut des Hautes Etudes Scientifiques, 81:73 205, 1995.
Maddalena Torricelli, Márton Karsai, and Lactitia Gauvin. weg2vec: Event embedding for temporal net-
works. Scientific Reports, 10(1):1 11, 2020.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
Clrike Von Luxburg. A tutorial on spectral clustering, Stalistics and computing, 17:395 416, 2007.
Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn Fung, Yin Li, and Vikas Singh. Nyströmformer: A nyström-based algorithm for approximating sell-attention. In Proceedings of the AAAI Conference Artificial Intelligence, volume 35, pp. 14138 14148, 2021.
Manzil Zaheer, Guru Guruganesh; Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, ct al. Big bird: Transformers for longer soquences. Advances in neural informalion processing systems, 33:17283 17297, 2020.
A Concentration theorems
In this appendix provide the enunciate and prove the Theorems motivating Equation 2 The first result describes the concentration of Z,/m around its mean.
Theorem A.1. Consider veclor a, E Rd and set (Vi...., ,nm of m independent random vectors in Rd.
Let
T yal < h = =mm(1) for all a. Letting Z; be defined as in Equation 1 then for all t 0
P
Z
E
Zi
S de (Vm7//44)))
Before proceeding with the proof of Theorem A.1 we will use in our demonstration.
let. us enunciate the following concentration theorem that
Theorem A.2 ( (Talagrand] 1995), (Ledoux 2001 Corollary 4.10). Given a random vector E [u, v|" with independent entries and a 1-Lipschitz (for the euclidean norm) and convex mapping g F RR one has the concentration incquality:
((g((w) - Rig(w)]] 2 1) S 4e - 2/4(a
VI>O:
Proof (Theorem A.1 Let ( E R be a vector with entries {xe Yo l Then, the vector ( satisfies the lypotthesi of Theorem A.g for all and for u = -u = h. We let g be:
)(wo())
16

---

Published in Transactions on Machine Learning Research (05/2025)
then immediately verify that Z, =m)(wl))). We are left to prove that g satisfies the hypotheses of Theorem A.2 well. Firstly, is conver because il is the sum of convex functions. We now compule the w,w' w' E [-h,, hl one obtains the following bound:
Lipschitz parameter. Considering w,
(0)
m
(
11
(c)
(a(w) - g(w')l
S
We
w/1
w w'1
m
m1
m
a=1
I
|m
||w
w'll
w
w'll
m
where in (a) we used the triangle inequality, in (b) we exploited the fundamental theorem of calculus, in (c) the |.1 is meant entry-wise and finally in (d) we used the Cauchy-Schwarta inequality. This set of inequalities implics that vmng is 1-Lipschitz and is suitable choice to apply Theorem A.2. Theorem A.1 is then easily obtained from a small play on and exploiting the relation between Z4 and ()w(
We remark that this result holds unchanged the case = Y. The proof can be easily adapted by singling out the negligible contribution given by |mll? S se (1), obtaining the same result.
Consider a vector æy E Rd and a set [y...... ymd of m independent random vectors in Rd. Let |oo yal S Om (1) for all a. Let fi denote the empirical distribution of x. ya and assume fi.m converges in distribution to fi, then,
Theorem A.3.
Zi m
= ["
lim
e' f((t).
m-+00
Also in this case, before proceeding with the proof, let us enunciate the dominated convergence theorem that allows one invert the integral and limit signs.
Theorem A.4 (Luxeebuurg 1971)). Let Fn be a sequence of Riemann integrabbe functions defined on bounded and closed interval aa b] which converges on [a, b) a Riemann- inttunnabtetteaelee function F. If there exists consland salisfying [F,(()) s M] for all [a, b) and for all n, then
lim dt (t = /
S" dt F(t)
dt lim Fx(t) no0
Proof (Theorem A.3 densities F.,,..., Pi where Pop denotes the integral of fi.
The values x yo = wa are a sequence of m random variables with cumulative
lim E
Zi
lim
m
E
w(l))
lim E (e
lim m 000
S* dt e' fim(()
m
m--000
771
m7
- -1" dt e² FFim(t))
-f" e'Fi(1) S* dt e(ft().
lim
e' (t)
(b) e' en((0 m00
000
In (a) we performed an integration by parts; in (b) we exploited the fact that convergence in distribution implies the pointwise convergence of the probability densily function and we applied Theorem A.4 for M ch
17

---

Published in Transactions on Machine Learning Research (05/2025)
B Derivation of Equation
7
We here report the derivation of Equation
from
(
-= Σ Pss log (oftMax(YYY") ks)
+ Σ yk yj
ijEV
iJEV
(a)
Σ Po)(YTT)ss
bg +
Σ
Psj log(Zz) + (YY
LJVV
WEV
L.EEV
Σ PYY
10
Σ
log( Z1)
1 Σ
(1n1 )a (YY?)
n
IEE
iEV
LJEV
Σ (PYY*)
log 74)
Σ
iEV
NEV
-r((P
+
log( 7r)) + - tr 1 TYY?)
ieV
1
-tr(Y +
log Z,) + - tr(Y T 1, 1Y)
iEV
where in (a) we used the softmax definition and rewrote y, y; = (YY") Doy in (b) we used Ejev Pij = 1; in (c) we used the fact that yy symmetric matrix; in (d) we leverage the trace definition; in (e) we use the property of the trace tr( (B) tr(BA).
C Derivation of the gradient
We here derive the gradient expression as it appears in Equation Note that in this derivation, the
quantities Na: Pa are considered as constants, in a stochastic gradient descent fashion. We observed that this gradient form achieves better results in fewer epoch. Let us first rewrite the loss function of Equation Following the passages detailed in Appendix we obtain
C=
Σ ijev
Ps
x xy + Σ log(Z.)
LEV
We now introduce the approximation of Equation
and rewrite
L Σ
Pij
+
Σ Σ
log
nua exp
2 Ma + T Tonnn
2 go Ano
LJEV
iEV
a=
Σ Kig 2a
=
Pij
Lig 2
+
log
na
exp
Dighaag
+ 2
11
ijeV 9=1
iEV
9,p=1
We now compute the derivative with respect to ake to obtain the respective gradient term.
1 1 Sgr [du..i + 85b2igg
O L = Σ Σ
Pij
n2)
L.EV q=11
Dmm Taa u nn tot as Eg d + Edp Po.op 1ou - du 6u)
Dest Taa e + ou Rax.
NEV
18

---

Published in Transactions on Machine Learning Research (05/2025)
We now use the notation Zia Ta et
No Ram and compute the sums.
Dz C
Σ
Pik
Vir
Σ
Paj -
1 Zj
11
n2
IEV
JEV
d
+
3ka
Na,r
1
Σ
Tkg Kkg Da. +
1
Σ
SkygSarer
Ea=1
Zka
2
2 9=1
9=1
Now we recall that Ex.1 Zia = Zi/n and Mai = Harr and that Qa = =52''
X )
-
1
1
X
+
(ZM
+
Σ
K
Zual (X (xealarr
n
kr
12
kr
(Zln)k
a=1
thus obtaining Equation
D Experiment implementation details
We here report some implementation details the experiments we conducted. This section complements the information in the main text when this insufficient reproduce our results.
D.1 Node embeddings
In the experiments we test the node embedding problem for the task of community detection. To do so, we work with synthetic graphs generated from the degree corrected stochastic block model (DCSBM) (Karrer & Newman 2011) that we here define.
Definition D.1 (DCSBM). Let w V -L........ be a class labeling function, where is the number of classes. Let P(w; = a) q- and consider two positive integers satisfying Cin > Cout 2 O. Further Let 0 ~ pa be a random variahle that encodes the intrinsic node connectivity, with E[ń] = and finite variance. For all E VVO, is drawn independenlly al random from pa- The entries of the graph adjacency matria are generaled independently (up to symmetry) at random with probability
P(Ag = =1) 0,0,
Cin
if wi) = w(i)
Cout
else
In words, nodes in the same community (w(f)= wwj)) are connected with higher probability than nodes in different communities. From straightforward calculation, the expected degree is E[d,] 0,, thus allowing one to model the broad degree distributions typically observed in real networks (Barabási & Albert 1999). Given this model, the community detection task consists in inferring the node label assignment from a realization of A. was shown that this is theoretically feasible (in the large n regime) and only if
(c- Cout) /v /E0A > 1. This is the parameters appearing in the main text. In the simulations the 0,'s are obtained by: i) drawing a random variable from uniform distribution between and 12; ii) raising it the power 6; ifi) normalize it so that E[Ø] 1. This leads a rather broad degree distribution, even if it maintains finite support.
For all three methods under comparison obtain the embedding vectors from A and then cluster the nodes into communitics by applying k-means and supposing that the number of communities is known. The algorithm of &Dall' (Dall'Amico Amico el al. 2019 2021) obtains the embedding by extracting a sequence of eigenvectors from a sequence of parameterized matrices that are automatically learned from the graph and does not require any parametrization. The DeepWalk and EDRep algorithms generate emboddings in 32 dimensions. We observed that the results are essentially invariant a large spectrum of d values for both embedding algorithms.
19

---

Published in Transactions on Machine Learning Research (05/2025)
D.2 Dynamically aware node embeddings
This experiment fealures three main steps: the creation of the supra-adjacency matrix from temporal network; ii) the creation of an embedding based on this matrix; iii) the reconstruction of a dynamical process taking place on the network. We now describe these steps in detail.
Definition of the supra-adjacency matrix
We consider a temporal graph collected by the SocioPatterns collaboration. This dataset describes face-to- face proximity encounters between people at conference. The data are recorded with a temporal resolution of 20 seconds and correspond to interactions within a distance of approximately 1.5 meters The dataset contains approximately 3000 different time-stamps. Following the procedure of (Sato et al. 2021 we ag- gregate them into 180 windows of approximately 15 minutes each. We thus obtain a. temporal graph a sequence of weighted temporal edges (i,j. E, wije), indicating that i, interacting
that is represented snapshot, | for cumulative time equal to Wijt- We say that a node is active at time / if it has neighbors time t. We let t.o be the time at which node is active for the a-th time. Given this graph, then build the supra-adjacency matrix that is defined as follows.
Definition D.2 (Supra-adjacency matrix). Consider temporal graph represented as sequence of weighted temporal edges (i. L, wige). We define a set of "lemporal nodes" given by all the pairs i, li, where i is a node of the temporal graph and li,a is the time of the a-th appearance of in the network. We denote this set D, formally defined as
D D= (ii iEV,DjEV. (i,i,tia) EE).
where V is the set of nodes and & of temporal edges. The cardinality of this set is D (D). We then define a directed graph with DD being the set of nodes. The directed edges are placed between
lia) + (i,liat1)
self connections
(i,ti,a) + (i,tg,61)
if t, = t3.6 and (i.j,t) E
(i,t + (i,ta,at1)
if ti t..6 and (i.j,t) E
The supra-adjacency matrix Asupra E DXD is the adjacency malrix of the graph we just defined.
Creation of the embedding
Given Asupra, the authors generate an embedding with the DeepWalk algorithm. Here, the random walker can only follow time-respecting paths, by construction of the matrix. We let the random walks have length equal to 15 steps and deploy also in this case the embedding dimension d = 30. For EDRep we use the row-normalized version of Asupra as As a result we obtain an embedding vector for cach (i, (,,))) pair.
Reconstruction of partially observed dynamic process
Following (Sao et al. 2021 we consider an epidemic process taking place on the temporal network. We use the SIR model (Keeling & Rohani 2011) which infected nodes (I) can make susceptible nodes (S) to transition to the infected state with a probability B3 if they are in contact. Infected individuals then recover (R) with a probability l and unable to infect or get infected. We run the SIR model letting all nodes to be in the S state the beginning of the simulation and having one infected node. The experiment is with 3 0.15 and 0.01 and it outputs the state of each node at all times, i.e., for all (i,4)
We then suppose to observe a fraction of these states (every node is expected to only be observe once) and obtain a binary variable for each (i,4,a) indicating whether node i was infected at time f; Exploiting the embeddings, we train a logistic regression model to predict the state the node each unobserved time and compare the predicted number of infected individuals against the observed ones. We repeat the experiment for 50 different realizations of the observed training set.
The experiment descrihed in CCttuto al. 2010). The data are shared under the Creative Commons Public Domain Dedication license and can be downloaded http:/ hwtp://ctoopateene..rrgo///eoa htwpsa/seeeees....oooooeeeeeeeee
20


Reviews:
[{"review_id": "fS3zVKbn5k", "reviewer": "Reviewer_JiP6", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n* Clearly written overall, easy to follow paper\n* Nice presentation of a simple method with diverse applications\n* Interesting clustering-based approximation to softmax\n\nMajor concerns:\n* W1: Presentation of Attention Score Matrix - The baseline methods compared in section 2 ([Random Feature Attention](https://arxiv.org/abs/2103.02143) and [Performers](https://arxiv.org/pdf/2009.14794) are used to approximate the attention score matrix in a transformer's attention layer. In my opinion, this is a significantly different setting than the one proposed in this paper, with different computational tradeoffs. I have the following concerns surrounding this: \n  * W1.1: Why are these the right baseline methods to consider for the setting used in this paper? What about considering more classic density estimation methods? What about simple sampling or top-k based approximations? Or simple matrix factorizations style approximations?\n  * W1.2: Another concern about these as baselines, is that they do not share the same $m\\rightarrow\\infty$ assumption in many use cases. \n  * W1.3:  In measuring the computational cost of the mixture-of-gaussian based method for softmax approximation (E.g. Equation 2), the $O(d^2)$  ($d$ as in dimension of the vectors), is discussed briefly, but then quickly dropped. This is a quite significant term. It is especially significant for the use case of approximating the attention matrix in a transformer layer. This seems to make the method well suited for the $d << m$ case, but less well suited for the $d < m$ case. Should we really not be so concerned about this quadratic term? I guess the thought is $d$ might be on the order of a few hundred so it is ok? Nevertheless, when comparing to methods without this term, it seems significant to mention. \n* W2: Approximation quality guarantees of EDRep \n * While the derivation of the EDRep algorithm is nice, it feels largely incomplete because it is missing, as I understand statements about either: generalization, quality of approximation / representational capacity, or convergence.\n  * It is not entirely clear to me the advantages of EDRep over alternatives, alternative models for the data etc. Especially for many of the applications that are considered in this paper, e.g. the \"2Vec\" use case (where countless alternative parameterizations have been considered).\n* W3: Missing related work: \n* While the paper does cover a good deal of related work, I think there are significant gaps for instance:\n  * Baharav et al, 2024 [https://proceedings.neurips.cc/paper_files/paper/2024/file/d52dbd66219dc4e432e0bd4f9c25c4c3-Paper-Conference.pdf] \n  * Tree-based methods such as [Blanc et al 2018](https://proceedings.mlr.press/v80/blanc18a/blanc18a.pdf)\n  * For Word2Vec relationship specifically: https://aclanthology.org/Q15-1016.pdf\n\nMore minor concerns:\n* The complexity of the proof is more limited than I would expect for TMLR, yet the paper reads in a largely theoretical way. I think a better presentation could be to strong emphasis the empirical contributions (once the above concerns are taken care of) and then using the theoretical analysis to support the empirically motivated work. \n* The proposed method is quite simple, but if presented in full level of rigor amidst the kinds of baselines and open questions described in above weaknesses, might be enough for a full paper contribution.\n* The choice of clustering algorithm seems important for the approximation quality, especially theoretically. Why is K-means the right choice given that the full covariance matrix is used? Please provide more explanation and details\n* It would be nice to give a bit of intuition about Equation 2, given this paper is generally accessible, i think this would be important to readers. E.g. covariance operator?\n* Some of the writing in the introduction, could be made more crisp: How common is the terminology: attention score matrix SoftMax(XY T ) ? I have seen it for an attention layer, but perhaps when defining it broadly like this, good to put citations? Also, I found the sentence \" The optimization of the attention scores motivates several embedding algorithms such as Word2Vec Mikolov et al. (2013) and the many “2Vec”-type algorithms that followed\". Hard to parse and a bit vague. It is not obvious at this point what you mean by optimization of the attention scores. The use of “several” and “many” is confusing also.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper studies approximations of the \"attention score matrix\", $\\textrm{SoftMax}(XY^T)$.\n\nThe paper proposes:\n1) a mixture of gaussian-based, linear in the number of rows of $Y$ approximation to the partition function used for normalizing the rows of $XY^T$ in the softmax function.\n2) An algorithm EDRep for learning embedding in the \"2Vec\" style of methods (e.g., matrix factorization).\n3) Performs analysis of the proposed approach and empirically investigates EDRep in diverse applications.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please comment on / address W1-W3 and address attention score matrix terminology.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "Z0DOF10FF9", "reviewer": "Reviewer_iPsW", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths:\n1. Paper provides some theory to justify their algorithm, although it is quite standard. The main idea is showing that the normalization constant in the denominator above converges fast to its mean which can be explicitly computed if we assumed embeddings come from a parametric mixture of Gaussians.\n\n2. Given experimental results are very rigorous. It measures a variety of metrics, computational costs and provide error margins\n\n3. Paper is well written and decently motivated. However, it missed the contemporary context (see below).\n\n4. Experimental results are very promising and show that the proposed approximation mechanism works well for the problems studied.\n\n5. EDRep algorithm provided seems novel (although straight forward) and it is effective in reducing computational cost\n\n### Weaknesses:\n1. Theory seems to be a missing both minor and major details. What is $h$ in Thm 2.2? How is (2) derived? How is (3) derived and why is K-means clustering the solution?\n\n2. Paper assumes that $\\mathcal{V}_i$ consist of embeddings which are independently chosen. This works when $\\mathcal{V}_i$ is large enough or if it is the whole universe of entities (like full vocab of tokens). However, when $\\mathcal{V}_i$ is restricted to small window of entities (e.g. context of words around a word in a sequence) independence may not hold true, hence the approximation technique might fail. Similarly, any kind of position encoding used in transformers would break the independence assumption.\n\n3. Many of the experiments (Fig 1 and 2) provide absolute error but fail to provide relative error which is more important for such calculations. \n\n4. Experiments only study quite old \\*2vec style methods and they do not study actual attention mechanism in deep learning models. Title is also misleading in that sense. It might also be true that \\*2vec style methods are not that powerful when compared to transformer/GNN style models.\n\n5. If the empirical entity interaction probability matrix $P \\in [0,1]^{n \\times n}$ is not sparse, then during training we would require the computation of $O(n^2)$ $q_j^{(i)}$'s which requires $O(n^2d)$ computations and the proposed attention mechanism won't produce linear complexity. In that case it is not clear how $P$ can be decomposed into product of $m$ sparse matrices to reduce the complexity.\n\n6. Paper do not discuss or compare with many works on reducing attention complexity [1, 2]. See [3] for comprehensive discussion. Therefore, it is not clear where the proposed method sits in the current literature space.\n\n7. It is not clear whether *2vec style baselines in Sec 4 use any computationally efficient techniques. For example, in standard word/node-2vec only a small number/subset of positive and negative examples are chosen per sample which effectively guards the complexity from blowing up quadratically. So it isn't clear how such simpler techniques compare against the proposed method.\n\n[1] Zaheer, M., Guruganesh, G., Dubey, K. A., Ainslie, J., Alberti, C., Ontanon, S., ... & Ahmed, A. (2020). Big bird: Transformers for longer sequences. Advances in neural information processing systems, 33, 17283-17297.\n\n[2] Dao, T., Fu, D., Ermon, S., Rudra, A., & Ré, C. (2022). Flashattention: Fast and memory-efficient exact attention with io-awareness. Advances in Neural Information Processing Systems, 35, 16344-16359.\n\n[3] Wan, Z., Wang, X., Liu, C., Alam, S., Zheng, Y., Liu, J., ... & Zhang, M. (2023). Efficient large language models: A survey. arXiv preprint arXiv:2312.03863.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "Paper studies \\*2vec-style embedding models like word2vec or node2vec. Although title has the term \"attention\" in it, this paper exclusively studies \\*2vec-style embedding models where pairwise interaction probabilities between two entities i and j are modelled as $q_{j}^{(i)} = \\exp(x_i^\\top x_j)/\\sum_{k \\in \\mathcal{V}} \\exp(x_i^\\top x_k)$ where $\\mathcal{V}_i$ is the set of $n$ entities under consideration and $x_i \\in \\mathbb{R}^d$ is the embedding of entity $i$. Note that authors do not particularly consider or run experiments on attention mechanisms in Deep Learning models which makes the title a bit misleading.\n\nComputing $q_{j}^{(i)}$ for all $i \\in \\{1, \\ldots, n\\}$ even for a fixed $j$ requires $O(n^2 d)$ computations as the denominator requires the computation of $n^2$ dot products. Authors propose to sidestep this computation by modeling distribution of current embedding vectors as a mixture of $\\kappa$ Gaussians. This reduces the complexity of computing the denominator to $O(n \\kappa d)$ for a fixed $j$.\n\nNext authors use this strategy to provide an algorithm EDRep for solving \\*2vec-style embedding problems more efficiently than naive implementations. Authors provide a set of comprehensive experiments to study the effectiveness of their mixture of Gaussian approximation and the EDRep algorithm on three \\*2vec-style embedding problems.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please address the concerns mentioned above.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "WTadOYpVMs", "reviewer": "Reviewer_2r1x", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**:\n\n1. The paper tackles an important problem of efficiently computing attention scores in linear time. \n2. The paper presents experiments in different settings with relevant baselines. \n\n\n**Weaknesses**:\n1. The paper is quite hard to follow especially because of the math notations used in it. Moreover, the motivation seems to change from the introduction. From the abstract, the reader expects an efficient algorithm for attention computation, but the main paper presents an optimization algorithm to extract the embeddings given an attention matrix.\n2. Several of the mathematical notations and assumptions need more clarification. See below:\n    * Is the convergence assumption in Theorem 2.2 realistic? Why should $x_i^Ty_a$ converge to the same value $\\omega_i$ for all $y_a$? \n    * $f_i$ is not defined in Theorem 2.2.\n    * The paper should have pointers to the corresponding proofs in the appendix.\n    * The proof of Theorem 2.2 is hard to follow. It starts off with $\\omega_{a}^{(i)}$ instead of $\\omega_i$. Moreover, $f_i$ or $F_{i, m}$ were never defined. \n    * The approximation in the unnumbered equation after Theorem 2.2 can vary quite a bit depending on the value of $h$. There should be some assumption about $h$ or error analysis of this approximation as a function of $h$. \n\n3. The main problem formulation in Section 3.1 is quite difficult to follow. See below:\n    * why do we care about a probability distribution over the embedding space? The paper so far discussed efficient computation of the attention normalization constant. \n    * The primary optimization objective is shown in Equation 4. What does $p_{0, j}$ indicate? Because this term is undefined, I did not understand the regularization term and its utility. \n\n4. Section 3.3. The description of the algorithm is unclear. Step 1, which vectors do we cluster using the k-means algorithm? \n\n5. The text refers to Figures 3A,B,C, and D but they are not defined in the diagram.\n\n6. Several pieces of text in the paper are unclear. See the comments section. \n\n7. Given that the main contribution of the paper is improving efficiency, there is little empirical support for the time improvement. Figure 2, only compares with the brute force approach. The legends are missing from the time results in Figure 3. \n\n    \n\n\n\n\n**Comments**:\n\n“Note that this approximation does not require the embedding vectors to be, themselves, drawn from a multivariate Gaussian distribution.” — the commas should be removed from this sentence. \n\n\n“We aim at giving a distributed representation of {p(i)}i∈V , in the form of an embedding matrix X ∈ Rn×d preserving the relational patterns encoded by the probability distributions.” — I’m not sure what this sentence implies.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper presents a method to efficiently compute the softmax attention scores in linear time. Particularly, it aims to efficiently compute the normalization constant for these attention scores, which is computationally time-consuming. The work relies on concentration bounds involving the normalization constant, assuming that the corresponding vectors have bounded norms. The paper showed that it is possible to improve efficiency by using a mixture of Gaussian parameterization and estimating the corresponding Gaussian parameters. Empirically they show that this formulation has a better approximation quality than existing baselines like Performers. Based on this observation, the paper presents an algorithm to learn feature representations given similarity values in the corresponding space.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. Please respond to the questions in the weakness section. \n2. The paper in its current state would require some rewriting to make it more presentable to the reader. \n    * Specifically, the paper should clearly state the problem that it is trying to solve up front in the abstract and introduction, whether it is an efficient attention mechanism or estimating feature vectors given a kernel space. \n    * The mathematical notations should be simplified to help the readers. For example, several notations are used when defining the probabilities $P$: $\\mathbf{p}_0$ $\\mathbf{p}_j^{(i)}$, $\\mathbf{p}^{(i)}$, $\\mathbf{p}_{0, j}$. These notations are difficult to follow and should be simplified.\n    * There should be more justifications for key assumptions in the paper. For example, it is unclear to me why $x^T_i y_a$ would converge to a value $\\omega_i$ that is independent of $a$. The paper should also refer to the appendix for the proof of a corresponding mathematical result. \n    * The experimental section also needs to be improved by incorporating detailed annotations for the figures.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper proposes a heuristic linear-time approximation method for efficiently computing normalization constants of softmax functions for embedding vectors with bounded norms.

Reviewers raised initial concerns regarding clarity, theoretical guarantees, and the selection of baseline comparisons. Specifically:
- confusion about whether the approach was designed for transformer architectures or more general embedding methods, recommending clearer framing and explicit statements of the scope.
- necessity for additional comparisons
- notational clarity

The authors addressed these concerns by clarifying their focus, explicitly distinguishing their contributions from transformer-specific methods, justifying their methodological choices, and improving overall manuscript readability. As a result reviewers considered the paper is thorough, is adequate in its descriptions and investigations of ideas, and is a sufficiently informative artifact of information to be shared with the TMLR community.

*Requested changes:*
That said, reviewers have remaining doubts about the choices of methods compared to in one of the principle results in the paper. In figure 1, the authors compare to RFM & Performers when approximating softmax normalization constants in an isolated (non end-to-end transformer) setting. However there is a body of other work on directly approximating the softmax that could have been compared (e.g., simple top-k baseline?, fitting hierarchical softmax or Blanc & Rendle (2018)? or more common matrix factorization / low-rank approaches?). The concern is that comparing with methods mainly intended for use in one way, evaluated for another way, can lead to confusion. The methods compared are designed for efficient attention computation in transformers. It is my understanding that effectively estimating the softmax normalization is not a necessary condition for these approaches, given that they are used in end-to-end systems. So please add a comparison with respect to softmax approximation literature.
Decision: Accept (poster)