Venue: ICLR 2021 Conference

Reviews:

final_score: rating: 6: Marginally above acceptance threshold
review_content: This paper proposes a kernel embedding for a set/sequence of features, based on the optimal transport distance to a set of references, leading to a fixed-dimensional embedding of variable length sequences.
The set of references can be obtained as cluster centers over the full dataset (unsupervised), or learned based on a downstream objective (supervised).
The method is somewhat related to a single layer of self-attention with an optimal transport map instead of dot product attention.
Experiments on classifying variable-length sequences are presented between protein, chromatin and NLP sentiment classification.

I recommend weak accept because (+) a well-motivated novel embedding/pooling architecture with (-) somewhat weak, initial proof-of-concept style results and (-) somewhat strange positioning against self-supervised methods and self-attention.

Strenghts:
* Well-motivated novel layer for embedding a variable-size set or sequence to a *fixed dimensional* embedding space (a point that's not directly apparent though, consider emphasizing this point)
* Relatively elegant formulation grounded in OT and kernel methods.
* The same framework allows an unsupervised and supervised variant.

Weaknesses:
* Experimental results are not very strong. For protein fold classification, I expect a comparison to transformer-based methods on proteins (eg ESM-1, https://github.com/facebookresearch/esm), simply average-pooled over the full sequence. For sentiment analysis, it is cool to see the unsupervised/linear probe on fixed features, beating the frozen BERT - however not competitive to BERT+finetuning which is how it is usually done.
* Given weak results, the impact of this method is not very clear. It is only shown on shallow single-layer OTKE, without clear discussion how this could be extended to deeper architectures.
* Given the shallow architecture, the paper is somewhat mispositioned, given the (a) introduction calling out limitations of self-supervised transformers and (b) comparison to self-attention in Sec 4: this suggests inserting the OTKE as a drop-in replacement for the dot product in a transformer model, trained with a self-supervised objective like MLM. Given the framing it seems such a logical step, it surprises me to see at least a discussion of it missing. This would be akin to recent work [1,2] of approximating self-attention by mapping to a fixed-dimensional space.
Alternatively, it seems like the method could be applied on top of strong/SotA MLM-transformer embeddings to enable better pooling than just [CLS] token or mean pooling (both protein fold prediction and sentiment classification).

Minor comments / questions:
* The labeling "unsupervised" in Table 2 and 4 is somewhat confusing/misleading, rather I'd name it "linear probe" or "supervised finetuning". Maybe re-clarify in intro of 5.1 "supervised == learning the z also"?
* What are the input features for protein fold classification? it is just mentioned they are 45-dimensional, not what is their nature (amino acid identity + pssm?)
* On end-to-end learning by back-propping through unrolled sinkhorn (Sec 3.4): this seems to be the method introduced in [3]? I think this paper should be cited in this context.
* For clarity: In both abstract, end of introduction, start of Sec 3.1 and Sec3.3, it is simply not clear to the reader what is the input/output of your method: a phrase like this would be much needed: "an embedding+pooling layer which aggregates a variable-size set or sequence of features to a fixed-size embedding".


- [1] arXiv:2006.03555 Choromanski et al. Masked Language Modeling for Proteins via Linearly Scalable Long-Context Transformers
- [2] arXiv:2006.04768 Wang et al. Linformer: Self-Attention with Linear Complexity
- [3] Genevay, A., Peyr´e, G., Cuturi, M., et al. (2018). Learning Generative Models with Sinkhorn Divergences. In AISTATS.



----
### EDIT: reply to authors' response
I thank the authors for their in-depth response and revision. The additional results comparing to pre-trained features definitely adds perspective, and it shows OTKE giving a very modest boost indeed, mostly over the weaker input features. The revision of the paper also improves the clarity, but my comment around Sec4 remains.
I stand by my original score for the 3 reasons I had originally listed.
review_title: A well-motivated novel embedding/pooling architecture with somewhat weak results

final_score: rating: 7: Good paper, accept
review_content: **Summary**:
The authors propose a new way to aggregate the embeddings of elements in a set (or sequence) by comparing it with respect to (trainable) reference set(s) via Optimal Transport (OT). The motivation to build such a pooling operation is derived from self-attention and the authors suggest an OT spin to that (e.g., the different reference sets/measures can be thought of as different heads in attention). This is, however, done in a principled way with the help of kernel embeddings and not just ad-hoc using the transport plan as the attention matrix.
 
**Pros**:
-They properly bridge the gap of OT with attention via their non-local pooling perspective. This allows them to have a similar feature aggregator, but which requires linear memory as compared to quadratic for the vanilla self-attention (although there have been several recent linear memory variants to it, any comments?). 
-The method results in an improvement over other baselines for several biology-based sequencing applications (also they demonstrate early results for NLP). 
-The paper is well written and clear.

**Cons**: 
-The empirical results are okay, but could have been better or tested more extensively. Also, there are some peculiarities about their empirical analysis. 
-This kind of technique has promise for use in NLP, but I am slightly worried that right now there are just too many "knobs" whose correct configuration will have to be worked out (anchor points, reference measures, nystrom, etc). However, I think this can be left to future work. 

*-> It would be great if the authors can answer the questions below:* 

Empirical analysis:
-In the experiments for Table 2, 3, 4, do you use the same number of heads for other baselines as the number of references in your case?

-It seems quite funny that, for their unsupervised results, the optimal number of supports p equal 1, which would basically imply that optimal transport wasn’t much of use and the features of x_i were just pooled based on an average. Do you have any comments on this and why this is happening?

-The NLP experiments are relatively weak and could still have been strengthened by testing on more tasks from the GLUE benchmark. Another surprising thing is that even though scores of your method unsupervised are higher than mean pooling of BERT features, the latter significantly outperforms in the supervised case. Do you have any explanation for this?

-Besides, why do you write in the text that NLP is “an a priori less favourable setting”?

-Can you give some precise estimates of how the runtime of this OTKE compares to dot product attention? Also, how expensive is the Nystrom procedure?

Miscellaneous: 
-This is a bit of nitpicking, but the authors are a bit lazy while citing, and in several places just cite the textbooks for OT and Kernels (which is not a problem). However, it is useful for the reader to additionally have the accurate references as well, e.g. for Sinkhorn’s algorithm perhaps also cite Sinkhorn & Knopp (1967) and Cuturi (2013); for OT theory, Villani 2008 etc.  
-Typos: "a set x and a reference z in X and a reference z”
-ISAB and PMA (unexplained acronyms)

**Conclusion**:  Overall, I think the idea is quite interesting, and carries the potential for OT-based building blocks in NLP, or even analysing the benefit provided by attention. Hence, I am in favour of accepting this paper.  

(PS: I might be slightly biased because I like OT and have also been thinking about ideas on similar lines connecting attention and OT.)
review_title: Nice and well-executed idea for feature aggregation based on OT (with connections to attention)

final_score: rating: 6: Marginally above acceptance threshold
review_content: The paper proposes a transport-based feature representation for the input of a set of vectors. The feature is defined through the transport to reference sets. Vectors in reference sets can be learned by supervised or unsupervised approaches. The relation with self-attention is also discussed. Experiments empirically show a good performance of the proposed method.

Transporting a set of inputs to trainable references is seemingly a more sophisticated extension of the classical basis function regression with trainable basis (basis corresponds to references), and thus, I think the basic idea would be reasonable and easy to have an intuition. My major concern is on about the settings of the number of references and supports, which would have a large number of possible settings. Except for that, I currently only have minor comments.

Discussing the selection of the number of references and supports in more detail would have been informative. In my understanding, these are most important hyper-parameters. Currently, no practical recipes and analysis of sensitivity to these settings are provided. How were they chosen in the experiments?

The number of supports can be different for each one of references. Therefore, there exist a huge number of possible settings even when only considering moderate sizes.

The optimized references seemingly have some interpretation about discovery of an important set of vectors. I'm not fully sure, but it seems to provide a different view point from self-attention, and thus, it is perhaps worth mentioning if possible. 

A naive extension could be to make multiple transportation layers (transport z to another references). Is it meaningless? (I guess, at least, some activation function is required) If it has some meaning, providing comments may be informative.

In eq.(5), 'R^k x p' should be 'R^p x k'?
review_title: Simple, but reasonable approach 

final_score: rating: 7: Good paper, accept
review_content: ##########################################################################

Summary:

The paper proposes a parametrized embedding based on optimal transport in the kernel space. The method is scalable and the parameters of the embedding can be learned in either unsupervised or supervised fashion. The paper demonstrates the effectiveness of the purposed approach using real data in bioinformatics and natural language processing. 

Overall, the proposed embedding seems to be novel and well-motivated. The empirical results are impressive. Also, the paper is well written and technically sound. I don't have major concerns about the paper. 


##########################################################################

Minor comments: 

- Contributions paragraph: "in either unsupervised and supervised fashion" --> "in either unsupervised or supervised fashion"
- The line below Eq. (1): the entropy should be - \sum P ( log P - 1 ) ? 
- Section 3.2 "Infinite-dimensional embedding in RKHS" line 1: "a reference z in X" repeated twice.
- Third line below Eq. (4): \Phi_z(x) instead of \Phi_z(z)
- First line in Section 5.2: "either unsupervised or supervised settings" --> "both unsupervised and supervised settings"
review_title: Interesting idea

Meta-review: All reviewers agreed that the paper proposes some interesting and novel ideas on the use of OT for pooling. It also provides some nice insights and strong experimental results. As suggested by one of the reviewer, a discussion about the impact of the number of references may be of interest though.


Decision: Accept (Poster)