Arbitrary Conditional Distributions with Energy

Abstract: Ryan Strauss, Junier B. Oliva

Reviews and author responses:

 Thanks for your helpful comments.

 The non-negativity requirement is placed on $\mathcal{E}(x)$, the energy. The unnormalized likelihood, $e^{-\mathcal{E}(x)}$, is therefore a value between 0 and 1. You are correct that we have no loss of generality though, due to the normalizing constant. This is a typical setup when using energy-based models to define probability distributions (see LeCun et al., A Tutorial on Energy-Based Learning, Equation 2). Thank you for clarifying your question about the normalizing flows. In theory, there is no inherent restriction on the type of distribution that can be represented with an appropriate change of variables and base distribution. In practice, however, flow models are limited to the invertible and efficient (for computing Jacobian determinants) transformations that are possible with neural networks (e.g. affine coupling). These limited transformations may not always be able to express the (potentially complicated) change of variables required given a base distribution (e.g. a standard Gaussian).

The authors propose a method for performing conditional density estimation through estimating an energy function and using importance weighting to normalize the density estimate. They focus on estimating single-dimensional distribution and make use of the product rule to estimate multivariate distributions. The authors show the method is empirically promising on data imputation tasks.  ## Summary of review

The authors approach seems reasonably well-motivated, and appears to empirically perform well. I am somewhat concerned by the heavy reliance on importance sampling to estimate intermediate quantities, and then taking products of/dividing by the results. I think it would be useful if the authors provided some numerical verification (perhaps in a simple case) that the estimates used are in fact reliable. I think the authors perhaps slightly over-emphasize the approach being order-agnostic, as this is only true during training, and as applied the model may depend heavily on ordering. I didn't see an experiment showing the degree to which estimated densities depend on order in the appendix, and this might be worth illustrating when the authors discuss strategies to mitigate this dependence. However, the authors are upfront in acknowledging that order dpendence at test time is a limitation at several points in the paper, and acknowledge there is more work to be done in this regard. Overall, I thought the paper proposed an interesting and reasonably simple method, and was enjoyable to read. 

## Questions for authors
- In the related work section, the authors claim they construct an 'order-agnostic weight-sharing technique'. However, the final conditional density estimates depend on an arbitrary ordering, as later acknowledged. What is the advantage of the training process not depending on choosing an ordering if applying the model still has this constraint?
- In line 114 the authors say the *energy* is non-negative. Later the energy is taken to be $\mathcal{E}$, which is the output of a neural netrowk if I understand correctly. Is it actually intended for $\mathcal{E}$ to be non-negative, is this a typo, or have I misunderstood?
- The authors justify the importance sampling estimator by the claim that importance sampling is effective in 1D. However, it seems the authors need to take products of the estimated normalizing constants in order to estimate the overall density. Does this not lead to an accumulation of errors as the dimension of the distribution you are trying to estimate increases? 
- The authors claim that normalizing flows place many restrictions on the types of distrbutions that can be modelled. Is this true even when one-dimensional distributions are estimated with the flow? Can you give concrete examples of what cannot be modelled/is difficult to model with flows?
- What is the effect of ignoring derivatives with respect to the samples from the proposal distribution and density in 8? Is the estimate of the graident still consistent (as the number of proposal samples goes to infinity)? Is there an effect on the variance of the estimator?
- In Equation 9 how many samples must be used in practice to obtain an accurare estimate of the mean? ## Limitations
The authors point to several limitations of the model considered. For example:
- missing information is modelled as uninformative, which is not a good assumption in some cases
- Estimation of the density of a joint distribution may depend on the order in which the product rule is applied (i.e. the network used to model conditional densities may not be consistent).

## Societal implications
The authors do not point to societal implications of their work. As the proposed method might be applicable to a wide range of tasks, this seems reasonable. Any societal implications would depend signficantly on which tasks the proposed method is applied to, and it is difficult to assess what these might be.

 Thank you for your reply; it has addressed my several of my concerns and I will raise my score from a 6 to a 7 accordingly. A few other notes regarding your reply:
