A NALYZING AND I MPROVING THE O PTIMIZATION L ANDSCAPE OF N OISE -C ONTRASTIVE E STIMATION
Anonymous authors
Paper under double-blind review
A BSTRACT
Noise-contrastive estimation (NCE) is a statistically consistent method for learning unnormalized probabilistic models.
It has been empirically observed that the choice of the noise distribution is crucial for NCE’s performance.
However, such observations have never been made formal or quantitative.
In fact, it is not even clear whether the difﬁculties arising from a poorly chosen noise distribution are statistical or algorithmic in nature.
In this work, we formally pinpoint reasons for NCE’s poor performance when an inappropriate noise distribution is used.
Namely, we prove these challenges arise due to an ill-behaved (more precisely, ﬂat) loss landscape.
To address this, we introduce a variant of NCE called eNCE which uses an exponential loss and for which normalized gradient descent addresses the landscape issues provably when the target and noise distributions are in a given exponential family.
I NTRODUCTION
Noise contrastive estimation (NCE) is a method for learning parameterized statistical models (Gutmann & Hyvärinen, 2010; 2012).
To estimate a distribution P ∗ , NCE trains a discriminant model to distinguish between samples of P ∗ and a known distribution Q of our choice, often referred to as the “noise” distribution.
If the function class for the discriminant model is representationally powerful enough, the optimal model learns the density ratio p ∗ /q , from which we can extract the density p ∗ since q is known (Menon & Ong, 2016; Sugiyama et al., 2012).
Compared to the well-studied maximum likelihood estimation (MLE), NCE avoids calculating the (often intractable) partition function,while maintaining the asymptotic consistency of MLE (Gutmann & Hyvärinen, 2012).
It is empirically well-documented that the choice of the noise distribution Q is crucial to both the statistical and algorithmic efﬁciency of NCE (Gutmann & Hyvärinen, 2010; 2012; Rhodes et al., 2020;
Goodfellow et al., 2014; Gao et al., 2020).
However, it has been observed in practice that even when following the standard guidelines for choosing Q , NCE can still yield parameter estimates far from the ground truth (Rhodes et al., 2020; Goodfellow et al., 2014; Gao et al., 2020).
Most recently,
Rhodes et al.
identiﬁed a phenomenon they call the “density chasm,” observing empirically that NCE performs poorly when the KL divergence between P ∗ and Q is large.
One example is when P ∗ , Q are both tightly concentrated unimodal distributions with faraway modes; the region between the two modes will have a small density under both distributions, thus forming a “chasm”.
While it makes intuitive sense that NCE does not perform well under such settings—since disparate
Q and P ∗ are easy to distinguish and do not require the model to learn much about P ∗ in order to do well on the classiﬁcation task—there has not been a theoretical analysis of this phenomenon.
In fact, it is not even clear whether the difﬁculty is statistical or algorithmic in nature.
In this work, we formally study the challenges for NCE with a ﬁxed Q with a focus on distributions in an exponential family.
We show that when the noise distribution Q is poorly chosen, the loss landscape can become extremely ﬂat: in particular, even when P ∗ and Q are two univariate
Gaussian with unit variance, the loss gradient and curvature can become exponentially small in the difference in their means.
We prove that this poses challenges for standard ﬁrst order and even second-order optimization methods, forcing them to take an exponential number of steps to converge to a good parameter estimate.
Thus, standard approaches to minimizing convex functions such as gradient descent—or even more advanced techniques such as momentum or Newton’s method—are not suited to the NCE objective unless Q is close to P ∗ in KL sense.
To remedy this issue, we study an alternative method for optimizing the NCE objective.
We consider instead Normalized Gradient Descent (NGD) whereby the gradient is normalized to have unit norm at each time step.
Perhaps surprisingly, we prove that this small modiﬁcation can overcome the problem of poor curvature in the Gaussian example.
In general, we show the number of steps for
NGD to converge to a good solution for the NCE loss depends on the condition number κ of the
Hessian of the loss at the optimum—the growth of this condition number is unclear for P ∗ and Q when they belong to an exponential family.
To address this, we propose the eNCE loss, a variant to NCE that replaces the log loss in NCE with an exponential loss, and we show that the resulting condition number is polynomial in the dimension and the parameter distance between P ∗ and Q when they belong to an exponential family.
Our proposed change of loss and optimization algorithm together form the ﬁrst solution that provides a provable polynomial rate for learning the parameters of the ground truth distribution.
Theoretically, both NCE and eNCE can potentially suffer from numerical issues during optimization when P ∗ and
Q are far—this is an interesting direction for future work.
Nonetheless, we ﬁnd this to be a simple and effective ﬁx to the ﬂatness of the loss landscape in many settings, as evidenced by experimental results on synthetic and MNIST dataset.
1. R ELATED W ORK
NCE and its variants have inspired a large volume of research in NLP (Mnih & Teh, 2012; Mnih &
Kavukcuoglu, 2013; Dyer, 2014; Kong et al., 2020) as well as computer vision (Oord et al., 2018;
Hjelm et al., 2018; Henaff, 2020; Tian et al., 2020).
It has been observed empirically that NCE with a ﬁxed noise Q is often insufﬁcient for learning good generative models.
The predominant class of approaches that have been proposed to overcome this issue aim to do so by not using a ﬁxed
Q but by iteratively solving multiple NCE problems with an updated Q , or equivalently updated discriminators.
This includes the famous generative adversarial network (GAN) by Goodfellow et al.
(2014), which uses a separate discriminator network updated throughout training.
In a similar vein, Gao et al.
(2020) also aimed to increase the discriminative power as the density estimator improves, and parameterize Q explicitly with a ﬂow model.
More recently, Rhodes et al.
(2020) proposed the telescoping density ratio estimation, or TRE, which sidesteps the chasm by expanding p ∗ /q into a series of intermediate density ratios, each of which is easier to estimate, leading to strong empirical performance—though their work carries no formal guarantees.
With respect to a ﬁxed Q , it remains an open question about what formally are the nature of the challenges posed by a poorly chosen Q , which could be statistical and/or algorithmic.
Various previous works have analyzed the asymptotic behavior of NCE and its variants (Gutmann & Hyvärinen, 2012; Riou-Durand et al., 2018; Uehara et al., 2020), but these do not provide guidance on the ﬁnite step convergence of NCE or its common variants.
The improvements to NCE in prior works are all borne out by the empirical observations of NCE practitioners, rather than motivated by theory, which is precisely the aim of this work.
Finally, we would like to note that prior work has proposed “generalized NCE" (Pihlaja et al., 2010;
Gutmann & Hirayama, 2011; Uehara et al., 2020), which relates the NCE objective to minimizing the Bregman divergence.
Generalized NCE says that we can design a family of training objectives by using different convex functions to deﬁne the Bregman divergence, and the proposed eNCE is an instance of the generalized NCE objective.
The difference between these prior work and ours is again the different focuses on asymptotic behavior versus ﬁnite step convergence.
P RELIMINARIES
The NCE objective Let P ∗ denote an unknown distribution in a parametric family { P θ } θ ∈ Θ , for some bounded convex set Θ , with P ∗ =
P θ ∗ .
Our goal is to estimate P ∗ via P θ for some θ ∈ Θ by solving a noise contrastive estimation task.
The noise distribution Q belongs to the same parametric family with parameters θ q ∈ Θ , so that Q = P θ q .
We use p θ , p ∗ , q to denote the probability density functions (pdfs) of P θ , P ∗ , and Q ; we may omit θ in P θ , p θ when it is clear from the context and write P, p instead.
Given P ∗ and Q , the NCE loss of P is deﬁned as follows:
Deﬁnition 2.1 (NCE Loss) .
The NCE loss of P θ w.r.t.
data distribution P ∗ and noise Q is:
[Equation]
The NCE loss can be interpreted as the binary cross-entropy loss for the classiﬁcation task of distinguishing the data samples from the noise samples.
Moreover, the NCE loss has a unique minimizer:
Lemma 2.1 (Gutmann & Hyvärinen 2012) .
The NCE objective in Deﬁnition 2.1 is uniquely minimized at P =
P ∗ , provided that the support of Q covers that of P ∗ .
Exponential family.
We focus our attention on the exponential family, where the pdf for a distribution with parameter θ is p θ ( x ) = exp (cid:16) θ (cid:62) ˜ T ( x )
− A
( θ ) (cid:17) , with ˜ T ( x ) denoting the sufﬁcient statistics and A ( θ ) the log partition function.
The partition function is treated as a parameter in
NCE, so we use τ to denote the extended parameter, i.e. τ :=
[ θ, α ] where α is the estimate for the log partition function.
We accordingly extend the sufﬁcient statistics as T ( x ) =
[ ˜ T ( x ) , − 1] to account for the log partition function.
The pdf with the extended representation is now simply p τ ( x ) =
exp( τ (cid:62) T ( x )) .
We will use the notation P θ and P τ interchangeably.
We will also use τ ( θ ) to denote the log-partition extended parameterization when the log partition function α properly normalizes the distribution speciﬁed by θ .
A compelling reason for focusing on the exponential family is the observation that the NCE loss is convex in the parameter τ :
Lemma 2.2 (NCE convexity) .
For exponential family p θ,α ( x ) = h ( x ) exp( θ (cid:62) ˜ T
( x ) − α ) , the NCE loss is convex in parameter τ :=
[ θ, α ] .
Lemma 2.2 has been stated under more general settings by Uehara et al.
; an alternative selfcontained proof is included in Appendix A for completeness.
Recall that Θ denotes the set of parameters without the extended coordinate for the log partition function.
We assume the following on distributions supported on Θ :
Assumption 2.1 (Bounded parameter norm) .
(cid:107) θ (cid:107) 2 ≤ ω , ∀ θ ∈ Θ .
Assumption 2.2 (Lipschitz log partition function) .
Assume the log partition function is β ZLipschitz, that is, ∀ θ 1 , θ 2 ∈ Θ , | log Z ( θ 1 ) −
log Z ( θ 2 ) |≤ β Z (cid:107) θ 1 − θ 2 (cid:107) .
Assumption 2.3 (Bounded singular values of the population Fisher matrix) .
There exist λ max , λ min > 0 , such that ∀ θ ∈ Θ , we have σ max ( E θ
[ T ( x ) T ( x ) (cid:62) ]) ≤ λ max , and σ min ( E θ
[ T ( x ) T ( x ) (cid:62) ]) ≥ λ min .
Assumption 2.4 (Smooth change in the Fisher matrix) .
Assume the maximum and minimum singular values of the Fisher matrix change smoothly.
Namely, there exist constants γ max , γ min > 0 s.t.
[Equation]
[Equation] (cid:107)≤ γ min Assumptions 2.2-2.4 can be viewed as smoothness assumptions on the ﬁrst, second and third order derivatives of the log partition function, and can in fact be viewed as introducing structural parameters of the distributions, rather than extra constraints.
For example, distributions with ﬂatter tails will have a larger λ max , which then translates to a slower rate in the results; distributions that are closer to being singular with have a smaller λ min , etc.
In particular, Assumption 2.3 says the singular values of the Fisher matrix E θ [ T ( x ) T ( x ) (cid:62) ] should be bounded from above and below.
It can be shown that the Fisher matrix is proportional to the Hessian of the NCE objective when using Q =
P ∗ , which means Assumption 2.3 can be interpreted as saying the NCE task can be solved efﬁciently under the optimal choice of Q .
O VERVIEW OF RESULTS
We ﬁrst provide an informal overview of our results, focusing on learning of exponential families.
Flatness of population landscape:
Our ﬁrst contribution is a negative result identifying a key source of difﬁculty for NCE optimization to be an ill-behaved population landscape.
We show that due to an extremely ﬂat landscape, gradient descent or Newton’s method with standard choices of step sizes will need to take an exponential number of steps to ﬁnd a reasonable parameter estimate.
We emphasize that though Gaussian mean estimation is a trivial task, its simplicity strengthens the results above : we are proving a negative result so that failures with a simpler setup means a stronger result.
Moreover, the results only apply to standard choices of step sizes, such as inversely proportional to the smoothness for gradient descent, or to the ratio between the smoothness and strong convexity for Newton’s method.
This does not rule out the possibility that a cleverly designed learning rate schedule or a different algorithm would work efﬁciently; the results are however still meaningful since gradient descent with standard step sizes is the most common choice in practice.
Overcoming ﬂatness using normalized gradient descent: Our second contribution is to show that the ﬂatness problem can be solved by a simple modiﬁcation to gradient descent if the loss is well-conditioned.
Speciﬁcally, we show that the convergence rate for normalized gradient descent is polynomial in the parameter distance and κ ∗ , the condition number of the Hessian at the optimum.
One immediate consequence is that in Gaussian mean estimation, for a target error of δ ∈ (0 , 1 β Z ] in parameter distance, NCE optimized with NGD achieves a rate of O ( 1 δ 2 ) , which is the same as the optimal rate achieved by MLE.
The remaining question is then whether κ ∗ is polynomial in the parameters of interests.
We show that κ ∗ can be related to the Bhattacharyya coefﬁcient between P ∗ and Q , which indeed grows polynomially in parameter distance under certain assumptions as detailed in Section 5.2.
Polynomial condition number for the eNCE loss: Our third and ﬁnal contribution is that if we modify the NCE objective slightly—namely, use the exponential loss in place of the log loss—then the condition number at the optimum is guaranteed to be polynomial.
We call this new objective eNCE .
Combined with the NGD result, we get that running NGD on the eNCE objective achieves a polynomial convergence guarantee.
We then provide empirical evidence on synthetic and MNIST dataset that eNCE with NGD performs comparatively with NGD on the original NCE loss, and both outperform gradient descent.
F LATNESS OF THE NCE LOSS
In this section, we study the challenges posed to NCE when using a badly chosen ﬁxed Q .
The main thrust of the results is to show that both algorithmic and statistical challenges can arise because the
NCE loss is poorly behaved , particularly for ﬁrst- and second-order optimization algorithms: when
P ∗ , Q are far, the loss landscape is extremely ﬂat near the optimum.
In particular, the gradient has exponentially small norm and the strong convexity constant decreases exponentially fast, limiting the convergence rate of the excess risk.
We further show that when moving from P = Q to P = P ∗ , the loss drops from Θ(1) to a value that is exponentially small in terms of the distance between
P ∗ and Q .
Consequently, common gradient-based and second order methods will take exponential number of steps to converge.
An important note is that our analysis is at the population level, implying that the hardness comes from the landscape itself regardless of the statistical estimators used.
Setup – Gaussian mean estimation : For the negative results in this section, let’s consider an exceedingly simple scenario of 1-dimensional, ﬁxed-variance Gaussian mean estimation.
We will demonstrate the enormous difﬁculty of achieving a good parameter estimate, even for such a simple problem—this bodes ill for NCE objectives corresponding to more complex models in practice, which certainly pose a much more difﬁcult challenge.
In particular, let P ∗ , Q, P be Gaussians with identity variance.
Let θ ∗ , θ q , θ denote the respective means, with θ ∗ being the target mean that NCE aims to estimate.
When the covariance is known to be 1, we can denote h ( x ) := exp( − x 2 2 ) , and parametrize the pdf of a 1d Gaussian with mean θ as p ( x ) = h ( x ) exp( (cid:104) τ
( θ ) , T ( x ) (cid:105) ) , 3 where the parameter is τ ( θ ) :=
[ θ, θ 2 2 + log √ 2 π ] and the sufﬁcient statistics are T ( x ) :=
[ x, − 1] .
4 We will shorthand τ ( θ ) when it is clear from the context.
Without loss of generality, we will assume θ q = 0 , and θ ∗ > 0 , and denote R :
= θ ∗ − θ q . We will write τ ∗ : = τ ( θ ∗ ) =
[ R, R 2 2 + log √ 2 π ] , and τ q := τ ( θ q ) =
[0 , log √ 2 π ] .
As a clariﬁcation, the results stated in this section will be in terms of R , hence the asymptotic notations Ω , O never hide dominating dependency on R .
4. P ROPERTIES OF THE NCE LOSS
We ﬁrst describe several properties of the NCE loss that will be useful in the analysis of ﬁrst- and second-order algorithms.
To start, we show that the dynamic range of the loss is large: that is, the optimal NCE loss is exponentially small as a function of R ; on the other hand, if θ is initialized close to θ q , the initial loss would be on the order of a constant.
Precisely:
Proposition 4.1 (Range of NCE loss) .
Consider the 1d Gaussian mean estimation task with mean θ ∗ , θ q ∈ R , and a known variance of 1.
Denote R := | θ q − θ ∗ | where R (cid:29) 1 , Then, the loss at θ = θ q is log 2 , while the minimal loss L ∗ is L ∗ ( R ) =
c exp( − R 2 / 8) for some c ∈ [ 12 , 2] .
The next shows we need to decrease the loss to be on an order comparable to the optimum value.
Namely, the loss is very ﬂat close to θ ∗ , thus in order to recover a θ close to θ ∗ , we have to reach a very small value for the loss.
Precisely:
Proposition 4.2.
Under the same setup as Proposition 4.1, for a given δ ∈ (0 , 1) , if the learned parameter τ satisﬁes (cid:107) τ − τ ∗ (cid:107) 2 ≤ δ , then L ( τ ) − L ( τ ∗ ) ≤ R exp( − R 2 / 8) δ 2 .
The way we will leverage Propositions 4.1 and 4.2 to prove lower bounds is to say that if the updates of an iterative algorithm are too small, the convergence will take an exponential number of steps.
Proposition 4.2 is proven via the Taylor expansion at θ ∗ : since the gradient is 0 at θ ∗ , we just need to bound the Hessian at θ ∗ .
We show:
Lemma 4.1 (Smoothness at P =
P ∗ ) .
Under the same setup as Proposition 4.1, the smoothness at
P =
P ∗ is upper bounded as σ max ( ∇ 2 L ( τ ∗ )) ≤ R √ 2 π exp( − R 2 / 8) .
We will also need a bound on the strong convexity constant (i.e. smallest singular value) at P
= P ∗ :
Lemma 4.2 (Strong convexity at P = P ∗ ) .
Under the same setup as Proposition 4.1, the minimum singular value at P = P ∗ is σ ∗ min ( ∇ [Equation] [Equation] =
[Equation]
Finally, in order to estimate the choice of the step size for standard optimization methods, we will also need a bound of the smoothness at P = Q :
Lemma 4.3 (Smoothness at P = Q ) .
Under the same setup as Proposition 4.1, the smoothness at
P = Q is lower bounded as σ max ( ∇ 2 L ( τ q )) ≥ R 2 2 .
Lemma 4.1, 4.2 are proved in Appendix D.4, and Lemma 4.3 is proved in Appendix D.5.
4. L OWER BOUNDS ON FIRST - AND SECOND - ORDER METHODS
With the landscape properties at hand, we are now ready to provide lower bounds for both ﬁrst-order and second-order methods.
For ﬁrst-order methods, we show that:
Theorem 4.1 (Lower bound for gradient-based methods) .
Let P ∗ , Q, P be 1d Gaussian with variance 1.
Assume θ q = 0 , θ ∗ > 0 without loss of generality, and assume R :
= θ ∗ − θ q (cid:29) 1 .
Then, gradient descent with any step size η = o (1) from an initialization τ = τ q will need an exponential number of steps to reach some τ (cid:48) that is O (1) close to τ ∗ .
Note, the maximum step size η = o (1) the theorem applies to is actually a loose bound: the standard setting of step size for gradient descent is η ≤ 1 /λ M for λ M := max θ ∈ Θ σ max ( ∇ 2 L ( τ ( θ ))) , which is Ω( R 2 ) by Lemma 4.3.
Theorem 4.1 helps explain why NCE with a far-away Q fails in practice, if we set the budget for the number of updates to be polynomial.
A natural remedy to the drastically changing norms of the gradients is to use methods that can properly precondition the gradient.
This motivates the use of second order methods, which adapt to the geometry of the loss and hence can potentially perform more competitively.
Unfortunately, standard second-order approaches are again of no help, and the number of steps required to converge remains exponential.
Consider Newton’s method with updates of the form η ( ∇ 2 L ) − 1 ∇ L .
At ﬁrst glance, this looks like it may solve the issue of a ﬂat gradient, since the
Hessian ∇ 2 L may also be exponentially small hence canceling out with the exponentially small gradient.
However, the ﬂatness of the landscape forces us to take an exponentially small step size η , resulting in the following claim:
Theorem 4.2 (Lower bound for Newton’s method) .
Let P ∗ , Q, P satisfy the same conditions as in theorem 4.1.
Let λ ρ :=
min θ ∈ Θ σ min ( ∇ 2 L ( τ θ )) , λ M := max θ ∈ Θ σ max ( ∇ 2 L ( τ θ )) .
Then, running the Newton’s method with step size η = O
( λ ρ λ M ) from an initialization τ = τ q will need an exponential number of steps to reach some τ (cid:48) that is O (1) close to τ ∗ .
Again, the condition η
= O (cid:16) λ ρ λ M (cid:17) follows the typical step size choice for Newton’s method, i.e. the step size should be upper bounded by the ratio between the global strong convexity constant and the global smoothness of the function, which is exponentially small for this setup by Lemma 4.2, 4.3.
N ORMALIZED GRADIENT DESCENT FOR WELL - CONDITIONED LOSSES
We have seen that due to an ill-behaved landscape, NCE optimized with standard gradient descent or Newton’s method will fail to reach a good parameter estimate efﬁciently, even on a problem as simple as Gaussian mean estimation, and even with access to the population gradient.
In this section, we will show that a close relative of gradient descent, normalized gradient descent (NGD), despite its simplicity, provides a ﬁx to the ﬂatness problem to exponential family distributions when the Hessian of the loss is well-conditioned close to the optimum.
Precisely, recall that the NGD updates for a loss function L is τ t
+1 =
τ t − η ∇ L ( τ t ) (cid:107)∇ L ( τ t ) (cid:107) 2 .
We assume that in a neighborhood around τ ∗ , the change in the shape of the Hessian H is moderate:
Assumption 5.1 (Hessian in a neighborhood of τ ∗ ) .
Under assumption 2.2 with constant β Z , assume that for any τ such that (cid:107) τ − τ ∗ (cid:107) 2 ≤ 1 β Z , it holds that σ max ( H ( τ )) ≤ β u · σ max ( H ( τ ∗ )) , and σ min ( H ( τ )) ≥ β l · σ min ( H ( τ ∗ )) , for some constant β u , β l > 0 .
The main result of this section states that NGD can ﬁnd a parameter estimate efﬁciently for exponential families, where the number of steps required is polynomial in the distance between the initial estimate and the optimum:
Theorem 5.1.
Let L be any loss function that is convex in the exponential family parameter and satisﬁes Assumptions 5.1 and 2.1 - 2.4.
Furthermore, let P ∗ , Q be exponential family distributions with parameters τ ∗ , τ q and let κ ∗ be the condition number of the Hessian at P
= P ∗ .
Then, for any 0 < δ ≤ 1 β Z and parameter initialization τ 0 , with step size η ≤ (cid:113) β l β u κ ∗ δ , performing NGD on the population objective L guarantees that after T ≤ β u κ ∗ β l · (cid:107) τ 0 − τ ∗ (cid:107) 2 δsteps, there exists an iterate t ≤ T such that (cid:107) τ t − τ ∗ (cid:107) 2 ≤ δ .
The main technical ingredient for proving Theorem 5.1 is the following Lemma:
Lemma 5.1.
Suppose Assumptions 2.2 and 5.1 hold with constants β Z , β u and β l .
Let L be a convex function with minimizer τ ∗ , and let g :=
∇ L ( τ ) .
For any δ ≤ 1 β Z , let γ = (cid:113)
β
l β u κ ∗ δ .
Then for all τ s.t.
(cid:107) τ − τ ∗ (cid:107) 2 ≥
δ , we have L ( τ ∗ + γ g (cid:107) g (cid:107) )
≤ L ( τ ) .
Lemma 5.1 explains the dependency on κ ∗ in the NGD convergence rate.
The intuition of the proof is that in a small neighborhood around τ ∗ , the set of parameters that have the same loss form an “ellipsoid", and by Taylor expansion, any two points in the same set will have a distance-to- τ ∗ ratio upper bounded roughly by √ κ ∗ .
The details are in Appendix C.2.
Proof sketch for Theorem 5.1 : the proof leverages two observations: the convexity of NCE loss on exponential family parameters, and that the Hessian in a neighborhood around τ ∗ changes moderately as stated in Assumption 5.1.
One can then show there exists a global constant γ such that for any τ t satisfying (cid:107) τ t − τ ∗ (cid:107) 2 ≥ δ , one step of NGD update guarantees a decrease of γ 2 in the squared error, which means NGD must have found an δ -close estimate within (cid:107) τ 0 − τ ∗ (cid:107) 2 γsteps.
The full proof is deferred to Appendix C.1.
5. E XAMPLE : 1 D G AUSSIAN MEAN ESTIMATION
It is relatively straightforward to check that NGD addresses the ﬂatness problem faced by Gaussian
mean estimation we considered in Section 4:
Corollary 5.1.
Let P ∗ , Q be 1d Gaussian with covariance 1 and mean θ ∗ =
R where R (cid:28) 1 , and θ q = 0 .
For any given δ ≤ 1 R and initial estimate τ 0 = τ q , NGD can ﬁnd an estimate τ such that (cid:107) τ − τ ∗ (cid:107) 2 ≤ δ , with at most O
( R 6 δ 2 ) steps.
Intuitively, the effectiveness of NGD comes from the crucial observation that though the magnitude for the loss and derivatives can be exponentially small, they share the same exponential factor, making normalization effective.
Formally, it can be shown that β u β l
=
O (1) (Appendix D.6).
Corollary 5.1 then follows from Theorem 5.1 and the curvature and strong convexity from Lemma 4.1, 4.2.
B OUNDS ON THE CONDITION NUMBER OF NCE
The convergence rate in Theorem 5.1 depends on κ ∗ , the condition number of the NCE Hessian at the optimum, and Hessian-related constants β u , β l in Assumption 5.1.
We now show that under the setup of Theorem 5.1, κ ∗ and β u , β l can be related to the Bhattacharyya coefﬁcient between P ∗ and
Q , which is a similarity measure deﬁned as BC ( P ∗ , Q ) :
= (cid:82) x (cid:112) p ∗ ( x ) q ( x ) dx .
As a result, we get the following convergence guarantee:
Theorem 5.2.
Suppose Assumptions 2.1, 2.3 hold with constants ω , λ max , and λ min .
Consider a NCE task with data distribution P 1 and noise distribution P 2 , parameterized by θ 1 , θ 2 ∈ Θ respectively.
Then for any given δ ≤ 1 R and initial estimate τ 0 = τ q , NGD ﬁnds an estimate τ such that (cid:107) τ − τ ∗ (cid:107) 2 ≤ δ within T ≤ C ·
[Equation]
steps, where C := 18 exp ( 2 β Z ) · ( λ max λ min ) 3 ·
[Equation]
In particular, when P ∗ , Q are not too far, we can further show a lower bound on BC ( P ∗ , Q ) :
Lemma 5.2.
For P 1 , P 2 parameterized by θ 1 , θ 2 ∈ Θ , if (cid:107) θ 1 − θ 2 (cid:107) 22 ≤ 4 λ max , then BC ( P 1 , P 2 ) ≥ 12 .
The proofs of Theorem 5.2 and Lemma 5.2 rely on analyzing the geodesic on the manifold of square root densities √ p equipped with the Hellinger distance as a metric; the details are deferred to
Appendix C.3 and C.4.
It is also worth noting that Theorem 5.2 only requires (cid:107) θ 1 − θ 2 (cid:107) to be smaller than a constant, rather than tending to zero as usually required for analyses using Taylor expansions.
Finally, we would like to note that although our analysis can be tightened, it is unlikely to remove such dependency since NGD only uses ﬁrst-order information.
Moreover, the condition number κ ∗ also affects the practical use of Newton-like methods, since matrix inversion is widely known to be sensitive to numerical issues when the matrix is extremely ill-conditioned.
It is an interesting open question whether a non-standard preconditioning approach might be amenable to this setting.
A NALYZING E NCE : NCE WITH AN EXPONENTIAL LOSS
The previous section proved that NGD can serve as a simple ﬁx to overcome the ﬂatness problem of
NCE for well-conditioned losses.
However, though we showed κ ∗ has a polynomial growth when the distributions P , Q ∗ are sufﬁciently close —it is unclear how κ ∗ behaves beyond this threshold.
In this section, we introduce a slight modiﬁcation to the NCE objective, which we call the eNCE objective, for which κ ∗ depends polynomially on some class-related constants.
This means though eNCE may still suffer from the ﬂatness problem, eNCE and NGD together provide a solution that guarantees a polynomial convergence rate .
Towards formalizing this, the eNCE loss is deﬁned as:
Deﬁnition 6.1 (eNCE Loss) .
Let ϕ ( x ) :
= log (cid:113) p ( x ) q ( x ) , and l ( x, y ) :
= exp( − yϕ ( x )) for y ∈ {± 1 } .
The eNCE loss of P θ w.r.t.
data distribution P ∗ and noise Q is:
[Equation]
It can be checked easily that the minimizing ϕ learns ϕ ( x ) = 12 log p ∗ q .
Moreover, each ϕ is associated with an induced distribution p , deﬁned by p ( x ) = exp( ϕ ( x ))
q ( x ) .
Relation to NCE :
Same as the original NCE loss (referred to as “NCE” below), eNCE learns to solve a distinguishing task between samples from P ∗ or Q .
The difference lies only in the losses, which have analogous forms: the NCE loss described in Def. 2.1 can be rewritten in the same form with l ( x, y ) :
= log [Equation] and ψ ( x ) := log p ( x ) q ( x ) .
The main advantage of the exponential loss is that the Hessian at the optimum is now guaranteed to be well-conditioned:
Lemma 6.1 (Polynomial condition number for eNCE loss) .
Under Assumption 2.3 with constants λ max , λ min , the condition number of the eNCE Hessian at the optimum is bounded by κ ∗ ≤ λ max λ min .
We can also show that eNCE satisﬁes part (ii) of Assumption 5.1.
Due to space considerations, the proof is deferred to Appendix B.1.
Lemma 6.2.
Under assumption 2.2 with constant β Z , for any unit vector u and constant c ∈ [0 , 1 β Z ] , the maximum and minimum singular values of H ( τ ∗ + c u ) satisfy assumption 5.1 with constants β [Equation]
= [Equation]
Lemma 6.1 and Lemma 6.2 together imply the Hessian is well-conditioned around the optimum.
Combined with Theorem 5.1, we have the main result of this section:
Theorem 6.1.
Let P ∗ , Q be exponential family distributions with parameters τ ∗ , τ q under Assumption 2.1-2.4.
Let λ max , λ min be constants for Assumption 2.3.
For any given δ ≤ 1 β Z and parameter initialization τ 0 , performing NGD on the eNCE objective guarantees that when taking
[Equation] steps, there exists an iterate t ≤ T such that (cid:107) τ t − τ ∗ (cid:107) 2 ≤ δ .
Proof.
Theorem 6.1 follows directly from Theorem 5.1, using the condition number bound from
Lemma 6.1 and constants from 6.2.
6. P ROOF OF L EMMA 6.
It can be checked that the Hessian at P
= P ∗ is H ∗ := ∇ [Equation] p )
[Equation]
Recall that θ ∗ , θ q , ˜ T denote the parameters and sufﬁcient statistics without the partition func- tion coordinate, and τ ∗ , τ q , T denote the extended version with the partition function, e.g. τ ∗ =
[ θ ∗ , log Z ( θ ∗ )] , T ( x ) =
[ ˜ T ( x ) , − 1] .
Then, we can rewrite H ∗ as:
[Equation]
The Lemma then follows from λ min I (cid:22) E θ ∗ + θ {+q 
+}
[ TT (cid:62) ] (cid:22) λ max I by Assumption 2.3.
E MPIRICAL VERIFICATION
To corroborate our theory, we verify the effectiveness of NGD and eNCE on Gaussian mean estimation and the MNIST dataset.
For MNIST, we use a ResNet-18 to model the log density ratio log( p/q ) , following the setup in TRE (Rhodes et al., 2020).
Results: For Gaussian data, we run gradient descent (GD) and normalized gradient descent (NGD) on the NCE loss and eNCE loss.
Figure 1 compares the best runs under each setup given a ﬁxed computation budget (100 update steps), where “best" is deﬁned to be the run with the lowest loss on fresh samples.
The plots show the minimum parameter distance min t ∈
[ T ] (cid:107) τ ∗ − τ t (cid:107) 2 for each step T .
We ﬁnd that NGD indeed outperforms GD, and that the proposed eNCE sees a further improvement over NCE while additionally enjoying provable polynomial convergence guarantees.
For MNIST, we can no longer compare parameter distances since τ ∗ is unknown.
Instead, we compare the result of optimization directly in terms of loss achieved, again under a ﬁxed computation budget (2K steps).
The results are shown in Figure 2, with NGD converging signiﬁcantly faster for both NCE and eNCE.
We note that eNCE can be numerically unstable, especially when P ∗ , Q are well separated.
Implementation details to prevent numerical issues are included in Appendix E.
C ONCLUSION AND D ISCUSSIONS
We provided a theoretical analysis of the algorithmic difﬁculties that arise when optimizing the NCE objective with an uninformative noise distribution, stemming from an ill-behaved loss landscape.
Our theoretical results are inspired by empirical observations in prior works (Rhodes et al., 2020;
Gao et al., 2020; Goodfellow et al., 2014) and provide the ﬁrst formal explanation on the nature of the optimization problems of NCE.
Our negative results showed that even on the simple task of
Gaussian mean estimation, and even assuming access to the population gradient, gradient descent and Newton’s method with standard step size choice still require an exponential number of steps to reach a good solution.
We then proposed modiﬁcations to the NCE loss and optimization algorithm, whose combination results in the ﬁrst provably polynomial convergence rate for NCE.
The loss we propose, eNCE, can be efﬁciently optimized using normalized gradient descent and empirically outperforms existing methods.
We hope these theoretical results will help identify promising new directions in the search for simple, effective, and practical improvements to noise-contrastive estimation.