Optimal transport maps for distribution preserving operations on latent spaces of Generative Models

Abstract: Generative models such as Variational Auto Encoders (VAEs) and Generative Adversarial Networks (GANs) are typically trained for a fixed prior distribution in the latent space, such as uniform or Gaussian.
      After a trained model is obtained, one can sample the Generator in various forms for exploration and understanding, such as interpolating between two samples, sampling in the vicinity of a sample or exploring differences between a pair of samples applied to a third sample.
      In this paper, we show that the latent space operations used in the literature so far induce a distribution mismatch between the resulting outputs and the prior distribution the model was trained on. To address this, we propose to use distribution matching transport maps to ensure that such  latent space operations preserve the prior distribution, while minimally modifying the original operation. 
      Our experimental results validate that the proposed operations give higher quality samples compared to the original operations.

Reviews and author responses:

Authors note that models may be trained for a certain distribution (e.g. uniform or Gaussian) but then "used" by interpolating or jittering known examples, which has a different distribution.  While the authors are clear about the fact that this is a mismatch, I did not find it well-motivated why it was  "the right thing to do" to match the training prior, given that the training prior is potentially not at all representative or relevant. The fact that a Gaussian/prior distribution is used in the first place seems like a matter of convenience rather than it being the "right" distribution for the problem goals, and that makes it less clear that it's important to match this "convenience" distribution. The key issue I had throughout is "what is the real-world problem metric or evaluation criteria and how does this proposal directly help"?

For example, authors cover the usual story that random Gaussian examples  lie on a thin sphere shell in high-d space, and thus interpolation of those examples will like on a thin shell of slightly less radius.  In contrast, the Uniform distribution on a hypercube [-1,1]^D in D dimensions "looks" like a sharp-pointy star with 2^D sharp points and all the mass in those 2^D corners. But the key question is, what are these examples being used for, and what are the trade-offs between interpolation (which tends to be fairly safe) and extrapolation of the given examples?

This is echoed in the experiments, which I found unsatsifactory for the same key issue: "What is the criteria for “higher-quality interpolated samples”? in the examples they give, it seems to be the sharpness of the images. Is that realistic/relevant? These are pretty images, but the evaluation criteria is unclear.




This paper is concerned with the mismatch between the input distribution used for training and interpolated input. It extends the discussion on this phenomenon and the correction method proposed by White (2016), and proposes an optimal transport-based approach, which essentially makes use of the trick of change of variables. The discussion of the phenomenon is interesting, and the proposed method seems well motivated and useful. There are a number of errors or inconsistencies in the paper, and the experiments results, compared to those given by SLERP, see rather weak. My big concern about the paper is that it seems to be written in a rush and needs a lot of improvement before being published. Below please see more detailed comments.

- In Introduction, the authors claim that "This is problematic, since the generator G was trained on a fixed prior and expects to see inputs with statistics consistent with that distribution." Here the learned generative network might still apply even if the input distribution changes (e.g., see the covariate shift setting); should one claim that the support of the test input distribution may not be contained in the support of the input distribution for training? Is there any previous result supporting this? 
- Moreover, I am wondering whether Sections 2.2 and 2.3 can be simplified or improved--the underlying idea seems intuitive, but some of the statements seem somewhat confusing. For instance, what does equation (6) mean?
- Note that a parenthesis is missing in line 3 below (4). In (6), the dot should follow the equation.
- Line 1 of page 7: here it would be nice to make it clear what p_{y|x} means. How did you obtain values of f(x) from this conditional distribution?
- Theorem 2: here does one assume that F_Y is invertible? (Maybe this is not necessary according to the definition of F_Y^{[-1]}...)
- Line 4 above Section 4.2: the sentence is not complete.
- Section 4.2: It seems that Figure 3 appears in the main text earlier than Figure 2. Please pay attention to the organization.
- Line 3, page 10: "slightly different, however..."
- Line 3 below Figure 2: I failed to see "a slight loss in detain for the SLERP version." Perhaps the authors could elaborate on it?
- The paragraph above Figure 3 is not complete.

The authors demonstrate experimentally a problem with the way common latent space operations such as linear interpolation are performed for GANs and VAEs. They propose a solution based on matching distributions using optimal transport. Quite heavy machinery to solve a fairly simple problem, but their approach is practical and effective experimentally (though the gain over the simple SLERP heuristic is often marginal). The problem they describe (and so the solution) deserves to be more widely known.

Major comments:

The paper is quite verbose, probably unnecessarily so. Firstly, the authors devote over 2 pages to examples that distribution mismatches can arise in synthetic cases (section 2). This point is well made by a single example (e.g. section 2.2) and the interesting part is that this is also an issue in practice (experimental section). Secondly, the authors spend a lot of space on the precise derivation of the optimal transport map for the uniform distribution. The fact that the optimal transport computation decomposes across dimensions for pointwise operations is very relevant, and the matching of CDFs, but I think a lot of the mathematical detail could be relegated to an appendix, especially the detailed derivation of the particular CDFs.

Minor comments:

It seems worth highlighting that in practice, for the common case of a Gaussian, the proposed method for linear interpolation is just a very simple procedure that might be called "projected linear interpolation", where the generated vector is multiplied by a constant. All the optimal transport theory is nice, but it's helpful to know that this is simple to apply in practice.

Might I suggest a very simple approach to fixing the distribution mismatch issue? Train with a spherical uniform prior. When interpolating, project the linear interpolation back to the sphere. This matches distribution, and has the attractive property that the entire geodesic between two points lies in a region with typical probability density. This would also work for vicinity sampling.

In section 1, overfitting concerns seem like a strange way to motivate the desire for smoothness. Overfitting is relatively easy to compensate for, and investigating the latent space is interesting regardless.

When discussing sampling from VAEs as opposed to GANs, it would be good to mention that one has to sample from p(x | z) not just p(z).

Lots of math typos such as t - 1 should be 1 - t in (2), "V times a times r" instead of "Var" in (3) and "s times i times n" instead of "sin", etc, sqrt(1) * 2 instead of sqrt(12), inconsistent bolding of vectors. Also strange use of blackboard bold Z to mean a vector of random variables instead of the integers.

Could cite an existing source for the fact that most mass for a Gaussian is concentrated on a thin shell (section 2.2), e.g. David MacKay Information Theory, Inference and Learning Algorithms.

At the end of section 2.4, a plot of the final 1D-to-1D optimal transport function (for a few different values of t) for the uniform case would be incredibly helpful.

Section 3 should be a subsection of section 2.

For both SLERP and the proposed method, there's quite a sudden change around the midpoint of the interpolation in Figure 2. It would be interesting to plot more points around the midpoint to see the transition in more detail. (A small inkling that samples from the proposed approach might change fastest qualitatively near the midpoint of the interpolation perhaps maybe be seen in Figure 1, since the angle is changing fastest there??)



Dear Reviewers, thanks again for your feedback and happy new year! 

Since two reviewers felt the experiments could be stronger, we have added a new revision with additional quantitative experiments (Section 3.3. and Table 2.) which compare the interpolation operations using Inception scores. These results mirror what was qualitatively observed in Section 3.2 -- namely that when compared with the original models, the linear interpolation gives a significant quality degradation (up to 29% lower Inception scores), while our matched operations do not degrade the quality (less than 1% observed difference in scores).

Regarding individual comments we refer to the individual responses previously posted. All other Figures, Tables and Sections referenced there have the same numbers as in the previous revised edition, so you only need to look at the latest revision.


Thank you for your review. 
We discuss your raised concerns and hope you reconsider the rating.

"It is not well-motivated why it is  'the right thing to do' to match the training prior, given that the training prior is potentially not at all representative or relevant..." "... [using the prior] seems like a matter of convenience ..."

