Bijective-Contrastive Estimation
Abstract
In this work, we propose Bijective-Contrastive Estimation (BCE), a classiﬁcation-based learning criterion for energy-based models.
We generate a collection of contrasting distributions using bijections, and solve all the classiﬁcation problems between the original data distribution and the distributions induced by the bijections using a classiﬁer parameterized by an energy model.
We show that if the classiﬁcation objective is minimized, the energy function will uniquely recover the data density up to a normalizing constant.
This has the beneﬁt of not having to explicitly specify a contrasting distribution, like noise contrastive estimation.
Experimentally, we demonstrate that the proposed method works well on 2D synthetic datasets.
We discuss the diﬃculty in high dimensional cases, and propose potential directions to explore for future work.
1. Introduction
Training energy-based models (EBM) via maximum likelihood estimation (MLE) is diﬃcult.
This is because likelihood evaluation under the density of the EBM requires computing the log-partition function, which involves an intractable integral.
For gradient-based optimization, one often needs to resort to Markov chain Monte Carlo (MCMC) methods in order to approximately sample from the energy model to estimate the gradients (Hinton, 2002).
However, MCMC methods are notoriously computationally expensive and hard to tune, which prevents EBMs from being widely adopted among the practitioners.
This computational bottleneck has therefore motivated the design of alternative training objectives of
EBMs, such as score matching (SM, Hyv¨arinen (2005)), denoising score matching (DSM, Vincent (2011)), noise-contrastive estimation (NCE, Gutmann and Hyv¨arinen (2010)), and Stein discrepancy minimization (Liu et al., 2016; Grathwohl et al., 2020).
In our attempt to devise a new training criterion for EBMs, we draw inspiration from the recent development of ﬂow-based methods.
Notably, Dinh et al.
; Nielsen et al.
(2020) demonstrate that the inverse model of a surjection map recovers the internal structure of the data distribution to a certain degree.
In order to reconstruct the data distribution with a higher precision, a chain of bijection and surjection maps needs to be composed, resulting in solving multiple classiﬁcation problems in sequel.
Motivated by this, we propose to use bijections to generate a collection of classiﬁcation problems and solve them by parameterizing the classiﬁers using a shared energy function.
We refer to the corresponding parameter estimation principle as Bijective-Contrastive Estimation (BCE).
As the EBM is trained to contrast the distributions induced by these bijections, it will learn to recover the ratio of the corresponding density functions.
By jointly solving a suﬃciently large collection of classiﬁcation problems with shared parameters, the EBM is guaranteed to recover the data distribution’s density up to a normalizing constant.
2. Bijective-Contrastive Estimation (BCE) 2.1. Generate Contrasting Distributions via Bijections
In this section, we ﬁrst show how to generate a binary classiﬁcation problem using a bijection and demonstrate how to solve the classiﬁcation problem with an energy model.
We then conjecture that while using the energy model to solve a single classiﬁcation problem will not be enough to recover the data density, we can do so by jointly solving suﬃciently many classiﬁcation problems.
Let x be a random variable following the density p ( x ), and deﬁne x (cid:48)
=
x + (cid:15) , where (cid:15) is a constant.
Then the density of x (cid:48) is equal to p ( x (cid:48) − (cid:15) ); see Figure 1.
We consider a binary classiﬁcation problem between these two distributions, which admits an optimal solution π ∗ ( x ) =
[Equation]
e. the blue curve).
This implies we can plug in a parametric density, or an unnormalized density using an energy model in place of p ( x ) as well as p ( x − (cid:15) ), and hope to recover p ( x ) by training the classiﬁer.
Nevertheless, we can’t recover p ( x ) solely from π ∗ .
This is because diﬀerent density functions p ( x ) can possibly generate very similar optimal solution.
For example, the left hand side of Figure 1 shows two possible p ( x )’s that result in almost indistinguishable π ∗ ’s for the same (cid:15) .
While the optimal classiﬁer generated by a single bijection may not be distinguishable for diﬀerent density functions, we conjecture that the family of optimal classiﬁers generated by a suﬃciently large collection of bijections will be uniquely dependent on the density function.
Consider the same classiﬁcation task above, but now with multiple (cid:15) ’s. In Figure 1 (right), we plot all the p ( x − (cid:15) )’s
(as well as the respective optimal classiﬁers π ∗ ) corresponding to diﬀerent values of (cid:15) ; the transparency of the curves reﬂects the magnitude of the (cid:15) value.
Unlike the single classiﬁcation case, we can now easily tell the two families of π ∗ ’s apart.
This reveals that with an increasing number of (cid:15) ’s, there will be less available densities that match all of the optimal classiﬁers of the generated problems.
With a suﬃciently large number of classiﬁcation problems, we hope to rule out all but one density function – the true density of the data distribution.
This would imply that we can learn the data density by parameterizing the classiﬁers as a function of the energy and the bijections, and jointly solving the classiﬁcation problems.
In the next section, we formalize and generalize these examples, and propose a new classiﬁcation-based learning principle for EBMs.
2.2. Bijective-Contrastive Estimation
Let x ∈ R
d ∼ p data ( x ) be a random variable representing the data, and t : R d → R d be a bijection.
The probability density of x t := t ( x ) can be conveniently expressed as p ( x t ) =
p data ( t − 1 ( x t ))
| J t − 1 ( x t ) | by the change-of-variable formula, where J f ( x ) denotes the Jacobian matrix of f wrt x .
Similarly to the examples in the previous section, let us ﬁrst consider a classiﬁcation problem to diﬀerentiate between x and x t for ﬁxed t when an observation is drawn from p data ( x ) or p ( x t ) with equal probability.
Let π ( x ) be a probabilistic classiﬁer π : R d → [0 , 1], solving the following objective:
[Equation]
The objective has a functional minimizer
[Equation]
That is, L ( π ∗ ; t ) ≤ L ( π ; t ) for any π .
Let p ( x ; f ) =
e − f ( x ) /Z
( f ), where f is the energy model and Z ( f ) is the normalizing constant.
We can parameterize the classiﬁer as
[Equation]
Plugging π ( x ; f ) into Equation (1), we can rewrite L as a loss functional of f
[Equation]
[Equation]
As discussed in the previous section, we conjecture that f ∗ satisfying e − f ∗ ( x ) ∝ p data ( x ) will be the only minimizer of L ( f ; t ) for a suﬃciently large collection of t .
Motivated from this, we propose to minimize an expected loss over a distribution of bijections:
[Equation] [Equation]
We refer to learning energy functions by minimizing Equation (4) (expectation taken over p data or an empirical distribution) as Bijective-Contrastive Estimation (BCE) .
We call the binary classiﬁcation form of Equation (4) binary BCE objective .
Let p ( t ) 1 denote the density of the bijection t .
The following theorem provides a positive answer to our conjecture: under some mild assumption on p ( t ) and the form of t , the minimizer f ∗ of the BCE loss is unique up to a normalizing constant and e − f ∗ ( x ) is proportional to the data density.
Theorem 1 (Uniqueness of BCE-based estimator)
Let F be the set of measurable functions.
For any x ∈ supp( p data ( x )) , deﬁne g x ( t ) =
t − 1 ( x ) .
If for any x ∈ supp( p data ( x )) and for any p ( t ) -almost sure set M , g x ( M ) c has Lebesgue measure zero, then p data ( x ) ∝ e − f ∗ ( x ) almost everywhere
[Equation]
Corollary 2 (Additive BCE)
If t ( x ) := x + (cid:15) where (cid:15) ∈ R
d is a random variable with a density function p ( (cid:15) ) that is non-zero everywhere, then the minimizer f ∗ of the BCE objective in Equation (4) satisﬁes p data ( x )
∝ e − f ∗ ( x ) almost everywhere.
The proofs are deferred to Appendix A. Theorem 1 shows that BCE is powerful enough to recover p data , up to a normalizing constant.
Moreover, according to Corollary 2, this holds true for very simple classes of bijections and distributions; for example, additive noise with a standard normal distribution.
2.3. Variants of objective functions
For additive bijections, we observe that t and t − 1 have the same density if p ( (cid:15) ) is standard normal.
In this paper, if the distribution of bijections satisﬁes p ( t ) = p ( t − 1 ), we informally refer to this property as inverse-symmetry .
Inverse-symmetry allows us to simplify Equation (4) to (neglecting a multiplier of 2)
[Equation]
In addition to additions, random permutations and matrix-vector product of data with orthogonal matrices are also simple bijections that are inverse-symmetric.
Inverse-symmetry also largely simpliﬁes the computation of the multiclass classiﬁcation loss, which would normally have quadratically many terms, to be
[Equation]
where t 1 , . .
, t n ∼ p ( t ) (neglecting a multiplier of n ).
We ﬁnd that training with multiclass
BCE tends to be more stable as it has smaller gradient variance.
2.4. Relation to noise contrastive estimation
The most closely related work to ours is noise-contrastive estimation (NCE, Gutmann and
Hyv¨arinen (2010)), which uses a tractable noise distribution q ( x ) and learns the energy function via the following objective,
[Equation]
If q ( x ) is deﬁned to be nonzero wherever p data ( x ) is nonzero, then e − f ∗ ( x )
= p data ( x ) if f ∗ minimizes the NCE objective 2 .
While NCE contrasts against a pre-deﬁned noise distribution q , BCE uses bijections to generate many contrasting distributions implicitly.
3. Experiment on 2D synthetic datasets
In order to demonstrate that by minimizing the BCE loss, we can learn the correct data density, we run density estimation experiments on six 2D synthetic datasets.
We use binary
BCE loss as in Equation (6), additive bijections for the perturbation, and standard normal distribution for the additive bijections.
We use a ReLU network with three 1,000-unit hidden layers, and train it using the Adam optimizer (Kingma and Ba, 2015) with β 1 = 0 .
9 , β 2 = 0 .
We run 50K iterations.
Mini-batch size and learning rate are set to 256 and 0.0001, respectively.
The results are presented in Figure 2.
We see that the energy-based model trained by BCE successfully learns the densities of a variety of 2D synthetic datasets.
4. Diﬃculties in modeling high-dimensional data
In preliminary experiments, we observe that the proposed method did not work well on high dimensional data.
On the MNIST dataset, we train the EBM via BCE and evaluate the quality of samples generated by running Langevin dynamics (Grenander and Miller, 1994), but the generated samples all fail to resemble the true data samples.
We conjecture that this is a result of using the cross entropy loss, which causes two types of vanishing gradients problems: the loss does not provide meaningful learning signals when t ( x ) is either too far from or too close to the data manifold.
The ﬁrst case occurs when we use a simple family of bijections for BCE.
For instance, adding an unstructured noise can easily push a highly structured data point oﬀ the data manifolds.
Hence the classiﬁers can achieve near-perfect accuracy while the learned unnormalized density is nowhere near the true one.
The gradients wrt the energy can also vanish when p ( t ( x )) is too close to the data density, e.g. when t is near identity.
Similarly to BCE, the ﬁrst type of vanishing gradient problem also prevents NCE from scaling up, occurring when q ( x ) is too far from p data ( x ).
Recently, Rhodes et al.
(2020) tackle this problem by generating a sequence of NCE problems, each of which contrasts two consecutive intermediate distributions interpolated between the data and a noise distribution.
When the two distributions are close enough, one can potentially mitigate the vanishing gradient problem.
In addition, generative adversarial networks (GANs, Goodfellow et al.
(2014)) also suﬀer from similar vanishing gradient problems, which can be addressed by regularizing the discriminator (Arjovsky et al., 2017; Roth et al., 2017).
These techniques can potentially be applied to BCE for modelling high-dimensional data.
5. Conclusion
In this paper, we propose a new classiﬁcation-based EBM training method, called BijectiveContrastive Estimation (BCE).
We prove that the data density (up to a normalizing constant) is the only minimizer of the BCE loss.
In the experiment, we show that the energy trained with BCE accurately models the 2D synthetic data distributions.
We discuss the diﬃculty of BCE in modeling high-dimensional data due to vanishing gradients, and suggest future directions for addressing this issue.
6. Acknowledgments
We thank Samsung for funding this work.
Appendix A. Proofs Proof of TheoremStep
In order to prove Theorem 1, we ﬁrst introduce another optimization problem equivalent to minimizing L BCE - b ( f ) in Equation (4).
Consider the following data generation process.
[Equation]
where ˜ x ∼ p data (˜ x ), t ∼ p ( t ), and y ∼ Bernoulli (cid:0) 12 (cid:1) .
Then the likelihood of x given y and t is p ( x | y = 1 , t ) =
p data ( x ) or p ( x | y = 0 , t ) =
p data ( t − 1 ( x ))
| J t − 1 ( x ) | by the change of variable.
Similarly, we deﬁne the “model” likelihood functions
[Equation]
Then
[Equation]
the minimization of the BCE objective in Equation (4) is equivalent to minimizing the expected KL-divergence between the true posterior p ( y | x, t ) and the model posterior q ( y | x, t ):
[Equation]
where
[Equation]
and q
[Equation]
| x
[Equation]
t
[Equation]
The equality holds since L BCE - b ( f ) is the cross entropy term in the expected KL and the entropy term of p ( y | x, t ) does not depend on f .
Step
Due to Equation (9), proving Theorem 1 is equivalent to showing that if g x maps any p ( t )-almost sure set to a Lebesgue-almost everywhere set, then
[Equation]
Proving from the LHS to the RHS is straightforward.
Hence, we only show the other direction.
Assume the RHS is true, since the integrand is non-negative, for almost every x , we have
[Equation]
which then implies for almost all t (depending on x ), the KL divergence is 0.
KL divergence is 0 if and only if the distributions are identical, which means p ( y | x, t ) =
q ( y | x, t ; f ).
That is, turning to the case where y = 1, we have that for almost every x (wrt the marginal probability p ( x ) =
(cid:82) p ( x, t ) dt ),
[Equation]
for almost all t (wrt the conditional probability p ( t | x )).
We can pick x so that p data ( x ) > 0, which will then imply
[Equation]
Since this holds true for p ( t | x )-almost all t , p data ( t − 1 ( x ))
∝ e − f ( t − 1 ( x )) almost surely.
Finally, since p ( t | x ) has the same support as p ( t ), a p ( t | x )-almost sure set is a p ( t )-almost sure set. Then by the assumption on g x , we have p data ( x ) ∝
e − f ( x ) almost everywhere on R d .
(cid:4)
Proof of Corollary
For any ﬁxed x ∈ supp( p data ( x )), we deﬁne g x ( (cid:15) )
= x − (cid:15) .
Since p ( (cid:15) ) is non-zero everywhere, Lebesgue measure λ is absolutely continuous wrt the probability measure ν corresponding to the density function p ( (cid:15) ).
As Lebesgue measure is translation invariant and g x is invertible, for any ν -almost sure set M , ν
( M c )
=
⇒ λ
( M c )
=
⇒ λ
( g x ( M ) c )
= 0.
Then the rest follows by an application of Theorem 1.
(cid:4)