- Why must the energy function be non-negative? The non-negativity requirement seems to be on $e^{-\mathcal{E}(x)}$; restricting $\mathcal{E}(x)$ to be non-negative seems to ensure that the unnormalized pdf has values less than $1$. However, perhaps this results in no loss of generality because of the normalizing constant? This point doesn't seem particularly important, but if the authors can clarify, it would be helpful.
- Even if the errors were in fact independent and zero mean, we would the sum of the errors to grow (just at a slower rate). In particular, the sum of $n$ iid variables with standard deviation $1$ is $\sqrt{n}$ (it is the mean of these random variables which has a standard deviation tending to $0$). However, your argument that the errors at least may not grow quickly seems reasonable, if perhaps speculative in terms of strong assumptions about the nature of the errors.
- Regarding normalizing flows: I understand that there are limitations on the parameterization of the network that accompany the use of normalizing flows. I intended to ask about whether this implies limitations on the distributions implied by base distribution and network. I am not aware of what the existing results in this area are, but perhaps it would be worth checking the literature to see if the claim that normalizing flows limit the flexibility of the implied distribution can be substantiated more precisely.
- Regarding importance sampling: Thank you for addressing my point, as at least one other reviewer (point 2a of review ENb9) asked about the efficacy of IS in this application and if sufficient samples were used, please consider including experimental evidence of how many samples are needed in a subsequent version of the paper. 

This paper introduces Arbitrary Conditioning with Energy (ACE) to estimate arbitrary conditional densities, using a Boltzmann distribution with an energy function parametrized by a neural network to estimate 1-dimensional conditional densities. An intractable normalizing constant is estimated with importance sampling and a proposal network. The method is then demonstrated on a few examples for density estimation under missingness and for imputation.   Strengths:

The paper is clearly written, and I find the setup interesting and useful especially under missing data. Only requiring the estimation of 1-dimensional densities is also easier to understand and computationally appealing. The work is an interesting combination of the setup of ACFlow [Li et al., 2020] and the method of Autoregressive Energy Machines (AEM) [Nash \& Durkan, 2019]. The experimental results also improve on ACFlow for density estimation, whilst being simpler.

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

Weaknesses/Questions:

I summarize the main weaknesses/questions below:

1.) I am not completely familiar with the energy function literature, but are there any guarantees that the integral $Z = \int \exp\left(-\mathcal{E} (\mathbf{x}) \right) d\mathbf{x}$ is finite? I do not see how a randomly-initialized network necessarily satisfies this property. 

2.) Although the idea is simple, I find the actual implementation of ACE to be potentially over-complicated:

a.) As we are only in 1 dimension, is there any issue with using numerical quadrature instead of importance sampling? The proposal adds a lot of complexity to training, and I'm not sure if 20 importance samples (as used in the paper) is more accurate than just estimating $Z$ on a small grid.

b.) I could be missing something, but why does the energy network need the inputs $(\boldsymbol{\phi}(\mathbf{x_o}; \mathbf{b}),\mathbf{b})$ in addition to the latent vector $\boldsymbol{\gamma}(u_i; \mathbf{x_o})$? Would the latent vector not suffice, as is done for AEMs?

c.) Why does the proposal network output multiple latent vectors, that is one for each $u_i$? If it contains the information in $\mathbf{x_{o}}$, then it seems better motivated to just use the same context vector for each $u_i$. 


3.) What is the motivation for training non-autoregressively, that is with 
$ \sum_{i=1}^{|u|} \log p(x_{u_i}\mid \mathbf{x_o}) $ instead of $\sum_{i=1}^{|u|}\log p(x_{u_i} \mid \mathbf{x_{o \cup u_{<i}}})$?  It seems that $\log p(\mathbf{x}_{\mathbf{u}} \mid \mathbf{x}_\mathbf{o})$ is better motivated, since that is being evaluated in the experiments. Also, is $\mathbf{o}$ selected at random for each gradient step? It wasn't clear to me from the paper. 

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

Overall:

I think the idea is good and has potential. However, some of the design choices seem over-complicated given the simple motivation, and either need to be simplified or justified better.  

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

References:

Li, Y., Akbar, S., \& Oliva, J. (2020, November). ACFlow: Flow models for arbitrary conditional likelihoods. In International Conference on Machine Learning (pp. 5831-5841). PMLR.

Nash, C., \& Durkan, C. (2019, May). Autoregressive energy machines. In International Conference on Machine Learning (pp. 1735-1744). PMLR. The authors have described some of the limitations of their method, such as the computational cost of averaging over multiple orderings of the dimensions. 

 Thank you for the clear and detailed reply. My concerns have been addressed, so I will be raising my score by 1 point. 

I just have a few recommendations for the revision below:

- I would recommend adding a sentence motivating importance sampling over the grid. In particular I liked the point about how simulating from the proposal is very cheap, so we save on the number of neural network passes overall.

