Venue: ICLR 2022 Conference

Reviews:

final_score: rating: 6: marginally above the acceptance threshold
review_content: 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.
main_review: The paper is well written and easy to understand. However, I do have some comments:
1) 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.

2)  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.


summary_of_the_review: Please refer above.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: 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.
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.

I think the paper can be enhanced with some further clarification.

1) 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$.

2) 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$?

3) 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.

4) 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?

5) 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.


summary_of_the_review: The paper proposes an interesting idea, and the experimental results are promising. There are some minor concerns to be clarified.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: 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. 
main_review: Strengths:
1. The connections between OTKE and EM is insightful in set representation learning, and differentiable EM is well motivated.
2. Experimental results are impressive and support the claims made in this paper well.

Weakness:
1. 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.

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.
review_title: null

final_score: rating: 8: accept, good paper
review_content: 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.
main_review: Pros:
- This work introduces a principled method for representing sets.
- 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.
- Good experimental results on varied datasets (NLP, bioinformatics, vision, synthetic).
- Sensitivity studies for different hyperparameters.

Cons:
- The proposed method may somehow lack of novelty since the idea of using prototypes has been very studied recently.

Questions and remarks:
- 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?
- Does DIEM learn the parameters of the prior distribution in the supervised setting? This could be more clear in the paper.
- 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?
- It could be great to provide more details on how to set the hyper-parameters for your method.
- Could you further discuss the impact of the prior depending on the task? Could we inject another prior/inductive bias here?
- 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.
- In the related work: "The limitation was found...": could you elaborate on this?

-----------------------

[1] Perceiver: General Perception with Iterative Attention (Andrew Jaegle and Felix Gimeno and Andrew Brock and Andrew Zisserman and Oriol Vinyals and Joao Carreira)

[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_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.
review_title: null

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)