Differentiable Expectation-Maximization for Set Representation Learning
Authors: Minyoung Kim
Venue: ICLR 2022 Poster (2022)
Keywords: Representation learning
Bayesian models
Mixture estimation
Optimal transport
Attention

Abstract: We tackle the set2vec problem, the task of extracting a vector representation from an input set comprised of a variable number of feature vectors. Although recent approaches based on self attention such as (Set)Transformers were very successful due to the capability of capturing complex interaction between set elements, the  computational overhead is the well-known downside. The inducing-point attention and the latest optimal transport kernel embedding (OTKE) are promising remedies that attain comparable or better performance with reduced computational cost, by incorporating a fixed number of learnable queries in attention. In this paper we approach the set2vec problem from a completely different perspective. The elements of an input set are considered as i.i.d.~samples from a mixture distribution, and we define our set embedding feed-forward network as the  maximum-a-posterior (MAP) estimate of the mixture which is approximately attained by a few Expectation-Maximization (EM) steps. The whole MAP-EM steps are differentiable operations with a fixed number of mixture parameters, allowing efficient auto-diff back-propagation for any given downstream task. Furthermore, the proposed mixture set data fitting framework allows unsupervised set representation learning naturally via marginal likelihood maximization aka the empirical Bayes. Interestingly, we also find that OTKE can be seen as a special case of our framework, specifically a single-step EM with extra balanced assignment constraints on the E-step. Compared to OTKE, our approach provides more flexible set embedding as well as prior-induced model regularization. We evaluate our approach on various tasks demonstrating improved performance over the state-of-the-arts.

Published as a conference paper at ICLR 2022
DIFFERENTIABLE EXPECIATION-MAXIIIZATTOO FOR SET REPRESENTATION LEARNING
Minyoung Kim
Samsung AI Center Cambridge, UK mi kim21 @mma 1. com
ABSTRACT
We tackle the set2vec problem, the task of extracting vector representation from an input set comprised of a variable number of feature vectors. Although recent approaches based on self attention such as (Sas)Transfoommes were very successful due to the capability of capturing complex interaction between set elements, the computational overhead is the well-known downside. The inducing- point attention and the latest optimal transport kernel embedding (OTKE) are promising remedies that attain comparable or better performance with reduced computational cost, by incorporating fixed number of learnable queries in attention. In this paper ap- proach the set2vec problem from completely different perspective. The elements of an input are considered i.i.d. samples from mixture distribution, and we define our set embedding feed-forward network as the maximum-a-postere estimate of the mixture which is approximately attained by a few Expectation- Maximization (EM) steps. The whole MAP-EM steps are differentiable operations with a fixed number of mixture parameters, allowing efficient auto-diff back- propagation for any given downstream task. Furthermore, the proposed mixture set data fitting framework allows unsupervised set representation learning naturally via marginal likelihood maximization aka the empirical Bayes. Interestingly, we also find that OTKE can be seen as special case of our framework, specifically a single-step EM with extra balanced assignment constraints on the E-step. Com- pared to OTKE, our approach provides more ficxible set embedding as well as prior-induced model regularization. We evaluate approach on various tasks demonstrating improved performance over the state-of-the-arte
(MAP)
INTRODUCTION
Effectively learning succinet and salient feature representation for complex structured data is a central part of recent deep learning. In particular the set structured data, arising in many important scientific tasks including bioinformatics and NLP tasks, poses several challenges to deal with varying numbers of elements and the constraints of permutation invariance, in contrast with conventional
instance-based representation learning (Zaheer et al. 2017 Edwards & Storkey 2017: Lee et al. 2019a Skianis et al. 2020). Perhaps the key successful set representation learning capture the interaction between elements effectively, implemented as self attention (Bahdanau et al. 2015) and
(set)transformer models (Vaswani et al. 2017 Lee al. 2019a).
Despite their superb performance in numerous application problems, one well-known drawback of the self-attention approaches is high computational (time and memory) overhead. Among others, the inducing-poii attention (Lee et al. 2019a) and the latest optimal transport kernel embedding (OTKE) (Mialon et al. 202T) are promising remedies that attain comparable or better performance with reduced computational cost. This is enabled by incorporating fixed number of trainable reference vectors as queries in attention, and especially OTKE, with the attention scheme based on optimal transport, showed state-of-the-ar results several tasks in bioinformatics and NLP.
Although OTKE is motivated from self attention, this paper we show that OTKE can be derived from a completely different perspective: it is the maximum likelihood estimate of Gaussian mixture model obtained by single Expectation Expectation-Mxxmmzzttio some constraints the E-step, which aims to model set elements i.i.d. data samples. From this
(EM) step (Dempster et al., 1977) with

---