- Some discussion on the skip-connection and multiple latent vectors would be useful motivation for readers who are familiar with AEMs (1 or 2 sentences should suffice).

- The stochastic approximation to the autoregressive score is an interesting point - maybe a comment in the paper that points to a derivation/discussion in the Appendix would be really helpful. I would also recommend discussing the random sampling of the vector $o$ earlier on if possible, e.g. in the paragraph starting Line 195. 


 Thank you for your thoughtful questions. We clarify these points below, and shall add clarifications to the manuscript accordingly. We hope that you recommend the article for publication as we also agree that this work has great potential and would be of interest to the ML community.

As with Nash, C., & Durkan, C. (2019), we did not observe any practical limitations based on the random-initialization of the energy network. One could guarantee a finite integral by utilizing the empirical support of the data to restrict the energy function.

Let us motivate the use of a proposal distribution below. We would like to obtain a good estimate of $Z$ using few (e.g. 20) samples since each sample incurs the overhead for evaluating a neural network to compute $\mathcal{E}(x_{u_i} | x_{o})$. A denser grid may yield a usable estimate of $Z$. However, since distributions may span over a large domain and may be non-smooth, importance sampling with a good proposal distribution will be much more sample efficient. This allows one to obtain as good an estimate as a grid with much fewer samples, which in turn will be much more computationally efficient since one will compute $\mathcal{E}(x_{u_i} | x_{o})$ on fewer points. (One would have to compute $\mathcal{E}(x_{u_i}=g_j | x_{o})$ on each grid point, $g_j$.) Note that our proposal distribution is a good approximator as indicated by the likelihoods we obtained empirically. Furthermore, sampling from $q(x_{u_i} | x_{o})$ is actually very efficient as it only requires one neural network evaluation to obtain the parameters of a 1d-GMM, which is efficient to sample from. Thus our importance sampling scheme nets much fewer neural network evaluations since a single pass to obtain the proposal distribution’s parameters saves the many evaluations of the energy network that would be required on a grid.

The energy network could operate over latent vectors or on observed features. We chose to include both because including inputs acts as a skip-connection (which often helps with a deeper network) and early experiments indicated that this led to better performance. Our motivation for using multiple latent vectors is that this allows them to represent the fact that two different unobserved features may depend on $x_o$ in different ways, and the model can start making such distinctions sooner (i.e. at the end of the proposal net instead of inside the energy net).

We have two primary reasons for training non-autoregressively. The first is a matter of speed per batch of instances. Training autoregressively would require evaluating every 1D conditional in the chain-rule decomposition of $p(x_u | x_o)$. This is certainly possible, however it would require effectively expanding the batch size during training to include all of those conditionals. For high-dimensional data, this can quickly become infeasible, or at best make training very slow. Our second motivation is that the non-autoregresive scheme used in the paper is effectively equivalent to stochastically approximating the autoregressive procedure. Randomly sampling one of the conditionals from the autoregressive chain of each example yields our training scheme. We do indeed randomly sample $o$ at each training step, so over the course of training, the model will see the same distribution of inputs under either approach.


 Thank you for your thoughtful comments and questions.

First, let us clarify the terms in which we will talk about a model’s sensitivity to ordering. A model can be *order-agnostic* (as ACE is), which is to say that autoregressive likelihoods can be evaluated in any order. As you point out, ACE is not what we will call *order-consistent*, meaning that evaluating the same likelihood with different orderings can produce different values. The advantage of being order-agnostic, even when not order-consistent, is tied to the advantages of arbitrary conditioning. By definition, since ACE can evaluate any arbitrary 1D conditional, it must be order-agnostic. That is, ACE can evaluate the 1d conditional of an $x_j$ given any subset of features $x_o$, thus it is able to evaluate likelihoods using any order. Hence, our motivation for this type of training is intertwined with our larger motivations for doing arbitrary conditioning.

It is correct that the energy, the function $\mathcal{E}$, is non-negative, and we parameterize it as the output of a neural network. In order to enforce the nonnegativity constraint, the network has a softplus activation on its final layer.

The accumulation of errors is not impactful due to the following. First, from a practical perspective, the curse of dimensionality makes it computationally infeasible to estimate the normalizer for the full distribution $p(x_u | x_o)$ over all dimensions in $u$ all at once. Thus an alternative scheme is needed. Also, while it is true that we are summing the log-normalizers across dimensions, the error present in the normalizer for a given dimension does not affect the estimate of the normalizer for any other dimension. That is, each dimension’s normalizer is computed independently of the others. Thus, we don’t expect to see a compounding effect in which errors accumulate and consequently exacerbate future errors. Finally, if the errors are symmetrically distributed around zero (which Figure 5 in Nash and Durkan 2019 would suggest), we can expect them to cancel out on average, rather than grow.