While true that the specific prior chosen is a matter of convenience, after it has been chosen it is *the prior that the model is trained for*. This is the standard practice when training GANs, so our point is that after you train your model you need to respect the prior you chose. So then you might say that the "wrong" prior was chosen, but it is well known that any distribution (in principle) can be sampled from via a mapping G applied to samples of a fixed (e.g. uniform) distribution z. See Multivariate Inverse Transform Sampling (e.g. slide 24 in https://www.slac.stanford.edu/slac/sass/talks/MonteCarloSASS.pdf ).


"...what is the real-world problem metric or evaluation criteria and how does this proposal directly help?"

The goal of this work is to improve upon how generative models such as GANs are visualized and explored when working with operations on samples. Too see why this is relevant in Section 1.1 (revised edition) we mention eight papers (out of many more) in the recent literature which use such operations to explore their models. A 'real-world' use case hinges on real-world use cases of generative models, but just to give an example you could imagine an application that allows a user to 'navigate' the latent space of a generated model to synthesize a new example (say logo/face/animated character) for use in some real world application. Such exploration of the model needs to allow for various operations to adjust the synthesized samples.


Regarding the 'usual thin sphere story' we note that the radius difference is quite significant, see Figure 2 (revised edition) which shows the radius distribution for the latent spaces typically used in the literature. Our approach completely sidesteps the issue.


For the experiments, we have added more examples of latent space operations and a discussion on the differences. A key property of our proposed approach is that it is 'safe': if you repeatedly look at some output of any operation (say e.g. midpoint of the matched interpolation), it will have exactly the same distribution as random samples from the model. Hence no matter what kind of image quality assessment you would use, it would be the (statistically) the same as for samples from the model without any operations.


Thanks for the feedback! 

While sticking to the main story we have significantly polished the paper. 
We have improved discussion of the experiments, acknowledging that there is not a  really noticeable difference between the SLERP heuristic and our matched interpolation in practice. This is perhaps not so surprising since SLERP does tend to match the biggest distribution difference (the norm mismatch) quite OK in practice (see Fig. 2 revised paper). 
Nonetheless, our proposed framework has many benefits which we have also better highlighted in the paper:
	- it gives a new  and well grounded perspective on how to do operations in the latent space of distributions
	- it is straightforward to implement, especially for a Gaussian prior (see Tab. 1 revised paper).
        - it generalizes to almost any operation you can think of, not just interpolation (see e.g. random walk in Fig. 11 (revised paper)).

Regarding specific comments:

- while the trained model -might- apply also for a different distribution, for the linear interpolation we typically see a clear difference. Note we do not claim that the supports of the distributions do not overlap - we only claim this for the distribution of the norms.

- we significantly simplified the explanation and motivation of Sec 2.1-2.2 (old version), removing the synthetic example (including eq (6)) and better focus on the (more relevant in practice) norm distribution difference - with detailed calculations moved to appendix. The subsections are merged into the intro of Sec 2 in the revised edition.
These changes were also in line with suggestions from AnonReviewer3 on simplifying the paper.

- p_{y|x} has been clarified in the text, it was referring to f(x) being a random variable where f(x)  is drawn from the conditional distribution over y given a fixed x. If this is unclear/confusing in our notation, we can also instead just cite the fact that KP is a relaxation of MP.

- Theorem 2: while the derivations would be easier if F_Y were invertible, it is not needed. F_Y is always monotonic, and F_Y^{[-1]} denotes the pseudo-inverse (hence the bracket [-1] ). See  https://en.wikipedia.org/wiki/Cumulative_distribution_function#Inverse_distribution_function_(quantile_function) and (Santambrogio, 2015) for more details.

- other typos/mistakes: should be fixed in revised version



Thanks for the feedback!
We followed your suggestion in the major comment and significantly polished and shortened the paper. 

 - as suggested, we focus on explaining the effect distribution mismatch through the norm distribution, moving unnecessary details to the appendix.

- we moved Lemma 1  to appendix as well as the detailed calculations of the examples, while summarizing the Gaussian case in Table 1.

- We now mention how simple the formulas end up in the Gaussian case. This is because the operators we consider are additive in the samples, which means the results of the operations are still Gaussian - requiring only a multiplicative adjustment for matching the variance.

- Working on the hypersphere is also a valid approach. This setting is very similar to our framework applied to  the Gaussian prior when taking the prior dimension towards infinity - and the projection to the sphere can be interpreted as the transport map. Note however by fixing points to lie exactly on the sphere one introduces a dependency between the coordinates (which means you can't do distribution matching coordinate-wise), but this dependency is very small since an i.i.d. Gaussian will already be on the sphere w.h.p. We actually tried this setting at some point before, but found it (surprisingly) less stable for DCGAN, e.g. resulting in collapse for the icon dataset.
- We adjust the motivation, as you mention interpolations and other operations are interesting on their own, and overfitting can be measured through other means.

- on VAEs vs GANs, we are currently only discussing the sampling in the test setting - where one only samples from p(z) ( see Figure 5 in https://arxiv.org/pdf/1606.05908.pdf )

- Typos/inconsistencies should now be fixed

- We added plots showing the  1D-to-1D monotone transport maps for Uniform and Gaussian, see Figure 3 revised edition.

- We will add a citation to  David MacKay for the mass distribution of a Gaussian. However we didn't find a nice reference which gives the same result for arbitrary distributions with i.i.d components.

- In Figure 15 in the appendix, we show example interpolations with twice as many points, so the transition is clearer. We note that the color may change sharply when interpolating between examples if the inbetween color is not 'realistic' for the data.

Meta-review: This paper exposes a simple recipe to manipulate the latent space of generative models in such a way to minimize the mismatch between the prior distribution and that of the manipulated latent space. Manipulations such as linear interpolation are commonplace in the literature, and this work will be helpful to improve assessment on that front.

Reviewers found this paper interesting, yet unpolished and incomplete. In subsequent iterations, the paper has significantly improved on those fronts, however the AC believes an extra iteration will make this work even more solid. Thus, unfortunately this paper cannot be accepted at this time. 

Decision: rejected-papers