Venue: ICLR 2021 Conference

Reviews:

final_score: rating: 6: Marginally above acceptance threshold
review_content: Summary:

This paper proposes a network architecture and a training model for learning Wasserstein barycenters of 512x512 discrete probability distributions. Results are compared with other approximation methods.

Reason for my score:

This paper is ell written, and the idea is interesting.  However, the elephant in the room is: How close are these learned barycenters from the actual barycenter? For me this is the main question, that cannot be answered with CNN approaches, and thus, makes it difficult to translate into uses for optimal transport techniques. This limits the impact that the paper can have on a broader audience.

The authors made a great effort to provide a rather complete review of the literature, and the provided experiments are sound. However.

1. The comparison of the results is made against another approximating method. Again, I would like to see some evidence that the learned model is actually generating good approximations of the Wasserstein Barycenter.

2. It is a less exciting if the loss function used is not OT itself but KL.

3. The main advantage of entropy regularized methods is that one can actually control how close the computed WB is from the actual WB. This is a main issue. In the proposed method, we have no way to know this, and I cannot avoid thinking this is just an informed guess.

4. I believe one main advantage of the proposed method would be to generate initialization points of entropy regularized approaches.

5. The proposed architecture is well explained. Why not other architecture, some experiments exploring other ways would be appreciated as well.

6. For some basic distributions it would have been more useful to compare the learned barycenter with the output of the linear program. At least in a couple of cases.  This can provide some evidence that the learned barycenter is actually a good approximation.


review_title: An interesting experiment with Wasserstein Barycenters

final_score: rating: 7: Good paper, accept
review_content: A modification of U-Net to compute Wasserstein barycenters

This paper presents a deep network architecture which learns to generate Wasserstein barycenters associated with a given set of input measures. The paper is clearly written and presents a fairly objective view on their work. It gives one more example of optimization task (here the convex minimization problem of Wasserstein barycenters) that can be replaced by deep neural networks with seemingly satisfactory results.
Overall, although the method is relatively straightforward, it has the merit of showing that fast approximations of linear approximations of Wasserstein barycenters can be performed at a cheap forward pass cost of a U-Net. Note that the method has no built-in information on the structure of the problem e.g. preservation of total mass for instance or more higher-level properties.  These fast approximations could be proven to be useful in other settings outside machine learning too.

The benefit of the proposed architecture is that it can be adapted to compute the barycenters for any given number N of inputs while only trained on pairs (N =2) of input measures. Apart from minor modifications in the pipeline, the architecture remains close to that of a U-Net with the modification that the barycenters weights are explicitly used to compute averages of the activations at different layers of the Deep CNN. This work is quite weak in terms of novelty (structure of the network for instance).

The experiments are compared against and trained with the GeomLoss code which seems to be state of the art for computing entropic regularity. They also compare with Wasserstein embeddings methods favorably. 

— References to literature in using deep learning for solving inverse problems would be welcome. Wasserstein barycenter is a particular instance of such class.

— Comments on the sensitivity of the results to training data would be welcome.

— I find the description of the algorithm for building the training data quite unclear. If I understand correctly, the training data is an approximation of Wasserstein barycenters between two measures. These approximations are simply computed with one step of the gradient descent of the Sinkhorn divergence and use an average of the corresponding transport maps. The cost of computation of these approximate barycenters are essentially two sinkhorn loops up to the given tolerance. This would have been interesting to have a comparison, in the supplementary material, of the difference between these approximations and more accurate approximations of Wasserstein barycenters such as given by entropic regularization with very small regularization parameter and discuss in more details what can be the impact of such inaccuracies. The authors say « While using more iterations of gradient descent yields more accurate results and removes this linearity, it also prevents easy combination and makes the dataset generation intractable ». I do not understand this statement. Cost of the gradient descent linearly scales with the number of steps, so what do the authors mean by « easy combination », the linear approximation step ?

Minor remarks: 
— the (not numbered) figure on page 6 called adjacent is almost impossible to read and it is not explained enough.
—  A dot is missing after KL-divergence on page 6.
— Why is DWE performing so badly in figure 3 in comparison with figure 4 ?
review_title: A modification of U-Net that computes approximate Wasserstein barycenters

final_score: rating: 3: Clear rejection
review_content: This paper proposes a numerical method to approximate 2d Wasserstein barycenters, that relies on a convolutional network. Namely, the proposed architecture is composed of
1.	Contractive paths, made of series of 2 convolutional layers, relu activation and a pooling step, characterized by weight sharing.
2.	An expansive path, symmetrically constructed, with upsampling by NN interpolation, and a final softmax activation.

Though fixed during training, the varying number of contracting paths in the test setting is able to predict the barycenter of a varying number of measures. The KL loss is used to train the network predicting the Wasserstein barycenter.
The method is assessed on the Quick-draw dataset, to produce barycenters of two measures and its generalization to barycenters of n measures.

Strong points of the paper include:
1.	The writing of the paper is clear.
2.	The goal to propose a speedup for the computation of Wasserstein barycenters with a varying number of input measures is pertinent.

Weak points include:
1.	The paper proposes an architecture tailored for very specific setting of 2d images. Its extensions to higher dimensional or non-image settings are limited and are not discussed.
2.	Though mentioned as closely related works, the paper is missing important baselines, for instance (Claici et al.,2018).
3.	The fact that only one dataset is considered undermines the paper, which focuses on numerics.
4.	The figure at the end of page 6 (which is lacking both title and legend) does not provide a relevant metric to assess the benefits of the proposed method. 
5.	The paper does not provide other metrics (such as inception scores) to assess the performance of the methods other than visually.

I recommend a reject on the aforementioned grounds.
review_title: Limited contributions

Meta-review: Given the reviewer's exchange with the authors, and my own examination of the paper, I don't think that it can be accepted in the present form.

First, since this paper aims at solving an optimization problem (for which existing methods exist, with theoretical guarantees) via a NN, it is important to compare appropriately to those methods, which is not done here.

Further, there are possible issues when applying these only to 2D data, and it is possible that it would not extend appropriately to other types of geometries, and costs in OT problems.
Decision: Reject