Even when estimating 1D conditionals, normalizing flows fall under certain restrictions. Namely, all transformations must be invertible and must have Jacobians that are efficiently computed. These restrictions inevitably place limitations on the network's architecture and capacity, which is why we advocate for an approach that allows the use of arbitrary networks.

Ignoring the derivatives of the proposal samples effectively means the proposal distribution is treated independently to the energy distribution (i.e. is not tuned to the normalizer) which leads to a proposal that attempts to fit the distribution (which in itself is desirable for importance sampling).

When estimating the mean of the energy distribution (Equation 9) , we have found that relatively few samples are needed in order to produce reliable estimates, likely due to the high quality of the proposal distribution. For example, on the Gas dataset, when using only 100 importance samples, the average variance of the estimates across 10 seeds is small (~0.003). In our reported results, we use from 3000 to 20000 samples, depending on the dataset.

 Thank you for your thoughtful comments and questions. We believe that the gain in performance stems primarily from decomposing the arbitrary conditioning problem into 1d conditional estimation, and parameterizing these 1d conditionals flexibly with an energy network. In terms of training cost, ACE performs similarly to ACFlow in most cases. Both methods can be trained in anywhere from a few hours to 1-2 days (on a single GPU), depending on the dataset. When evaluating likelihoods, ACE (in contrast to general ACFlow models) can compute the autoregressive components in parallel, but may be more memory intensive.

 Thank you, we are glad that you found the work to be clever, and we hope you recommend our paper for publication. Let us clarify your question. When evaluating $p(x_2, x_3 | X_1=x_1)$, one must evaluate the likelihood at some particular values for features; i.e., one would evaluate $p(X_2=x_2, X_3=x_3 | X_1=x_1) = p(X_2=x_2|X_1=x_1) p(X_3=x_3|X_1=x_1,X_2=x_2)$. This is akin to the case with a Gaussian where one may ask the likelihood of observing $X_2=x_2, X_3=x_3$ given that $X_1=x_1$. Thus, when computing the normalizer $Z_{{x_3}|X_1=x_1, X_2}$, we must normalize for the case when $X_2=x_2$ (i.e., $Z_{{x_3}|X_1=x_1, X_2=x_2}$), since we are seeking to normalize the likelihood $p(X_3|X_1=x_1,X_2=x_2)$ (where we are given that $X_1=x_1, X_2=x_2$). Thus, the proposal distribution samples are only used for estimating the respective normalizers for 1d conditionals and are not percolated up to condition; only the values that one is evaluating the likelihood on ($X_1=x_1, X_2=x_2, X_3=x_3$) are ever used in conditioning. When imputing one is no longer evaluating a likelihood, but is instead drawing a sample of $X_u$  given $X_o=x_o$, or computing $E[X_u | X_o=x_o]$. We shall add some of these details to further clarify such subtleties. With regards to the number of missing variables in the proposal (or energy) network, it is not the case that there is always the same number. During training, we randomly sample the missing features for each batch (see lines 244-246 for additional details), so the networks are exposed to observed sets of different sizes.

In this paper, the authors propose to compute p(x_u|x_o) for any combination of observed variables (x_o) and unobserved varaibles (x_u) with a single NN. The authors actually propose to compute p(x_u|x_o) = prod p(x_{u_i}|x_o,x_{u_1} … x_{u_{i-1}}). Basically, they train a 1D conditional estimator that that can depend on any of the inputs. 

The proposed mechanism is based on two networks. A Proposal network that takes as input a mask for the observed inputs and a vector with the inputs for the observed variables at the right positions. The output of the proposal network is the components for a mixture of Gaussians ($\omega$) and a latent vector for each missing component of x. The second network takes the latent vector and estimates the probability.   I must say that the algorithm is very clever, and it seems to perform well compared to other methods. I would be in favor of having the paper presenting at the conference. But I have a few questions that I do not understand how the algorithm works at inference time, and I would like the authors to be able to explain how their solution performs. Let me start with a very simple example, this would probably be a good simple example for the readers too. Let us assume that we have 4 variables x_1, x_2, and x_3 coming from a multivariate Gaussian. Let’s assume that we train the proposal network with two missing variables and one observed. Side question: Do we always have to have the same number of missing variables in the proposal network, I would say so, but please let me know if I am wrong about this? 