Published as a conference paper at ICLR 2022
perspective we propose several interesting extenssion/genrralizatoon such as: i) we perform multiple EM steps instead of a single one, ii) unlike treating the references (model parameters) just as an initial iterate of the EM, we also place prior distribution on the references incorporate the prior effect to the output embedding in a more principled Bayesian manner. With this generalization, our set embedding function becomes (differentiable) ferwaforward network defined as a multi-step EM for the maximum-a -aaommmu------tiiioo posterior (MAP) estimate of a Gaussian mixture. Furthermore, the proposed mixture set data fitting framework allows unsupervised set representation learning naturally via marginal likelihood maximization aka the empirical Bayes.
On several set/sequence representation learning problems including biological sequence classification problems and NLP tasks, our approach shows considerable improvement over the state-of sifte-offtheeartt the- including OTKE. Our main contributions are summarized below:
1. We derive OTKE by novel EM-based mixture fitting perspective, which allows to gener- alize it to a more flexible and principled set embedding function that can offer more modeling options, prior-induced model regularization, and unsupervised learning via empirical Bayes.
2. The whole MAP-EM steps are differentiable operations with fixed number of mixture parameters, enabling efficient auto-diff back-propagation for any given downstream task.
3. Our approach shows improved results over state-of-the-are on various set embedding tasks in bioinformatics and NLP that require effective interaction modeling among set elements.
2 BACKGROUND ON OТКE
OTKE Mialon et al. 202 is a computationally efficient trainable model for representing variable- length set, by aggregating a variable number of features/clements in a set into a fixed-size vector. The input to the model is set S = {xI.... 2n.} (n varies from instance to instance), where each s; E I is an input feature vector which may be an output of some feature extractor network, thus involving some trainable parameters in it. For instance, in the OTKE paper. they often consider parametric Nyström approximation of the RKHS Φ(x) (Williams & Seeger 2001). We abuse the notation throughout the paper. denoting by both raw input element and an extracted feature o(z). where the distinction will be clear from context.
The OTKE model's s set embedding function, emb: : + RD where S, is defined as follows:
n
emb(S) Vp. cat EQu
where cat() indicates concatenation of vectors in column, and hence the output is of dimension D = d p. The trainable parameters of the embedding function are the p reference vectors Z = {ZI,.... } with zj E Rdd, where they serve as queries to determine the weights Q in the attention form (.. Instead of following the popular dot-product strategy (Vaswani et al.] (017)). the weight matrix Q is the (unique) solution of the optimal transport problem (Villani) 20088; Cuturi, 2013) between S and Z with the cost matrix C' defined as the negative kernel Cij -k(2,, 24), that is,
F 1 p' Vj and Σ Eon, 1/n, Vi min Ec s.L. Σou
where the objective is augmented with the entropic term (impact c), H(Q) = - Eig Qis log Qúj
allowing the efficient Kinopo (SK) algorithm to be applied CCuturi 2013).
This OT-based weighting is shown to be more effective than the dot-product in the biological application domain (Mialon et al. 202I). The key difference from dot-product is the balanced assignment constraints Σ, Quj I/p for all j 1,..., P, and the main motivation for this is to have all references zj contribute equally to describing the set S. This may be reasonable strategy in certain scenarios to prevent small number of references from affecting the attention dominantly.
Note that ( can be solved by a few fixed point matrix scaling iterations (SK algorithm): The SK algorithm finds the optimal solution as Q = (aag(u)ADiagin where Arj exp( -Oa-- /e and the vectors E R. and E TR. are the fixed points of uq (Au)), UJ (AT for i = 1...., n,
2

---

Published as a conference paper at ICLR 2022
j i= P. The fixed point iteration usually converges quickly after few iterations. Since the whole SK iterations constitute the feed-forward pass for the final output (T). hence OT is differentiable. Note that as the transport cost in OT is defined as the negative kernel, it becomes simply the dot- product in the RKHS, i.e., A(r,, =) = x = (with the RKHS embedding notation, ((x) 2), where = lives in the same RKHS. OTKE can be extended multi-head network, for which we the above OTKE procedure multiple times independently. each with different parameters Z, then concatenate the outputs.
OTKE has several notable advantages over the self attention. First, it is computationally efficient, scalable to very large dataset. Unlike the quadratic cost of the self attention, OTKE uses a fixed number (p) of references serving as queries. Secondly, it allows unsupervised learning where the model parameters Z can be learned without target (class) labels via the Wasserstein (or k-means alternatively) clustering of set elements as per the OT problem. This unsupervised learning can also serve as pre-training for the potentially unknown target downstream tasks. We also note that regarding the first benefit, the formula (T) is also very similar to the inducing-point attention of SetTransformer (Lee et al. 2019a) (known as the PMA layer) except for the OT-based weighting
scheme. In Appendix B.2 we provide a proof showing that PMA is indeed special case of OTKE.
3 EXPECTATION MAXIMIZATION VIEW OF OTKE
While OTKE motivated and derived from the attention mechanism, in this section we derive the OTKE embedding formulas ( from completely different perspective, namely the EM-based maximum likelihood learning of the Gaussian mixture for set data. Specifically we consider a mixture of p Gaussians with equal mixing proportions and fixed shared spherical covariances, that is,
p(x10) = -
(3)
j=1 p((=)) p(a|c=5)
where e denotes the component latent variable, and the learnable parameters are only the means 0 = --.,..., 2p set S (Z1..... } under the model is log p(510) E, log p(x](0).
) By regarding the set elements v.'s as i.i.d. data samples, the log-likelihood of the
Now we perform one E-step in the EM toward maximizing the log-likelihood, which is derived by the following Jensen lower bound:
log (100) = Σ log Ep(o, = =) Σ Σgen ((ii)) log pla;, G j) i= i=10=1
(4)
q(ili)
1=1 j=1
where v(jli) is the variational distribution and the bound is tight if ((lii = p(c; = j ll..). Although the latter can be derived in a closed form, we aim to formulate an optimization problem for q()))..
We use a joint variational form instead, i.c.. 93 == 4(ilt) g(i) with g(i) == 1/n. Then maximizing
the lower bound in (4) with respect to the (n p) matrix q = (gt)) is equivalent to solving:
min Σ C9 (q),, where C² log p(=; c = j) ||., 1,1 + const. 17
(5)
2e
By having the quadratic kernel ((x 2) = - 1/ 2|| which is (conditionally) positive defi- nite (Schölkopf & Smola 2002), we see that (5) reduces (2) except for the first constraints E; 4-j = 1/p (the other constraints Σ, lij = 1/n are automatically met by construction). Note that the quadratic kernel approximates the dot-product (in the RKHS) if the norms of the points are roughly constant.
Hence the conventional E-step of the EM algorithm coincides with the OT-based attention weight computation (2) without the balanced assignment constraints Σ; qiz 1/p. Conversely, if we consider to maximize the lower bound (4) with the extra balanced assignment constraints on the variational distribution q = (qtz). the corresponding E-step exactly reduces to the OT solution 2
Once the E-step is done (i.c.. q is found), we perform the M-step, that is, maximizing the lower bound in with respect to 0 while q is fixed. This admits a flosed-form update equation,
Ti
(j = 1, 1....,p).
(6)
3

---

Published as a conference paper at ICLR 2022
Note that if we imposed the balanced assignment constraints, the denominator of (6 is constant (1/p), and by concatenating 2 s, exactly reduces to the OTKE formula (I) (up to a constant factor). But, even without the constraints, the M-step in (6) suggests to normalize (balance) the scales of the transformed vectors Σ; (tjT; across j = 1,..- P. according to the reference marginals Σ. Pi (how many points are assigned to 2;).
This establishes a very close connection between OTKE and the EM-updated mean parameters of the Gaussian mixture. Note that our derivation showed that OTKE is recovered by just one EM step, and there is possibility of extending it by taking multiple EM steps. Furthermore, this opens up other interesting generalization of OTKE through the (EM) mixture data fitting perspective. We summarize below reasonable and interesting generalization/exteniii of OTKE in four folds:
1. Instead of fixing mixing proportions and (shared) covariances in the Gaussian mixture as in (3), we can also incorporate them as parameters of the model in conjunction with the means.
2. Instead just one EM step, we can take multiple (k) EM steps with k as hyperparameter.
3. It becomes an option whether the E-step is solved with the balanced assignment constraints (solving the OT problem) or not (regular E-step). This enables more flexible modeling.
4. OTKE can be seen as (single-step) EM which starts from the references zj 's as initial iterate. Not just using them as an initial iterate, we can also place a prior distribution on the Gaussian mixture parameters, that is, p(f) which serves learnable model parameters. This allows us to define the embedding as the posterior distribution emb(S) = p(0]S).
These extensions of OTKE are incorporated into single framework, which becomes a novel set embedding function viewed as differentiable MAP Expectation-aximizzainn Maximizattonn for Gaussian mixture.
4 DIFFERENTIABLE EM (DIEM)
We consider a Gaussian mixture model with non-fixed mixing propor- tions and covariances:
A
θ
p 00)
Jy Na; Vi), 0 (mj, V3Dg-1
L, L,
Each set
where the covariances V, are typically chosen as diagonal (d d) Figure 1: DIEM's graphical matrice to reduce the computational overhead. Given the mixture the set elements are considered i.i.d., p(S)0) ) I[>=> p(ar, (0). Now we impose prior distribution on f where the Dirichlet-Normal-nne Inverse- Wishart (Dir-NIW for short) would be fairly natural choice:
model diagram.
A, (() Dir(m; II NIW(=,, A,
A A= (ug, EdpEEd>
(8)
This defines graphical model (whose diagram shown in Fig. T where A generates mixture parameters 0's, one for each set S, and each 0 generates #1, T'n E S. We emphasize that the reason we consider the prior is that we can express the process of how the model parameters impact the final output in more principled manner, beyond just having them as an initial iterate. For simplicity, the prior parameters 11, X, are shared across the components. That is, Dir(m; 7) x II, . where is scalar, and NIW((z;, V/) = =N(zj: Vg/A) IWWV;; E., for all j 1, ...., ,,, where X, are scalar, M1 E Rd, and Σ, are positive diagonal matrices.
Then we define our set embedding function emb(S) + RD as the MAP estimate of the Gaussian mixture denoted by MM (S) arg maxe p(0]S, A), which is approximated by the k-step EM update. And the prior parameters A will constitute the model parameters of the embedding function emb(S) to trained. The EM update equations for the MAP can be derived in closed forms. More specifically, from the following lower-bound of the log-posteriog ("=o" means equal up to constant)
("=."mean
p(x.,.)
log (0) = log]()]A) + log p(S10) 2 log ((]A( + Σ a(ilt) log
In some of experiments in Sec. used fixed spherical covariances instead which performed better.
4

---

Published as a conference paper at ICLR 2022
S
1" EM
2dd EM
k" EM single EM step
100
A
A
Made of 8 P(O]A)
EM step
A single
0
EM step
single
6
emb A (
((2)
S
Multi Head Pooling (MHP)
emb(S; A)
DIEM(S)
Figure 2: Network diagrams of our Differentiable EM (DIEM) model. (Left) 0 emb(S; A) as a feed-forward network that performs the k-step MAP EM starting from the mode of the prior 6". (Right) Multi-head representation by MHP where we propose three strategies. See text for details.
the E-step maximizes the lower bound with respect to g (with 0 fixed), which yields
(z Vi) E-step: ((li)) Σ3 (M Vi)
or a(ili)
(o 1). (10)
The former formula (10) the regular E-step without the balanced assignment constraints, and the latter is the solution with the constraints where SK(C, e) indicates the solution of the OT problem via the SK algorithm (c.f., (2 and 5 Note that we multiply the latter by n to have conditional distributions (())..
The M-step is the oog-prior regularized weighted ligeliilldood maximization with respect to 9 (with q fixed), which admits closed-form solution (Appendix B. for derivations)
M-step: TíJ >. ((10) + Ty
2/ E. a(ib)) + MOTA
n n+ppTn
Σ.9(10 +
V E - (z))2) + E1 + (w3 - (z
(11)
Σ.
where all operations including squaring are element-wise. The values originate from the prior, defined as: T, = 7 10-, 7> = X, Ti 1, and т, = v + + Note that r's represent the strength of the prior impact: turning them off to 0 completely ignores the prior and leads to the maximum likelihood M-step update. Instead of dealing with them differently, simplify it by collapsing them into a single hyperparameter T (= Tr
= Tv). Then the simplified M-step becomes:
9010) +T T/
n2 p 2, E, ((li)a +ATT Vy E, (()D) + (E, 13)) E1 +T
(zz) (12)
E.90
Now we control 7 to impose the prior impact in the set embedding: smaller (larger) for lower (higher, resp.) prior impact. We treat t as the hyperparameter (e.g., chosen by validation) while A {u. Σ} constitutes the learnable prior parameters.
As an initial iterate 8 of the EM, we use the mode of prior p( A). which equals
1 Initial 0": 3 P
Σ
2 Pjj
(13)
++1
Note that 0 is function of A. We deal with multiple (k) EM steps, 6 0 + following E-step (10) and M-step 12). and the final 0* is a function of A. Overall, from A to the final iterate 8¹ constitutes a feed-forward pass, and it is composed of all differentiable operations. The final output emb( S; A) = 8 ~ 8 AM (S; A) is then transformed to vector representation by pp pp
0h each
concatenating the parameters in 0* in fixed order (e.g., cat ( * V/k m2 =k Vi ....... The ferd-forward network of emb(S: ; A) visualized Fig. (Left).
kk Vk)
Multi-head representation. We can also extent the model to have multi- representation, which simply done by replicating the process H times with different prior parameters A !!...., " See Fig. 2 (Right). We have H outputs from the MAP-EM set embedding networks, denoted by ((h) for h T,..., H. Then the final output pooling from these ((h)'s, and we propose three different options for this multi-head pooling (MHP): i) PC [parameter-concatenattie g(h)' 's in fixed order, ii) SB (select-best) computes the marginal data likelihood log p (S(AA) for each h and select/retumn the one with the largest score only, and iii) SB2 outputs the additional one-hot encoding of the selected mixture's 's ID in addition to SB's output. Note that PC makes sense in that we have the supports of the input set covered by all H mixtures together, while SB strategy is motivated
simply concatenates
5

---

Published as a conference paper at ICLR 2022
from the conventional practice in EM learning where EM algorithm multiple times with differential initial iterates and take the one with the highest likelihood score. SB2 additionally links the output to the specific mixture parameters selected.
Unsupervised learning. DIEM can be typically used in supervised learning, where the output of (SE is fed into subsequent layers (including the classification head) in a prediction network, and updated in an end-to-end learning manner. Beyond this, DIEM can be trained with only sets (S,,..., ] without target labels. This naturally done by marginal data likelihood maximization for mixture learning, where we maximize Σi= log (og)(S with respect to the prior parameters A. We can perform EM for optimizing the objective (in this case, the off-line EM until convergence with the entire training sets, without the auto-diff backprop mode), or alternativey the k-means clustering. The latter leads to the unsupervised learning algorithm similar to that of OTKE (Mialon et al. 2021
5 RELATED WORK
Set representation learning, Prior to the attention-base approaches, the neural statistician (Edwards & Storkey 2017) aims build model that learns the statistics of an input set, while DeepSet (Zaheer et al. 2017) has the element-wiss -isee nonlinear transformation followed by simple mean or max pooling to aggregate features. Although the universality of this network architecture was partially proved, the
limitation was found (Wagstaff et al.) 2019).. and network structures more suitable for exchangeability and permutation invariance were suggested (Garnelo et al. 2018 Bloem-Reddy & Teh 2019).
Attention-based methods. Self attention was shown to be very successful for modeling set-structured data (Vinyals et al., [016] Yang et al.) 2018 [sse et al.) 2018 Kim et al. 2019). While there were several sophisticated attempts to reduce the quadratic computational cost of self attention (Wang) cet al. 2020 Kitaev et al.] 2020), the inducing-poine approaches that introduce learnable parameters queries are especially suitable for set inputs (Lee al.) 2019a Pritzel et al. 2017: [Skinnis et al.) 2020 Mialon et al. 2021). The attention mechanism was analyzed rigorously in ( (Tsai et al. 2019) providing unifying kernel interpretation. Applying multiple EM steps in our DIEM is architecturally similar to the recent Perceiver model (Jaegle et al. 2021) with multiple latent-attention layers.
Mixture- based representation. There was an attempt to define the similarity between two images (image as set of fixed feature vectors, e.g., Fisher vectors (Perronnin et al. (0100) using mixture models about a decade ago (Liu & Perronnin, 2008). Similar ours, they represent each image as MAP estimated mixture density, and the similarity between two images is measured as a divergence between the corresponding mixture densities. However, instead of treating the pipeline of whole EM steps as a differentiable object and learning the prior parameters as in our DIEM, they merely apply single MAP-EM step to obtain fixed density for each image, where the initial mixture for EM also fixed from global offline learning with entire training images.
6 EVALUATION
We evaluate our DIEM model empirically on two different types of tasks: counting and clustering problems (Sec. 6.1
6.2) to verify the model 's capability of learning general set representations by modeling interaction between set elements, ii) large-scale biological sequence classification and NLP to test the performance of the proposed model on real-world problems in both supervised and unsupervised settings.
tasks (Sec. 6.3 6.4 6.5
The hyperparameters in our DIEM include: p (the mixture order), H (the number of heads), & (the number of EM steps). (prior impact), and the multi-head pooling strategy (either of PC, SB, or SB2). We report the results of the best combinations that are selected by cross validation. The empirical study on the impact of these hyperparameters the mixture modeling: learn both means and covariances for counting, clustering, and SCOP 1.75 datasets, while covariances are fixed for SST-2 and DecpSEA.
summarized Appendix CC4] The other option is
6.1 OMNIGLOT UNIQUE CHARACTER COUNTING
This task, originally devised in (Le et al. 2019a), is useful to verify capability of modeling interaction between set elements. From the OMNIGLOT dataset (Lake et al. 2015) which contains 1,623

---

Published as a conference paper at ICLR 2022
Table 1: Unique character counting on OMNIGLOT. The averaged 0/1 accuracy results.
Method
Small set
Large set
DeepSet (Zaheer al. 2017
0.4617 0.0076
0. 0.1927 0.0091
Dot-prod Attn Yang et al. 2018
Ilse et al.
2018
0.4471 + 0.0076
N/A 0.3191 + 0.0050
SefTransformer (SAB+ + PMA)
[ee et al.
2019a
0.6037 + O. 0.0075
OTKE (Mialon al.
2021
0.5754 0.0130
0.3352 0.0098
DIEM (Ours)
0.7153 0.0067
0.4440 + 0.0069
different handwritten characters with 20 examples each, we build set by randomly selecting n images, and the task is to predict the number of unique characters in the set. Following the data construction protocol from (Lee et al. (2019a) we split the 1,623 characters into training/validaationnnss splits so that the test dataset contains only unseen characters. Specifically, we randomly choose the number of images n uniformly from { [Nmim..... N Nmaa): J: then choose the number of unique characters randomly from {Cmins...,,?? , Then we randomly sample c different classes (characters), and sample n images from those classes while ensuring that all classes have at least one image. We form two experimental settings: i) Small set (Niin 6, Nuaz 10, Cmin 1) originally used in (Lee et al. 2019a) and ii) (a more challenging) Large set (N (Nmin = 11, N max = 30, Cmin = 5).
For the models, we adopt similar feature extractor architecture as (Lee et al. 2019a): first apply four Cony-BN-Red
layers to each (element) image to have feature representation <(z), then perform the set embedding emb(S) whose output is fed into fully connected layer return the Poisson parameter A. The final loss function is the negative Poisson log- log-likelihood - log p(c|A). The batch size is 32 (sets) for both datasets.
Results. We report the 0/1 prediction accuracy averaged over 10 random runs in Table 1 OTKE, after validation, uses (p 100, H = 1) for the small set and (p = 50, H - 2) for the large set. Our DIEM takes (p- 20, H = 3, k = 2. 10 3, PC) for the small set and (p = 50, H = 2, k = 2, r = 10 PC) for the large set. For both datasets, our DIEM outperforms the competing approaches by large margin. Unlike OTKE, we have prior regularization and multiple EM steps which appear to be important to improve the performance.
6.2 AMORTIZED CLUSTERING
The task of amortized clustering is to learn function S(S) - O where S = fx1,-.... -........... Fn is (data) set of n points, and θ = {ax m, C} is the parameters of the mixture of Gaussians, p(z; Q) = 2zyN((;; mj, Cÿ), best fit to the data. We follow the setup similar to (Lee al. 2019a) where instead of providing the supervised data of pair ((5, ))), we are only given the set data S, and the loss function is the negative log-likelihood logg(( ()) Following (Lee et al. 2019a), we form two datasets: i) Synthetic dataset where n 2D data points are sampled from a randomly generated mixture of four Gaussians, and ii) CIFAR- 100 images where we first sample four classes out of 100, and sample n images from the classes to form a set S. In both cases n E {100, 500). For CIFAR-100, each image represented as a 512-dim vector from a VGG network pre-trained on the training set (Simonyan & Zisserman 2014).
For our DIEM (and OТKE), we first apply a simple linear layer transformation to have features (zz). Then the output of the set embedding emb(S) is fed into fully connected layer to have the Gaussian mixture parameters ④ For the competing DeepSet and SerTransformer models, we faithfully follow the network structures in (Lee et al. 2019a). The details are summarized in Appendix C.2
Results. The results are reported in Table 2 where the test log-likelihood scores and the adjusted rand index (ARI) are shown for the synthetic and CIFAR- 100 datasets, respectively. As references, we also report the oracle performances: (Synthetic) the true mixture used to generate the data, and CIFAR-100) the full offline EM-trained mixture model. Although SetTransformer greatly outperforms DeepSet, our DIEM exhibits further improvement over SerTransformer by large margin. Relatively poor performance of OTKE may be due only one-step EM update, compared to our DIEM that used h: = 3 EM steps. For instance, DIEM also showed degraded performance when k is used (Fig. 10 in Appendix C.4). Another reason of success of DIEM is the prior-induced
regularization (r 0.01) whereas OTKE has no such regularization (Fig. I in Appendix
This is the typical sctup for the research topic called deep amortized clustering LLe et al. 2019h Pakman 2020 Genevay et al. 2019) which is different from our general set representation learning
7

---

Published as a conference paper at ICLR 2022
Table 2: Amortized clustering on 2D synthetic data and CIFAR-100. For CIFAR- 100, since it was difficult to obtain exactly the same pre-trained VGG network as in (Lee et al. 2019a), we ran all models with our own pre-trained VGG network (about 71% test accuracy compared to 68.54% from (Lee et al. 2019a)).
Synthetic
CIFAR-100
Method
I Log-likelihood f
Adjusted rand index 1
Oracle
-1.4726
0.9842
DeepSet (mean pooling) Zaheer
2017
1.7606 + 0.0213
0.5736 + 0.0117
DeepSet (max pooling) Zaheer et al.
2017
1.7692 + 0.0130
0.5463 + 0.0154
Dot-prod Attn Yang et al. 2018 Isse ct al.
2018
-1.8549 + 0.0128
N/A
SelTransfomer (SAB PMA) Lee et al. 2019a
-..115 + 0.0046
0.9246 + 0.0113
SerTransformer (ISABI6 6 + PMA) (Lee et al 2019a
1,5000 = 0.0068
0.9381 + 0.0122
OTKE (p- = 4, = 5) Mialon et al. 2021
1.7803 + 0.0028
0.8207 0.0074
DIEM (p
5, k 3, r 0.01. SB2) (Ours)
-1.4873 + 0.0018
0.9770 + 0.0019
Table 3: SCOP 1.75 classification accuracies (top 1/5/10) for unsupervised and supervised learning.
Method
Unsupervised
Supervised
DeepSF HHou et al.
2019
N/A 81.8 0.8 92.8
73.0/ 90.3 94455
CKN
Chen et al.
2019a
0 0.2
84.1 0.1 0.2/
RKN
Chen a..
20196
N/A
85.3 9.2
SefTranstormer [Lee al.
2019a
N/A N/A
79.2 14245
Rep-the-Set (Skianis
2020
84.5 0451006///444444//9/5/0
OTКE (Mialon et al.
2021
85.8 + 0.2 95.3 + 96.
DIEM (Ours)
86.4 /9. 97.1 0.1
90.5
6.3 PROTEIN FOLD CLASSIFICATION TASK ON SCOP 1.75
Protein fold classification is the well-known task in bioinformatics where the goal is to predict the fold class for given protein sequence. We use the preprocessed data of the SCOP version 1.75 and 2.06 from (Hou et al., 2019 Mialon et al. (02I), which consists of 19,245 sequences (14.699/2,013 trining/validattin from SCOP 1.75 and 2,533 test from SCOP 2.06). Each input protein is sequence of amino acids whose length ranges from tens to thousands, and each amino acid is represented as a 45-dimensional vector based on the PSSM, secondary structure, and more. It is multi-classification task to predict one of the 1,195 different folds to which a protein sequence belongs.
For the models, we adopt the overall network architecture from (Chen al. 2019a Mialon et al. 2021): each input sequence is treated as a set S, each set element ((x) is a learnable Gaussian kernel mapping on 10-mers via Nyström RKHS approximation with anchor points (we use == 1024 for unsupervised learning, and L = 512/ 128 for supervised learning), and the output of the set embedding emb(S) is fed into the final softmax linear classification layer. For the competing methods, we replace the set embedding emb( S) by: the global mean pooling layer for CKN (Chen al. 2019a), the OTKE layer (Mialon et al. 2021). and our DIEM layer. In the unsupervised learning. we sequentially train-then-fix the first two layers without class labels: train (fx) via k-means, fix it, do unsupervised training of emb(S), and fix it, after which the last classification layer is trained with class labels. We use the k-means clustering for unsupervised learning. In the supervised learning case, the whole network is trained and-to-end.
Results. Classification accuracy (top 1/5/10) for unsupervised and supervised learning are shown in 3 In our DIEM, we fix the covariances V as identity and perform h: = OT E-step for the unsupervised learning. while we learn priors for V and take ho = 2 regular E-steps for the supervised learning. The details of the hyperparameters are summarized in Table [2]le Appendix C.3] The result shows that our approach outperforms all ttate-of-the-arts for both learning settings. Compared to OTKE, the increased accuracy of our DIEM can be mainly attributed to the prior-induced model regularization (T 10 for unsupervised learning) and multiple EM steps 2 for supervised learning). See also Fig. 12 in the Appendix for the hyperparameter impact analysis.
6.4 TASK: SENTIMENT CLASSIFICATION ON SST-2
The dataset (Socher et al., 2013) consists of 70,042 movie reviews with positive/negaatie binary sentiment. We follow the experimental protocol from (Mialon et al. 202T) where the original 67,349
8

---

Published as a conference paper at ICLR 2022
Table SST-2 sentiment classification accuracies for unsupervised and supervised learning.
Method
Unsupervised
Supervised
BERT [CI.S] emhedding (Devlin al.
2019
84.6 + 0.3
90.3 + 0.1
BERT mean pooling (Deviin et al. 2019
85.3 + 0.4
90.8 4 0.1
SelTransformer (Lee ct al.
2019a
N/A
87.9 0.8
Approximate Rep-the-Sct Skiani ct al.
2020
N/A
86.8 + 0.9
Rep-the-Set (Skianis ct al.
2020
N/A
87.1 + 0.5
OTKE (Mialon et al.
2021
86,8 +: 0.3
88.1 + 0.8
DIEM (Ours)
87.6 + 0.2
88.7 + 0.4
training data are split into 80%/20% training/validation sets, and the 872 validation reviews form test set. This is because test evaluation the original test set requires online submission in the GLUE (Wang et al. 2019) leaderboard. Similar to (Mialon et al. [021]. we adopt the -68-dimensional word vectors from the pre- pre-trained BERT model (Devlin al. 2019 Wolf et al. 2019),, and use the
Gaussian RKHS mapping (zz) approximated by the Nyström method with 2048 (unsupervised) or
64 (supervised) filters. Overall network architectures and learning options are also similar to (Mialon)
ett al. 202T)
Results. Table 4 summarizes the results. For the unsupervised learning, the proposed DIEM attains the best accuracy, even outperforming the pre-trained BERT model with the popular [CLS] embedding or mean pooling. Unlike OTKE that takes large number of parameters (p 300, H 1). our DIEM has far fewer parameters (p 3, H = 2, k = 1 with PC multi-head pooling). When we trained OTKE with the same complexity architecture, only attained accuracy 86.6 = 0.5. The improved accuracy of DIEM over OTKE is mainly due to the prior-induced regularization (we use = 10-3) and the regular E-step (with the OT E-step, the performance is on par, 86.3 0.8). For the supervised learning, both DIEM and OTKE has p 30, H 4. and the performance difference mainly comes from the regular E-step taken by DIEM (T = 10 6 SB). Although DIEM performs comparably well, it falls short of the fine-tuned BERT model. As alluded in (Mialon et al. 202T), the set embedding approaches may not be very attractive in this case due to short sentences (small cardinality sets).
6.5 CHROMATIN PROFILE DETECTION ON DEEPSEA
Finally, we test DIEM on the large-scale DeepSEA Table 5: Chromatin profile detection
dataset (Zhou & Troyanskaya 2015) with about five mil- results on DeepSEA.
lion genomic sequences. It is multi-label classification task
Method
auROC
auPRC
to predict 919 chromatin profiles. As in (Mialon et al.) (021). the overall network architecture follows the ID Conv layer from DeepSEA, while the pooling layer after the Conv is replaced by the DIEM layer. The RKHS feature mapping is
DeepSEA
0.933
0.342
OTKE
0,936
0.360
DIEM (Ours)
0.936
0.360
not used (identity mapping), and the Gaussian positional encoding was adopted. The supervised learning results are shown in Table where we have marginal improvement over DeepSEA, and the same performance OTKE using the (p 64, H 1, k 2, 1) hyperparameter choice.
7 CONCLUSION
In this paper we proposed novel differentiable EM model for set representation learning. Our model is built from the perspective of fittingg a Gaussian mixture model to the set data that are viewed as i..d. samples, which offers more flexibility and prior-induced model regularization in principled Bayesian manner. The proposed model is also shown to generalize the recent set embedding models based on optimal transport and attention, leading a computationally efficient model with superb performance on tasks in bioinformatics and NLP. The mixture fitting perspective can potentially solve other important problems in set representation learning. For instance, deciding the number of components/feferences can be translated into the mixture order selection problem, and tackled by well-known information criteria methods. This will be left as interesting future research.
3 As in Mialon et al. 2021 the cost matrix in OT is multiplied by the RBF-like position similarity kernel Pt, exp(- expp-(x// h/pl/a² Alternatively tried random feature extension (Rahimi & Recht 2008 for the position index i, and concatenated them o(x.), however, the results were more less the same.

---

Published as a conference paper at ICLR 2022
REFERENCES
Dzimitry Bahdanau, Kyunghyun Cho, and Joshua Bengio. Neural machine translation by jointly learning align and translate. In International Conference Machine Learning. 2015.
B. Bloem-Reddy and Y. W. Teh. Probabilistic symmetry and invariant neural networks, arXiv preprint arXiv:1901.06082 2019.
Dexiong Chen, Laurent Jacob, and Julien Mairal. Biological sequence modeling with convolutional kernel networks. Bioinformatics, 35(18)33294 3302, 2019a.
Dexiong Chen, Laurent Jacob, and Julien Mairal. Recurrent kernel networks. In Advances in Neural Information Processing Systems, 2019b.
Marco Cuturi. Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances. In Advances in Neural Information Processing Systems, 2013.
A. P. Dempster, N. M. Laird, and D. B. Rubin. Maximum likelihood from incomplete data via the em algorithm. Journal of the Royal Statistical Society. Series (Methodological) 39(1):1-33) 1977.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In In Proceedings of the North American Chapter of the Association for Computational Linguistics (NAACL), 2019.
Harrison Edwards and Amos Storkey. Towards neural statistician. In International Conference on Learning Representations, 2017.
M. Garnelo, D. Rosenbaum, C. J. Maddison, Ramalho T., D. Saxton, M. Shanahan, Y. W. Teh, J. Rezende, and S. M. A. Eslami. Conditional neural processes. International Conference on Machine Learning, 2018.
Aude Genevay, Gabriel Dulac-Aroldd, and Jean-Philippe Vert. Differentiable Deep Clustering with Cluster Size Constraints. arXiv preprint arXiv: /9100099000 2019.
Jie Hou, Badri Adhikari, and Jianlin Cheng. DeepSF: Deep convolutional neural network for mapping protein sequences to folds. Bioinformatics, 34(8):1295-1333. 2019.
M. Ilse, J. M. Tomezak, and M. Welling. Attention-based deep multiple instance learning. In International Conference Machine Learning, 2018.
Andrew Jaegle, Felix Gimeno, Andrew Brock, Andrew Zisserman, Oriol Vinyals, and Joao Carreira. Perceiver: General Perception with Iterative Attention. In International Conference on Machine Learning, 2021.
H. Kim, A. Mnih, J. Schwarz, M. Gamelo,, A. Eslami, D. Rosenbaum, O. Vinyals, and Y. W. Teh. Attentive neural processes. In International Conference on Learning Representations, 2019.
D. P. Kingma and L. J. Ba. Adam: A Method for Stochastic Optimization, 2015. International Conference on Learning Representations (ICLR).
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The Efficient Transformer. In International Conference on Learning Representations, 2020.
B. M. Lake, R. Salakhutdinov, and J. B. Tenenbaum. Human-level concept learning through proba- bilistic program induction. Science, 1332- 2015.
Juho Lee, Yoonho Lee, Jungtaek Kim, Adam Kosiorek, Seungjin Choi, and Yee Whye Teh. Set trans- former: framework for attention-based permutation-invariaan neural networks. In Proceedings of the 36th International Conference on Machine Learning, pp. 3744-3753, 2019a.
Juho Lee, Yoonho Lee, and Yee Whye Teh. Deep Amortized Clustering. arXiv preprint
arXiv: /909. 13433, 2019b.
Yan Liu and Florent Perronnin. A similarity measure between unordered vector sets with application to image categorization. In JEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2008.
10

---

Published as a conference paper at ICLR 2022
Grégoire Mialon, Dexiong Chen, Alexandre d' Aspremont, and Julien Mairal. A trainable optimal transport embedding for feature aggregation and its relationship to attention. In International Conference Learning Representations, 2021.
Ari Pakman, Yueqi Wang, Catalin Mitelut, JinHyung Lee, and Liam Paninski. Neural Clustering Processes. In International Conference Machine Learning, 2020.
F. Perronnin, J. Sánchez, and T. Mensink. Improving the fisher kernel for large-scale image classifica- tion. In European Conference on Computer Vision, 2010.
A. Pritzel, B. Uria, S. Srinivasan, A. Puigdomenech. O. Vinyals, D. Hassabis, D. Wierstra, and
C. Blundell. Neural episodic control. In International Conference on Machine Learning, 2017.
Ali Rahimi and Benjamin Recht. Random features for large-scale kemell machines, 2008. In Platt, J.
C., Koller, D., Singer, Y., and Roweis, S. T. (eds.), Advances in Neural Information Processing Systems 20.
Bernhard Schölkopf and Alexander J. Smola. Learning with Kernels. MIT Press, Cambridge, 2002.
K. Simonyan and Zisserman, A Very deep convolutional networks for large-scale image recognition. arXiv preprint arXiv:1409. 1509, 2014.
Konstantinos Skianis, Giannis Nikolentzos, Stratis Limnios, and Michalis Vazirgiannis. Rep the set: Neural networks for learning set representations, 2020. AI and Statistics (AISTATS).
Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D. Manning, Andrew Ng, and Christopher Potts. Recursive deep models for semantic compositionality over a sentiment treebank. In Empirical Methods in Natural Language Processing (EMNLP), 2013.
Yao-Hung Hubert Tsai, Shaojie Bai, Makoto Yamada, Louis-Philippe Morency, and Ruslan Salakhut- dinov. Transformer dissection: unified understanding of transformer's attention via the lens of kernel. In Empirical Methods in Natural Language Processing (EMNLP). 2019.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems, 2017.
Cédric Villani. Optimal Transport: Old and New. Springer, 2008
O. Vinyals, Bengio, and Kudlur. Order matters: Sequence to sequence for sets. International Conference Learning Representations, 2016.
E. Wagstaff, F. B. Fuchs, M. Engelcke, I. Posner, and M. Osborne, On the limitations of representing functions on sets. arXiv preprint
2019.
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R. Bowman. GLUE: A multi-task benchmark and analysis platform for natural language understanding. In International Conference on Learning Representations (ICLR), 2019.
Sinong Wang, Belinda Z. Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-Attention with Linear Complexity. arXiv preprini arXiv:2006.04777 04768, 2020.
Christopher K. I. Williams and Matthias Seeger. Using the Nyström method to speed up kernel machines. In Advances in Neural Information Processing Systems. 2001.
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Remi Louf, Morgan Funtowicz, and Jamie Brew. Huggingface's 's transformers: State-of-the-art natural language processing. arXiv preprint 191.77191, 2019.
B. Yang, S. Wang, A. Markham, and N. Trigoni. Attentional aggregation of deep feature sets for multi-view 3D reconstruction. arXiv preprint arXiv:1808.0077 2018.
M. Zaheer, S. Kottur, S. Ravanbakhsh, B. Poczos, R. R. Salakhutdinov, and A. Smola. Deep sets. In Advances in Neural Information Processing Systems, 2017.
Jian Zhou and Olga G. Troyanskaya. Predicting effects of noncoding variants with deep learning- based sequence model. Nature methods, 12(10):931 1934, 2015.
11

---

Published as a conference paper at ICLR 2022
Appendix
We provide some formal derivations, experimental details and additional results that are skipped in the main paper.
A COMPLEXITY AND RUNNING TIME
In this section we analyze the complexity, especially the number of parameters and time complexity for forward pass, of our DIEM. We also compare the wall clock running time for OTKE (Mialon et al. 202T) and our DIEM with different numbers of EM steps.
A.1 COMPLEXITY
Notation. First, we recall the notation: The input set S = {x1....,, Fn contains n d-dimensional feature vectors 5; E IRd We consider the OTKE model with p references [= (each 2j E IR²) as parameters. To be comparable, DIEM has p Normal- Inverse- Inverse-Wishart (NIW) parameters (ız, 2,
(Hy E Rd, D is a (d x d) diagonal matrix), each of which defines prior for each component (22, V2) in the Gaussian mixture. For simplicity, we assume single-head representations (i.e., H = 1) for both models, since multi-head (H heads) representation increases all compute and memory linearly by H.
Numbers of parameters. The numbers of parameters are: pd (OTKE) vs. 2pd (DIEM), Hence DIEM requires only twice as many parameters as OTKE.
Time complexity. First, one forward pass in OTKE consists of: [01] computing the cost/kernel matrix C which is (n x p) with C22- = (x,,, 2,), [O2] solving the OT problem (2) by the Sinkhorn- Knopp (SK) matrix scaling algorithm, and [O3] performing attention-weighted sums (1) to have the final embedding. The least amount of time for [O1] is O(npd) to inspect each element of x; and 2j once (e.g.. linear kernels and squared- -iatteddisttnee based kernels require O(npd)). Denoting the number of SK iterations by M, [O2] takes OnnpM), while [O3] requires O ((ppd). Thus, the time complexity for one forward pass in OTKE is ((p M d)).
For DIEM with & EM steps, [D1] each E-step (10) amounts to computing the Gaussian likelihoods, and additionally running the SK algorithm if the OT E-step is used. [D2] Each M-step (12) mainly consists of the weighted sums. [DI] takes ((ppd) for the regular E-step, and ('pp(M + d)) for the OT E-step, while [D2] requires only (mpdd. Hence, the time complexity for one forward pass in our DIEM is (kru((M + d)), only As times as much as OTKE.
We summarize it in Table 6 The actual wall clock running times compared in the next section.
A.2 (FORWARD PASS) RUNNING TIME
To measure the forward pass time for OTKE and our DIEM, we consider first the large-scale SCOP 1.75 dataset: the input feature dimension d = 512 (from Gaussian RKHS) and the set cardinality 1091. To see the impact of the set cardinality, we inject random noise features to the sets to make twice and four times larger sets (n = 2182 and n = 4364). For fair comparison, both OTKE and DIEM use one head (H 1) with the number of references (or mixture components) p = 100. The batch size is 128. We run all models on the same machine, Core i7 3.50GHz CPU and 128 GB RAM with a single GPU (RTX 2080 Ti).
Fig. 3 compares the average per-batch forward pass time for OTKE and DIEM with different EM steps (k- = to 4). Although DIEM has an option to choose the E-step type (either regular E-step or OT E-step), we take all OT E-steps to have fair comparison with OTKE. (For the running time comparison between these two E-step types, please refer to Fig. As shown in Fig. 3 increasing the number of EM steps in our DIEM incurs only small amount of computational overhead. The impact of the set cardinality n on the running time exhibits similar trends for both OTKE and DIEM.
In Fig. 4 we compare running time of OTKE, SerTransformer (Lee et al. 2019a), and DIEM (with = EM steps). Here we use a minimal network structure for SetTransformer: one SAB block for the encoder, one PMA and one SAB block for the decoder. For larger sets (x2 and x4) we reduce
12

---

Published as a conference paper at ICLR 2022
Table 6: Numbers of parameters and time complexity of OTKE (Mialon et al. 202T and DIEM. We assume single-head models (H = 1). DIEM follows k EM steps. Here, = imput set cardinality, imput feature dimension, p number of references (OTKE) or mixture components (DIEM), and M number of SK iterations. Time complexity for one forward pass.
OTKE
DIEM
Number of parameters
pd
2pd
Time complexity
( + dd)
(kmpp M + d))
15
14
Forward (msec)
13
10
9-
DIEM k=1 (1)
DIEM k=2 (11)
8-
k-3 (111)
k=4 (1111)
OTKE
1091(x1)
X2 cardinality (n)
x4
Figure 3: Average (per-batch) forward pass time (in milliseconds) for OTKE and DIEM with different EM steps (k 1 to 4) on the SCOP 1.75 dataset. In addition to the original set cardinality n 1091 (denoted by x I on the left corner), we increase n by two (x2) and four times (x4). To be fair with OTKE, we use OT E-steps during the & EM steps, denoted by sequences of I's in the legend (e.g.. k = 4 (1111) indicates that four OT E-steps are used during k = EM steps).
the batch size for SerTrrnsformer to fit into the GPU memory. We need not reduce the batch size for OTKE and DIEM they fit into the GPU memory for all three cardinalities.
The result shows that OTKE and DIEM are scalable to large sets whereas SerTransformer takes significantly longer time, which well aligned with the fact that self-attention suffers from quadratic computational cost. In Fig. 5 we compare the running time between OT E-steps and regular E-steps in our DIEM. The maximum number of Sinkhorn- Knopp iterations for solving the OT problem is set to 10. While the regular E-step is consistently faster than the OT E-step, the increased overhead of the OT E-step seems to be minor compared the overall running time.
For smaller datasets, we consider the OMNIGLOT counting (large set) and CIFAR 100 clustering datasets. The running time of OТКE and DIEM on these two datasets depicted in Fig. 6 For the OMNIGLOT large set, we set 50, H = 2, d = 64 for both models where n E {11,-.. .1...... 300 and batch size is 32 sets. For CIFAR-I00 clustering, we test with p = 4, H = 5, d = 512 for both models where π. E {100, ... 500 and batch size is 10 sets. Our DIEM takes regular E steps. We vary the number of EM steps & in DIEM from 1 to 5. The additional forward pass time incurred by extra EM steps in DIEM looks permissible, the overall running time of DIEM being constant factor comparable to that of OTKE.
B MATHEMATICAL DERIVATIONS
In Appendix B.1 we provide detailed derivations for MAP EM update equations [0 in the main paper. In Appendix B.2] formally show that the satuuiinn-piit attention module of SerTransformer (Lee et al. 2019a) known as PMA, is actually a special case OTKE (Mialon et al. 2021].
13

---

Published as a conference paper at ICLR 2022
(111)
35
SetTransformer
30
25
20
Forward
1091(x()
x2
Set cardinality (n)
Figure 4: Running time comparison among OTKE, Set and DIEM. Average (per-batch) forward pass time (in milliseconds) on the SCOP 1.75 dataset with three different set cardinalities (original, x2, and x4). DIEM takes three EM steps = 3 (111) indicates that three E-steps are used during k 3 EM steps). SetTransformer uses minimal network structure: one SAB block for the encoder, one PMA block and one SAB block for the decoder.
k=3 (000); Regular E-steps
14
k=3 (111) OT E-steps
Forward-pon time (msec)
12
10
T
8.
6
1091(x1)
Set cardinality (n)
Figure 5: OT E-steps vs. regular E-steps our DIEM. Average (per-batch) forward pass time on SCOP 1.75 dataset with three different set cardinalities (original, x2, and x4). We compare DIEM with = (111) (all OT E-steps) with (000) (all regular E-steps).
OMNIGLOT counting (large set)
CIFAR-100 clustering
10
30
10
Forward-pear time
pass
35
30
OTKE (K=M(
k-2
k-4
k=5
OTKE
DIEM [K=]]
k-2
kG5
OTKE and DIEM (k=1-5)
OTKE and DIM (k=1-5)
Figure 6: Running time comparison between OTKE and DIEM with k = 1
EM steps on
(Left) OMNIGLOT TTO counting (large set) and (Right) CIFAR-100 clustering.
14

---

Published as a conference paper at ICLR 2022
B.1 MAP EM DERIVATIONS
E-step is optimization of the lower bound in (9) with respect to g(jl)) with 0 fixed. Hence it is identical to the maximum likelihood E-step for the lower bound of (4). Now derive the M-step, maximizing the lower bound in (9) with respect to with ((ili) fixed. That is,
arg max L( == (()A( + Σ ((11)) log p(..())..
(14)
The log-prin pegpprior term of the Dir-NIW density can be expanded as ("=e" refers to equality up constant):
log ((]AA) =e
1/ + + 2
log VV 2 (2) - my) V (2) - -15) TT
2
+ (m (1-1))
1)_ log 73-
(15)
The weighted log-likelihood term in (14 can be written as:
Σ
((jli) logg(xt.))))
Σ a(ib log
1
log [Vg]
2 (z; (+)==,)[0 V, (z; 2))) (16)
2
Combining
15 and (16,, we take derivatives with respect to 0 and solve for 0 as follows:
0L 1 + E. a(il)) const
E + +1-
(17)
vn, Jj Jj
n n p. (1-
OL
E,q(i1ir; +
-V, (zz 14) - Vv Σ (011)(z - x.)) = 0 = 2) Ec0(/i) +X
(18)
OL
((u+d+ d d 2+ E ((11)) V, M22 11)² 2- -E, Eg(11((0 - =3)²²
0V,1
V = M(2 22)2
(19)
+
Ex0 - + (2/)
(20)
Σ. ++
where the equality between 19 and 20 be easily verified from the optimal 2 in (18
B.2 PMA OF SETTRANSFORMER IS SPECIAL CASE OF OTKE
The PMA (inducing-poits attention) module of SetTransformer (Lee et al., (019a) is an attention block that incorporates a fixed number of learnable parameters as queries, and plays an important role in SerTransformer to have a fixed size vector representation for variable-length input set. Formally, PMA module has learnable p vectors = {z....., } and for given set S {xI,.... Øn). performs the (dot-product) attention operation with each zj query and S as keys and values. Hence each 2j is transformed as follows:
exp (W22,) ( (Wz))/Va
E exp ((W (Wx:)/Va) Wx.
(21)
where W2, and W are the weight matrices for queries and keys/values, respectively. Since 2, itself a learnable parameter vector, it is redundant to introduce W and we can subsume it under 2j, which leads to
exp (z] (Wz.)/VV)) Wz.
PMA:
(22)
EL exp (3] ( (wx)/Va)
Next, we build a special OTKE module by having the RKHS feature mapping ((z;) = Wz,, which corresponds to the linear kernel k(z,,, 2y) = 2] ((xx) 2 Wa.. Furthermore, we consider the OT
15

---

Published as a conference paper at ICLR 2022
Table 7: (OMNIGLOT Small) Unique character counting the OMNIGLOT small The averaged 0/1 accuracy results. In DIEM, k 2(00) (or k = (()) means 2 (or 1) EM steps used both with regular E-steps.
Method
Small set
DeepSet 7ahherr al. 2017
0.4617 0.0076
Dot-prod Attn (Yng al. 2018 llse al.
2018
0.4471 0.0076
SelTransformer (SAB + PMA)
Lee ct al.
2019a
0.6037 0.0075
OTKE (p 100, H
Mialon
2021
0.5754 0.0130
OTKE 50, H 2) Mialon al.
2021
0.5686 0.0187
OTKE (p 20, Н = 3) Mialon et al.
2021
0.5216 0.0182
DIEM (p 20,
& 1(0), 10". SB) (Ours)
0.6304 0.0056
DIEM (p 20, H & 2(00), 10- (B) (Ours)
0.6514 0.0074
DIEM (p 20, 11 3, & 1(0), - SB2) (Ours)
0.6229 = 0.0114
DIEM (p 20,
K 2(00), T 10 SB2) (Ours)
0.6065 0.0087
DIEM (p 20, H 3, & 1(0), 10", PC) (Ours)
0.6083 0.0194
DIEM (p 20, H 3, & 2(00), r 10 PC) (Ours)
0.7153 0.0067
problem between Z and S without the second constraints Σ/=1 Qi = 11/n, vi. That is,
min Σ -c s.t.
Vj.
(23)
f=1
with Cy = -k((r;,
= -2 Wr We can solve 23
using the Lagrange multiplier,
L: b: Co cH(Q)
Σ M
Qis- 1/p)
(24)
We set the derivative to 0 and solve for Qıj as follows
ac
aQ Ci e(log Qi + 1) X, =
Qij exp ( Casle+
(25)
for some constant aj. Applying the constraints Σm= Qij 1/p determines ay and Qij as:
exp [2] (Wax)/c)
Qij " exp (2) (Wx:)) (Wx))/e)
(26)
p P1 exp Ci/ee)
By choosing the entropic regularizer trade-off 2j - vp E. (xx)(x;), we get
V and applying the OTKE embedding formula
exp (z] (Wa:)\VV WXi,
OТKE:
(27)
E exp (2] (Wz,)/vd)
which equivalent to PMA 22 up to constant factor. Hence PMA of SetTransformer can be viewed OTKE module without the equal element assignment constraints Σ Qij 1/n, Vi.
C EXPERIMENTAL DETAILS AND ADDITIONAL RESULTS
C.1 OMNIGLOT UNIQUE CHARACTER COUNTING
(Original) Small set. This experimental setup is the same as that of (Lee et al. 2019a). The sets of OMNIGLOT images are generated by: set cardinality characters & ~ (1....,n). The batch is composed of 32 sets. The results are summarized in Table 7
{6, 7, 8, 9, 10} and the number of unique
)
Large set. This large experiment is built from the sets generated by: set cardinality (11, 12, ... 30} and the number of unique characters k ~ n......n.. The batch is composed of 32 sets. The results are shown in Table
16

---

Published as a conference paper at ICLR 2022
Table 8: (OMNIGLOT Large) Unique character counting on the OMNIGLOT large set. The averaged 0/1 accuracy results. In DIEM, k = 2(00) means 2 EM steps used both with regular E-steps.
Method
Large set
DeepSet 7aheer al. 2017
0.1927 + 0.0091
Se(Transformer (ISAB32 PMA) Lee et al. 2019a
0.2456 + 0.0068
SetTransformer (ISAB64 PMA) Lee et al.
2019a
0.2836 0.0087
SedTransformrr (SAB + PMA) (Lee al. 2019a
0.3191 + 0.0050
OTKE (p = 50, H = 2) [Mialon et al. 2021
0.3352 0.0098
DIEM (p- 50,
2, k 2(00), r 10
SB2) (Ours)
0.3887 0.0071
DIEM (p= 50, H 2, k 2(00), r 10- 5, PC) (Ours)
0.4392 0.0056
DIEM (p 50, H 2, k 2(00), T 10 PC) (Ours)
0.4440 + 0.0069
DIEM (p- 50,
k 2(00), r 10 PC) (Ours)
0.3581 0.0121
Table Amortized clustering on synthetic data. The averaged test log- -og-liiklihood scores are shown. In DIEM, k k= 2(00) (or k 3(00)) means 2 (or 3) EM steps used all with regular E-steps.
Method
Synthetic
Oracle
-1.4726
DeepSet (mean pooling) Zahherr et al.
2017
- 1.7606 + 0.0213
DeepSet (max puoling) Zaheer ct al.
2017
- -1.7692 + 0.0130
Dot-prod Attn YYnng et al. 2018 lsse et al. Se(Transfommr (SAB PMA) Lee al.
2018
-1.8549 0.0128
2019a
- -1.1145 + 0,0046
SecTransformer (ISAB16 + PMA) (Le et al. OTКE (p- 4, HI 5) Mialon
2019a
- -1.5009 + 0.0068
2021
-1.7803 0.0028
DIEM (p 4, II 5, 3(000),
10
SB2) (Ours)
- -14483 + 0.0021
DIEM (p =0 11 = 5, = 3(000), r = 0.01, SB2) (Ours)
-1.4833 + 0.0018
DIEM (p = H = =5,k 3(000), r SR2) (Ours)
-1.4960 0.0039
DIEM (p- - H - 5, - 2(00), - 0.01, SB2) (Ours)
- -1.5082 + 0.0032
Network architectures and learning options. The feature vectors ((x) of images (set elements) are formed by applying four Conv(64,3,2.BB) ReB. layers. For OTKE and our DIEM the output of the set embedding emb((S) is fed into the fully-connected layers FC(64, - FC(I softplus) to return the (scalar) Poisson parameter. For Se(Transformer (bx) is fed into SAB(64,4) ( SAB(64,44) PMA(1,8.8) - FC (softplu) layers where SAB(d.h) means the set attention block with d/h units/- heads, and PMA(k;,d,h) is the Pooling Multihead Attention layer with & vectors and d/h units/heads. When the inducing-point attention is used, S SAB(64,4) blocks are replaced by ISAB(m, 64,44) that has i inducing points (m 32 or 64) with 64/4 units/heads. For DeepSet, ((x) is fed into FC(64,ReLJU) -C((4 Mcan- Pool- - FC(64. - FC(1 (.ooffll)))))) layers. For all models, we use the Adam optimizer Kingma& & Ba 2015) with learning rate 10 4 and batch size 32 sets, until 200K iterations.
C.2 AMORTIZED CLUSTERING
Synthetic 2D data clustering. The results are summarized in Table 9
Synthetic small set data clustering. We test the models on sets with smaller cardinality: N (50, 100} instead of the previous N ~ {100, .., 5003. The results are shown in Table 10
CIFAR-100 clustering. Results are shown in Table 11
Network architectures and learning options: For OTKE and our DIEM, we first apply a linear layer FC(dh.-) to the din-dimensional input vectors to have features (bxx), where dh = 128, din = 2 for the synthetic dataset and the small synthetic dataset, and dh = 256, din = 512 for the CIFAR- 100. The output of the set embedding emb(S) is then fed into FC(dh, dout = 4 + 4 din + 4 for a, m, and C. respectively, in the order-4 spherical-covariane mixture parameter output θ {a, mnmm.. C}. We also apply softmax and softplus layers for and C, respectively. For SefTransformer, the dimediimesioonal inputs are fed into SAB(dh,4) - SAB(dh,44 (M,) SAB(dh,4)- SAB(dh,4) FC(doul,-) layers, whereas for the inducing-point
FC(dout,-) where Gaussian
17

---

Published as a conference paper at ICLR 2022
Table 10: Synthetic small set data clustering. The averaged test log-likelihood scores are shown. In DIEM, k 3(000) means 3 EM steps used all with regular E-steps.
Method
Small Synthetic
DeepSet (Zaheer al. 2017]
2.0955 0.0054
Se(Transformer (ISAB16 PMA) Lee ct al.
2019a
1.5841 + 0.0077
SeTTransfomer (ISAB32 PMA) Lee al.
2019a
11.5277 0.0083
SetTransformer (SAB PMA) (Lee al.
2019a
-1.099 + 0.0068
DIEM (p = 4, = 5, = 3(000), r = 10 6, .SB2) (Ours)
- -1,4297 + 0.0027
DIEM (p 4, 5, 3(000), r 10-5,SB2) (Ours)
-1.308 0.0022
DIEM (p H, 5, k 3(000), 10 10"4,BB2) (Ours)
-1.4274 0.0015
DIEM (p = 41 = 5, k 3(000), r 10 SB2) (Ours)
-1.4301 + 0.0030
DIEM 4p=4,H- 5, 3(000), 10 SB2) (Ours)
-1.4253 0.0010
DIEM (p = 4, = 5, k 3(000). 10 SSB) (Ours)
-1.4343 + 0.0015
Table 11: Amortized clustering on 100. The averaged test adjusted rand index (ARI) scores shown (the higher the better). In DIEM, k 2(00) (or 3(00) means 2 (or 3) EM steps used all with regular E-steps.
Method
CIFAR- 100
Oracle
0,9842
DeepSet (mean pooling)
Zaheer el al.
2017
0.5736 0.0117
DeepSet (max pooling) 7aheer
2017
0.5463 0.0154
Se(Transformer (SAB PMA) [ee al.
2019a
0.9246 +: 0.0113
SerTransformer (ISAB 1 PMA) (Lee et al
2019a
0.9381 0.0122
OTKE II Mialon et 2021
0.8207 0.0074
DIEM (p H, = 5, 2(00), 0.01. SB2) (Ours)
0.9709 + 0.0013
DIEM (p - -4 = 5, = 3(000), r 0.01. SB2) (Ours)
0.9770 0.0019
DIEM (p H, 5, 3(000), 0.01, SB) (Ours)
0.9688 0.0022
DIEM (p H, = 5, A 3(000), 10 SB2) (Ours)
0.9680 + 0.0011
DIEM (p- = 4, H = 5, 3(000), 7 1.0, SB2) (Ours)
0.9630 0.0024
18

---

Published as a conference paper at ICLR 2022
Table 12: SCOP 1.75 classification accuracies (top 1/5/10) for supervised learning. The hyperparame- ter L in OTKE and DIEM indicates the number of Nyström anchor points for the approximation of Gaussian RKHS embedding. In DIEM, k 2(00) means 2 EM steps used both with regular E-steps. The hyperparameter L is the number of anchor points used in the Nyström RKHS approximation of the feature (zz) through the Gaussian kernel mapping on 10-mers (subsequences of length 10).
Method
Supervised
DeepSF (Hou et al.
2019
/3.0/90.3/
CKN (Chen ct al.
2019a
84.1 0.1/94.3 + 0.2/ 0.1
RKN
(Chen et al.
2019h
85.3 0.3/ 15/995000000999900 0.2//6.5 0.1
SerTranstormer
(Lee et al.
2019a
79.2 1.6/91.5-1
0.6
Rep-the-Set (Skiams al.
2020
84.5 0.6 1.4/990:::4499.--- 0.4
OTKE (p- 50, //
128) (Mialon et al.
2021
85.2 00==0.4/999 0.1.:.......5...6...0.00 0.2
OTКE (p 10. // -= 512) (Mialon ct al.
2021
88.7 0.7 0.1
DIEM (p 50, 1,K k 2(00), T 10
128) (Ours)
DIEM (p 50, II 1,K 2(00), r 10 8,- L 512) (Ours)
89.6 0.3/
attention, the first two SAB(dh.4) blocks are replaced by ISAB(m,dh.44) For DeepSet, the layers are composed of hig FC(di,RelLU) layers, Mean-Pool, and (ka - 1) FC(dh,,ReLU), followed by FC(dout,-), where Ky for synthetic and ka 6 for CIFAR-100, For all models, we the Adam optimizer with learning rate 10 3 (10 44 for CIFAR- 100) and batch size = 10 sets, until 50K iterations.
C.3 PROTEIN FOLD CLASSIFICATION TASK ON SCOP 1.75
Unsupervised learning. The best performing DIEM model found with (p 100, H
1, = 10 following the OT E-step. See also Fig. 12 for the results of different prior strength values.
Supervised learning. The results for several different hyperparameter choices are shown in Table 12
Network architectures and learning options. We faithfully follow (Mialon et al. 2021 for the overall network architectures and learning options. For instance, the Adam optimizer is used with batch size 128 for 100 epochs, and the learning rate initially 0.01 and halved if there is no decrease in validation loss for 5 consecutive epochs.
C.4 EMPIRICAL STUDY HYPERPARAMETER IMPACTS
We empirically analyze the impact of various hyperparameters used in our DIEM model. They summarized as follows.
OMNIGLOT (small). We first study the impact of prior strength 7 and multi-head pooling strategy.
and the multi-head shown in Fig. 7 We have the best accuracy when r is the middle value 10-3 compared to the extreme ones (10- lowest prior impact and 10 highest impact). Regarding the multi-head pooling strategy, PC (parameter concatenation) performs the best for small т, while SB2 (select best with the selected mixture ID one hot encoding) is slightly better than PC for 10". Next we test the impact of the number of EM steps and the prior impact т. For this we fix the multi-head pooling strategy as SB2 or PC. As shown in Fig. g having multiple k EM steps consistently outperform single EM step. Unlike i the single EM step models 10".
While fixing the number of EM steps as k = 2. we vary T E { [10 -6 10 100-3 -3 100)
pooling strategy E (P,,SB, SB2). The results
attained the best performance with the strongest prior impact
Synthetic clustering. We also test the impact of multi-head pooling strategy and prior strength on this synthetic clustering dataset for our DIEM model (p = 4, H 5, & = 3). As shown in Fig. 9 SB2 strategy outperforms the other two strategies for all 5 values. For SB2, the impact of appears to be minor, while for the other two strategies, we have opposite behavior; larger performs better for PC, and smaller т improves the performance for SB. Fig. 10 visualizes the impact of the number of EM steps, where the model is with = 4, H 5. 0.01. SB2). It clearly shows that having multiple EM steps significantly improve the performance over the single EM step. There is only a little improvement from & = to k
19

---

Published as a conference paper at ICLR 2022
0.72
I
PC
0.70
S82
Accuraty
0.68
0.66
0.64
0.62
0.60
T= 10
T
T= 100
r (prior impact)
Figure 7: Impact of the prior strength hyperparameter and the multi-head pooling strategy for OMNIGLOT counting (small). The number of EM steps k = 2 fixed.
tau=
tau=le-6
0.70
0.70
tau= 1e-3
tau=1e-3
tau=1e0
tau=1e0
0.65
0.65
0.60
0.60
0.55
0.55
0.50
0.50
k=1
==2
k=1
k=2
(the number of EM steps)
k (the number of EM steps)
Figure 8: Impact of the number of EM steps & and the prior impact т for OMNIGLOT counting (small). The multi-head pooling strategy fixed as SB2 (LEFT) or PC (RIGHT).
-1.450
11.475
-15500
Log-likel/hood
-1.525
-1.550
-1.575
-1.600
PC
-1.625
SB
SB2
11.550
T= 10
T= 10
T= 10"
t (prior impact)
Figure Impact of the prior impact hyperparameter and the multi-head pooling strategy for synthetic clustering.
20

---

Published as a conference paper at ICLR 2022
-1.5
Log-likelihood
-1.6
-1.7
-1.8
k
k=2 (the number of EM steps)
k=
Figure 10: Impact of is (the number of EM steps) for synthetic clustering. The model is with
(p = 4, H 5, = 0.01, SB2).
Adjusted index
0.975
0.970
0.965
0.960
10-6
10-5
10-4
10-3
10-2
10-1
10
(prior impact)
Figure 11: Impact of the prior strength hyperparameter 7 on the CIFAR-100 clustering dataset with the DIEM model = 4, H 5. le 3, SB2).
CIFAR- 1-00 clustering, We vary the prior impact hyperparameter T from 10-6 to 10 and the results are depicted in Fig. 11 It indicates that choosing a moderate prior trade-off (10-2) yields the best accuracy.
SCOP 1.75 protein fold classification. We test the protein fold classification dataset the impact of the prior strength 7. For the unsupervised learning setting, we vary and the results are shown in Fig. 12 We have similar trend CIFAR- 100 clustering, where the moderate strength performs the best. Finally, we compare the types of E-steps performed for the supervised learning setup. For the DIEM model (p 50, H k = 2, r 10 run in two different ways: the regular E-steps and the OT E-steps (i.e., without and with imposing the balanced assignment constraints, respectively). The results shown in Fig. 13 and the regular E-steps are more effective than OT E-steps for this particular dataset.
10-3
128), the k EM steps
21

---

Published as a conference paper at ICLR 2022
OTKE
86.4
Accuracy (top 1)
86.2
85.8
85.6
10
10
10
10
10
10
100
<prior impact)
Figure 12: Impact of the prior strength hyperparameter on SCOP 1.75 unsupervised learning with the DIEM model (p 100, H 11, 1(1)..
91
Regular E-steps OT E-steps
90
Accuracy (top 1)
85
Figure 13: Comparison of regular E-steps and OT E-steps SCOP 1.75 supervised learning.
22


Reviews:
[{"review_id": "lJEa62FlRCH", "reviewer": "Reviewer_7HuD", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper is well written and easy to understand. However, I do have some comments:\n1) It is obvious to see that DIEM achieves better results than OTKE baseline in terms of offline evaluation metrics, such as accuracy, log-likelihood score. And as the author mentioned, the improvements come from the multiple steps EM algorithms. If this is the case, has the runtime been increased? In addition, the author also mentioned that OTKE-type methods would reduce the computational cost compared with attention (Set)Transformer. Based on two arguments, the running time probably should be compared between different baselines.\n\n2)  DIEM doesn't have better results than OTKE on the largest DeepSEA dataset, which would influence the practical performance of DIEM on the large-scale NLP/Bioinformatics tasks.\n\n", "summary_of_the_paper": "The paper is well written. In this paper, the author proposed an EM-based algorithm, DIEM, for set representation learning. The author first provides the equivalence between the OTKE representation learning algorithm and a single-step EM algorithm with extra balanced assignment constraints on the E-step. Then DIEM is developed and consistently outperforms/competes with OTKE algorithms in different empirical studies with the assistance of multiple EM steps and extra regularization. And DIEM is applicable both for supervised and unsupervised settings.", "summary_of_the_review": "Please refer above.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "g7zbL5oTfH8", "reviewer": "Reviewer_Y94U", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Pros:\n- This work introduces a principled method for representing sets.\n- The OTKE method is derived in a principled manner. An interesting consequence is that the choice of the number of reference can be made using the existing litterature of mixture fitting.\n- Good experimental results on varied datasets (NLP, bioinformatics, vision, synthetic).\n- Sensitivity studies for different hyperparameters.\n\nCons:\n- The proposed method may somehow lack of novelty since the idea of using prototypes has been very studied recently.\n\nQuestions and remarks:\n- What is the intuition of doing multiple EM steps in terms of embedding? Can this be related to the recent Perceiver [1] architecture? What is your view on this?\n- Does DIEM learn the parameters of the prior distribution in the supervised setting? This could be more clear in the paper.\n- The paper claims that the method has low computational complexity but it seems that this claim is not detailed in the paper (apart from remarks on the number of prototypes). Could you elaborate on the complexity of the EM steps?\n- It could be great to provide more details on how to set the hyper-parameters for your method.\n- Could you further discuss the impact of the prior depending on the task? Could we inject another prior/inductive bias here?\n- Features given by protein language models such as ESM [2] can greatly improve results for SCOP 1.75. In fact, this may be the actual state-of-the-art for this dataset (see Table 5 in OTKE paper). Transfer learning is however orthogonal to the method proposed here but it is worth having this in mind.\n- In the related work: \"The limitation was found...\": could you elaborate on this?\n\n-----------------------\n\n[1] Perceiver: General Perception with Iterative Attention (Andrew Jaegle and Felix Gimeno and Andrew Brock and Andrew Zisserman and Oriol Vinyals and Joao Carreira)\n\n[2] Biological structure and function emerge from scaling unsupervised learning to 250 million protein sequences (Rives, Alexander and Meier, Joshua and Sercu, Tom and Goyal, Siddharth and Lin, Zeming and Guo, Demi and Ott, Myle and Zitnick, C. Lawrence and Ma, Jerry and Fergus, Rob)", "summary_of_the_paper": "This work proposes a new embedding for sets of features, an important problem since many data modalities can be seen as such (images, sentences, etc.). More precisely, a set is represented by the output means of an EM algorithm for fitting the input set with a mixture of gaussians. The authors draw a new connection to an existing method for set embedding (OTKE). Moreover, their method achieves good experimental results.", "summary_of_the_review": "The paper seems sound and provides new insights for set representation, with convincing experiments. I tend to recommend acceptance but it would be great if the authors could answer my questions.\n\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "Ol9DnAUw86", "reviewer": "Reviewer_xBWk", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n1. The connections between OTKE and EM is insightful in set representation learning, and differentiable EM is well motivated.\n2. Experimental results are impressive and support the claims made in this paper well.\n\nWeakness:\n1. Time complexity or empirical wall-clock time is needed to give a thorough analysis of differentiable EM. It will be helpful to present the time complexity (or empirical wall-clock time) of differentiable EM, since it takes several EM steps and costs more time compared to OTKE.\n", "summary_of_the_paper": "This paper discusses that optimal transport kernel embedding (OTKE) can be regarded as a single expectation-maximization (EM) step towards the maximum likelihood estimate of Gaussian mixture models under mild conditions. Motivated by the finding, this paper proposes differentiable EM, which can be regarded as a generalized version of OTKE with prior and several EM steps. Experiments on OMNIGLOT unique character counting, amortized clustering in CIFAR-100, protein fold classification on SCOP 1.75, sentiment classification on SST-2 and chromatin profile detection on deepsea demonstrate the effectiveness of differentiable EM on set representation learning. ", "summary_of_the_review": "This paper presents a novel idea about set representation learning. Experiments cover multiple tasks and support the claims well. Though more analysis on time complexity is needed, I think this paper is above the acceptance threshold.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "MKSOsZ_fv8J", "reviewer": "Reviewer_emDK", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Overall, I like the paper; it is well written, and the interpretation of the set-embedding procedure as an EM iteration indeed makes sense. It is also good to see the authors derive a novel algorithm from their re-interpretation. The experiments are diverse and thorough, and as far as I can see, they seem to be reproducible with all the details provided in the appendix.\n\nI think the paper can be enhanced with some further clarification.\n\n1) In my opinion, it is quite important to compare the number of parameters when comparing different set embedding methods; for instance, in (Lee et al., 2019), they set the number of parameters for DeepSets and Set transformers roughly the same. How many parameters were used for the proposed method? I hope to see the parameter counts at least in the appendix. It would also be helpful to compare the wall-clock time for the forward passes; especially, for the proposed method, it is worth checking the inference time w.r.t. the number of EM iterations $k$.\n\n2) There are quite a few hyperparameters or options for the proposed model; the number of mixture components $p$, number of EM iterations $k$, prior hyperparameter $\\tau$, and the way of pooling (PC, SB, or SB2). Judging from the appendix, the performance of the proposed approach is quite sensitive to the choice of these hyperparameters. I'm also quite confused with three options for the pooling; is there any guide for which one to choose? Was any of those three pooling methods dominant in general? It is quite hard to directly compare the effect of individual choices of the hyperparameters because the results so far is not controlled experiments for the hyperparameters. Does the performance generally saturate with the number of mixture components $p$ or the number of EM steps $k$?\n\n3) Have you considered using generative models other than a mixture of Gaussians? I guess the primary reason for the choice is its conjugacy, but probably we can think of other conjugate pairs for the mixture components.\n\n4) Collapsing the hyperparameters $\\tau = \\eta-1 = \\lambda = 1 = \\nu + d + 2$ is weird; for instance, $\\nu + d + 2$ cannot be equal to one. Can you elaborate on this?\n\n5) How important is the step to initialize the parameters as the mode of the posteriors? What happens with the randomly initialized parameters or learning them as well with gradient descent? For instance, if the mixture components are not conjugate so the MAP parameters are not easily estimated, then we may consider different options.\n\n", "summary_of_the_paper": "This paper proposes a novel set embedding method inspired by the EM algorithm. Treating each element in a set as i.i.d. samples from a mixture of Gaussians, the procedure of computing pairwise similarities between the elements and prefixed set of reference vectors corresponds to the computation of responsibilities in E-step for the mixture of Gaussians, and the embedding step using the similarities corresponds to the parameter update in M-step. The previous approaches such as OTKE can directly be interpreted with this EM view (plus balanced assignment constraint). Based on this reinterpretation, the paper proposes a novel set-embedding method extending previous methods in various ways; 1) use multiple steps of EM updates, 2) learn parameters other than reference vectors (covariances and mixing proportions), 3) learn the initial value of the parameters by placing prior distributions on them. The resulting algorithm entitled DIfferentiable EM (DIEM) is demonstrated to excel in various set-to-vec tasks.", "summary_of_the_review": "The paper proposes an interesting idea, and the experimental results are promising. There are some minor concerns to be clarified.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This work proposes a new embedding for sets of features. A set is represented by the output means of an EM algorithm for fitting the input set with a mixture of Gaussians. The authors draw a new connection to an existing method for set embedding (OTKE). Moreover, their method achieves good experimental results.

There is general consensus among the reviewers that the paper is sound, well-written and provides new insights for set representation, with convincing experiments.

The authors have answered to most comments raised by the reviewers and have revised the paper accordingly.

I recommend acceptance as a poster.
Decision: Accept (Poster)