At test time, we are given x_1 = X1 and we compute:
P(x_2, x_3|X_1=X1) = p(x_2|x_1=X1) p(x_3|x_1=X1,x_2)

The first thing that, we need to do is to get N samples from q(x_2|x_1=X1) and Z_{{x_2}|x_1=X1}, using equation 5.

Now, how do I compute, Z_{{x_3}|x_1=X1, x_2}. This partition function is no longer a number but a function of x_2. Do I compute this for a specific value of x_2, such as a sample from p(x_2|x_1=X1) or the maximum or the true value that x_2 takes? I guess it is the last one, using this procedure would only work if we are computing the likelihood for the full sample. So Z_{{x_3}|x_1=X1, x_2} is not a function of x_2 but a constant with x_2=X2.

I think this is more evident when looking at the imputation in which there is no access to X2 and the authors use the following approximation:
P(x_2, x_3|X_1=X1) = p(x_2|x_1=X1) p(x_3|x_1=X1)
And use the mean p(x_2|x_1=X1) and p(x_3|x_1=X1) to impute the variables. For example, using this to generate the top half of an MNIST image would be horrible and would only provide a blurred image. In other image processing would even be worst. 
 Not needed

The authors propose a strategy to estimate any arbitrary conditional distribution by using an the chain rule:  they estimate all one-dimensional distributions.  While this would require learning exponentially many distributions, a weight sharing scheme is used to limit the number of parameters that need to be learned from data.  The resulting approach represents each univariate distribution in the form exp(neural network) and the normalizing constant is estimated via importance sampling (a proposal distribution can also be estimated if desired).  They show that this simple approach is flexible and outperforms more sophisticated approaches on real-world data sets.  Overall, I quite like the approach. It's a simple and clever combination of modern techniques that should be easy to replicate.  My only real criticisms are related to the experimental evaluation (see below) and these are minor.

Originality:  The work appears to be a novel approach for density estimation tasks.

Quality:  1) I don't personally love using likelihoods to evaluate continuous models (despite the prevalence of these types of evaluations).  While the average conditional likelihood is definitely better, it still isn't great.  I've seen plenty of examples with sum-product networks and similarly expressive models where high likelihoods does not correspond to good model performance on real tasks.  Thankfully, the imputation/predictive performance is also better for ACE.  Honestly, I'd suggest moving the likelihood stuff to the appendix and replacing it with larger figures (the bar charts in particular are difficult to read) and more experiments.

2)  I actually find Figure 2 somewhat perplexing.  Why does ACE seem to do so much better in the missingness .1 regime than the other competitors? The performance seems more comparable in the other two settings.  Thoughts?

3) While the current experiments are sufficient, in my opinion, for publication, I think it would make an even stronger case if the authors could include a few more interesting data sets (or even one practical domain) for which this approach succeeds where other approaches fail significantly.

4) No mention is really made of wall-clock comparison between the methods (or I missed it).  Could the authors comment a bit?

Clarity:  The paper is well-written and easy to follow with only minor typos.

Significance: The simplicity of the approach lends itself to a variety of different applications. n/a + see above

Meta-review: This paper proposes an approach for performing conditional density estimation of the form $p(x_u|x_o)$ for any arbitrary choice of the observed ($x_o$) and unobserved ($x_u$) variables by using an energy function parametrized by a neural network to estimate the 1D distributions.

Three reviewers recommend acceptance, and one reviewer believes the paper is below the acceptance threshold. After reading the reviews and the rebuttal, I recommend this paper for acceptance. I also encourage the authors to improve the paper by considering the reviewers' comments. In particular:
+ Add an analysis of the number of samples needed for the importance sampling approximations.
+ As noted by reviewer AEG1, the imputation method may not be appropriate for certain applications because the imputation is performed using the marginal means, e.g., the means of $p(X_2|X_1=x_1)$ and $p(X_3|X_1=x_1)$, therefore, ignoring the dependences between $X_2$ and $X_3$. This should be mentioned/discussed in Section 4.4.2 or 5.1.2. 
+ Add a sentence motivating importance sampling over the grid.
+ Add some discussion on the skip-connection and multiple latent vectors.
+ Include a derivation or discussion of the stochastic approximation of the autoregressive score in the Appendix.

Decision: accept