Title: Provable benefits of annealing for estimating normalizing constants: Importance Sampling, Noise-Contrastive Estimation, and beyond
Venue: NeurIPS 2023
Accepted: yes
Keywords: noise-contrastive estimation

annealed importance sampling
Abstract: Recent research has developed several Monte Carlo methods for estimating the normalization constant (partition function) based on the idea of annealing. This means sampling successively from a path of distributions which interpolate between a tractable "proposal" distribution and the unnormalized "target" distribution. Prominent estimators in this family include annealed importance sampling and annealed noise-contrastive estimation (NCE). Such methods hinge on a number of design choices: which estimator to use, which path of distributions to use and whether to use a path at all; so far, there is no definitive theory on which choices are efficient. Here, we evaluate each design choice by the asymptotic estimation error it produces. First, we show that using NCE is more efficient than the importance sampling estimator, but in the limit of infinitesimal path steps, the difference vanishes. Second, we find that using the geometric path brings down the estimation error from an exponential to a polynomial function of the parameter distance between the target and proposal distributions. Third, we find that the arithmetic path, while rarely used, can offer optimality properties over the universally-used geometric path. In fact, in a particular limit, the optimal path is arithmetic. Based on this theory, we finally propose a two-step estimator to approximate the optimal path in an efficient way.

Provable benefits of annealing for estimating normalizing constants: Importance Sampling, Noise-Contrastive Estimation, and beyond

Abstract

Recent research has developed several Monte Carlo methods for estimating the normalization constant (partition function) based on the idea of annealing. This means sampling successively from a path of distributions that interpolate between a tractable "proposal" distribution and the unnormalized "target" distribution. Prominent estimators in this family include annealed importance sampling and annealed noise-contrastive estimation (NCE). Such methods hinge on a number of design choices: which estimator to use, which path of distributions to use and whether to use a path at all; so far, there is no definitive theory on which choices are efficient. Here, we evaluate each design choice by the asymptotic estimation error it produces. First, we show that using NCE is more efficient than the importance sampling estimator, but in the limit of infinitesimal path steps, the difference vanishes. Second, we find that using the geometric path brings down the estimation error from an exponential to a polynomial function of the parameter distance between the target and proposal distributions. Third, we find that the arithmetic path, while rarely used, can offer optimality properties over the universally-used geometric path. In fact, in a particular limit, the optimal path is arithmetic. Based on this theory, we finally propose a two-step estimator to approximate the optimal path in an efficient way.

1 Introduction

Recent progress in generative modeling has sparked renewed interest in models of data that are defined by an unnormalized distribution. A prominent example is energy-based models, which are increasingly used in deep learning [1], and for which there are a variety of parameter estimation procedures [2-5]. Another example comes from Bayesian statistics, where the posterior model of parameters given data is frequently known only up to a proportionality constant. Such models can be evaluated and compared by the probability they assign to a dataset, yet this requires computing their normalization constants (partition functions) which are typically high-dimensional, intractable integrals.

Monte-Carlo techniques have been successful at computing these integrals using sampling methods [6]. The most common is importance sampling [6] which draws a sample from a tractable, "proposal" distribution to integrate the unnormalized "target" density. Noise-contrastive estimation (NCE) [3] uses a sample from both the proposal and the target, to compute the integral. Yet such methods suffer from high variance, especially when the "gap" between the proposal and target densities is large [7-9]. This has motivated various approaches to gradually bridge the gap with intermediate distributions, which is loosely referred to as "annealing". Among them, annealed importance sampling (AIS) [10-12] is widely adopted: it has been used to compute the normalization constants of deep stochastic models [13, 14] or to motivate a lower-bound for learning objectives [15, 16]. To integrate the unnormalized "target" density, it draws a sample from an entire path of distributions between the proposal and the target. While annealed importance sampling has been shown to be effective empirically, its theoretical understanding remains limited [17, 18]: it is yet unclear when annealing is effective, for which annealing paths, and whether AIS is a statistically efficient way to do it.

In this paper, we define a family of annealed Bregman estimators (ABE) for the normalization constant. We show that it is general enough to recover many classical estimators as a special case, including importance sampling, noise-contrastive estimation, umbrella sampling [19], bridge sampling [20] and annealed importance sampling. We provide a statistical analysis of its hyperparameters such as the choice of paths, and show the following:

First, we establish that using NCE is more asymptotically statistically efficient - in the sense of how many samples from the intermediate distribution need to be generated - than the importance sampling estimator, but in the limit of infinitesimal path steps, the difference vanishes.

Second, we find that the near-universally used geometric path brings down the estimation error from an exponential to a polynomial function of the parameter distance between the target and proposal distributions.

Third, we find that using the recently introduced arithmetic path [21] is exponentially inefficient in its basic form, yet it can be reparameterized to be in some sense optimal. Based on this optimality result, we finally propose a two-stage estimation procedure which first finds an approximation of the optimal (arithmetic) path, then uses it to estimate the normalization constant.

2 Background

Importance sampling and NCE The problem considered here is computing the normalization constant 1 , i.e. the integral of some unnormalized density f 1 ( x ) called "target".

Importance sampling and noise-contrastive estimation are two common estimators for that integral which use a random sample drawn from a tractable density p 0 ( x ) called "proposal"(Table 1, column 3). In fact, they are part of a larger family of Monte-Carlo estimators which can be interpreted as solving a binary classification task, aiming to distinguish between a sample drawn from the proposal and another from the target [22]. These estimators are summarized in Table 1. Each estimator is obtained by minimizing a specific binary classification loss that is identified by a convex function ϕ ( x ) . For example, minimizing the classification loss identified by ϕ IS ( x ) = x log x yields the importance sampling estimator. Similarly, ϕ RevIS ( x ) = -log x leads to the reverse importance sampling estimator [23], and ϕ NCE ( x ) = x log x -(1 + x ) log((1 + x ) / 2) to the noise-contrastive estimator.

Annealed estimators Annealing extends the above "binary" setup, by introducing a sequence of K +1 distributions from the proposal to the target (included). The idea will be to draw a sample from all these distributions to estimate the integral of the target f 1 ( x ) .

These intermediate distributions are obtained from a path ( f t ) t ∈ [0 , 1] , defined by interpolating between the proposal p 0 and unnormalized target f 1 : this path is therefore unnormalized. Different interpolation schemes can be chosen. A general one, explained in Figure 1, is to take the q -mean of the

General (

q

∈

]0

,

1]

)

f

t

(

x

) =

(

(1

-

t

)

p

0

(

x

)

q

+

tf

1

(

x

)

q

)

1

q

Geometric (

q

→

0

)

f

t

(

x

) =

p

0

(

x

)

1

-

t

×

f

1

(

x

)

t

Arithmetic (

q

= 1

)

f

t

(

x

) = (1

-

t

)

p

0

(

x

) +

tf

1

(

s

)

q

0

q

=1

p

0

p

1

Figure 1: q -mean paths between the proposal and target distributions. The geometric and arithmetic paths are obtained as limit cases. Here, the proposal (red) is a standard Gaussian. The target (blue) is a Gaussian mixture with two modes, and same first and second moments as the proposal. The path of distributions interpolates between blue and red.

proposal and target [21]. Two values of q are of particular interest: q → 0 defines a a near-universal path [18], obtained by taking the geometric mean of the target and proposal, while q = 1 defines a path obtained by the arithmetic mean.

Once a path is chosen, it can be uniformly 2 discretized into a sequence of K + 1 unnormalized densities, denoted by ( f k/K ) k ∈ [0 ,K ] with corresponding normalizations ( Z k/K ) k ∈ [0 ,K ] . In practice, samples are drawn from the corresponding normalized densities ( p k/K ) k ∈ [0 ,K ] using Markov Chain Monte Carlo (MCMC). This sampling step incurs a computational cost, which is paid in the hope of reducing the variance of the estimation. It is common in the literature [17, 18] to assume perfect sampling , meaning the MCMC has converged and produced exact and independent samples from the distributions along the path, which simplifies the analysis.

Estimation error A measure of "quality" is required to compare different estimation choices, such as whether to anneal and which path to use. Such a measure is given by the Mean Squared Error (MSE), which is generally tractable when written at the first order in the asymptotic limit of a large sample size [24, Eq. 5.20]. These expressions have been derived for estimators obtained by minimimizing a classification loss [22] and are included in Table 1. They measure the "gap" between the proposal and target distributions using statistical divergences. Note also that the estimation error depends on the normalized target density (column 4), while the estimators are computed using the unnormalized target density (column 3). Further details are available in Appendix A.

3 Annealed Bregman Estimators of the normalization constant

The question that we try to answer in this paper is: How should we choose the K distributions in annealing (the target is fixed), and how are their samples best used? To answer this, we will study the error produced by different estimation choices. But first we define the set of estimators for which the analysis is done.

Definition of Annealed Bregman Estimators We now define a new family of estimators, which we call annealed Bregman estimators (ABE) ; the motivation for this terminology will become clear in the coming paragraphs. We will show that this is a general class of estimators for computing the normalization using a sample drawn from the sequence of K +1 distributions. For ABE, the log

normalization log Z 1 is estimated additively along the sequence of distributions

Defining the estimation in log-space is analytically convenient, as it is easier to analyze a sum of estimators than a product. Exponentiating the result leads to an estimator of Z 1 . We naturally extend the binary setup ( K = 1 ) of Chehab et al. [22] and propose to compute each of the intermediate log-ratios, by solving a classification task between samples drawn from their corresponding densities p k/K and p ( k +1) /K . Each binary classification loss is now identified by a convex function ϕ k ( x ) and defined as

where the regression function r k ( x ; β k ) is parameterized by the unknown log-ratio β k

For the true β ∗ k , it holds β ∗ k = log( Z ( k +1) /K /Z k/K ) and r k ( x ; β ∗ k ) = p ( k +1) /K ( x ) /p k/K ( x ) . The convex functions ( ϕ k ) k ∈ [0 ,K -1] which identify the classification losses are called "Bregman" generators, hence ABE. As mentioned above, we assume perfect sampling and allocate the total sample size N equally among the K estimators in the sum.

Hyperparameters The annealed Bregman estimator depends on the following hyperparameters: (1) the choice of path q ; (2) the number of distributions along that path K +1 (including the proposal and the target); (3) the classification losses identified by the convex functions ( ϕ k ) k ∈ [0 ,K -1] .

Different combinations of these hyperparameters recover several common estimators of the logpartition function. In binary case of K = 1 this includes importance sampling, reverse importance sampling, and noise-contrastive estimation, each obtained for a different choice of the classification loss [22]. To build intuition, consider K = 2 so that we add a single intermediate distribution p 1 / 2 to the sequence. Using the importance sampling loss ( ϕ 0 = x log x ) for the first ratio, and reverse importance sampling ( ϕ 1 = -log x ) for the second ratio, recovers the bridge sampling estimator as a special case [20]

Alternatively, we can use these classification losses in reverse order: reverse importance sampling ( ϕ 0 = -log x ) for the first ratio, and importance sampling ( ϕ 1 = x log x ) for the second ratio, and recover the umbrella sampling estimator [19] also known as the ratio sampling estimator [25]

Another option yet, is to use the same classification loss for all ratios. With importance sampling ( ϕ k = x log x, ∀ k ∈ J 0 , K -1 K ), we recover the annealed importance sampling estimator [10-12]

The family of annealed Bregman estimators is visibly large enough to include many existing estimators, obtained for different hyperparmeter choices. This raises the fundamental question of how these hyperparameters should be chosen, in particular in the challenging case where the target and proposal have little overlap and the data is high dimensional . To answer this question, we will study the estimation error produced by different hyperparameter choices.

4 Statistical analysis of the hyperparameters

We consider a fixed data budget N and investigate how the remaining hyperparameters are best chosen for statistical efficiency. The starting point for the analysis is that as ABE estimates the

normalization in log-space, the estimator is obtained by a sum of independent and asymptotically unbiased estimators [26] given in Eq. 1 and thus the mean squared errors written in table 1 are additive. (Recall, the independence of these estimators is because new samples are drawn for each estimation task.) Each individual error actually measures an overlap between two consecutive distributions along the path, and annealing integrates these overlaps.

4.1 Classification losses, ϕ k

Given the popularity of annealed importance sampling, we should first ask if the importance sampling loss is really an acceptable default. We recall an important limitation of importance sampling: its estimation error is notoriously sensitive to distribution tails [27]. Without annealing, it is infinite when the target p 1 has a heavier tail than the proposal p 0 . When annealing with a geometric path, for example between two Gaussians with different covariances p 0 = N ( 0 , Id ) and p 1 = N ( 0 , 2 Id ) , the geometric path produces Gaussians with increasing variances Σ t = (1 -t/ 2) -1 Id and therefore increasing tails. Hence, the same tail mismatch holds along the path. Note that this concern is a realistic one for natural image data, as the target distribution over images is typically superGaussian [28] while the proposal is usually chosen as Gaussian.

This warrants a better choice for the loss: In the binary setup ( K = 1 ), the NCE loss is optimal [20, 22] and its error can be orders of magnitude less than importance sampling [20]. This optimality result has been extended to a sequence of distributions K > 1 [29, eq. 16]. We further show that in the limit of a continuous path, the gap between annealed IS and annealed NCE is closed and we provide their estimation error:

Theorem 1 (Estimation error and the Fisher-Rao path length) For a finite value of K , the optimal loss is NCE:

In the limit of K →∞ (such that K 2 /N → 0 ), NCE, IS, and revIS converge to the same estimation error, given by the Fisher-Rao path length from the proposal to the target:

where the Fisher-Rao metric I ( t ) := E x ∼ p ( x ,t ) [( d dt log p t ( x )) 2 ] is defined as the Fisher information over the path, using time t as the parameter.

This is proven in Appendix B. Note that in this theorem, to take limits successively in N → ∞ then in K →∞ , we assume that N grows at least as fast as K 2 . While the NCE estimator requires solving a (potentially non-convex) scalar optimization problem in Eq. 2 and IS does not, this is the price to pay for statistical optimality. In the following, we will keep the optimal NCE loss and will indicate the dependency of the estimation error on ϕ NCE with a subscript, instead. We highlight that our theorems in this paper apply to the MSE in the limit of K → ∞ : their results hold the same for the IS and RevIS losses by virtue of theorem 1. Just as in the binary case, while the estimator is computed with the unnormalized path of densities (Eq. 2), the estimation error depends on the normalized path of densities (Eq. 8).

4.2 Number of distributions, K +1

It is known that estimating the normalization constant using plain importance sampling ( K = 1 ) can produce a statistical error that is exponential in the distance between the target and the proposal [30]. We show that in the binary case, NCE also suffers from an estimation error that scales exponentially with the parameter-distance between the target and proposal dimension.

In the following, we consider a proposal p 0 and target p 1 that are in an exponential family with sufficient statistics t ( x ) . Note that exponential families have a rich history in statistical literature: they are classic parametric families of distributions [31] and certain of them have universal approximation capabilities [32, 33]. An exponential family is defined as

where Z ( θ ) = ∫ exp( ⟨ θ 1 , t ( x ) ⟩ ) is the partition function. We will consider that the (unnormalized) target density f 1 is what we call a simply unnormalized model defined as

Note that in general, a pdf can be unnormalized in many ways: one can multiply an unnormalized density by any positive function of θ and it will still be unnormalized. However, the simple and intuitive case defined above is what we base the analysis below on.

For exponential families, the log-normalization log Z ( θ ) is a convex function ("log-sum-exp") of the parameter θ [34], which implies 0 ≼ ∇ 2 θ log Z ( θ ) . In our theorems, we use the further assumptions of strong convexity with constant M , and/or smoothness with constant L (gradient is L -Liptschitz):

For exponential families, the derivatives of the log partition function yield moments of the sufficient statistics, and the hessian ∇ 2 θ log Z ( θ ) = Cov x ∼ p [ t ( x )] is in fact the Fisher matrix. We can interpret our two assumptions: Eq. 11 can be seen as a form of "strong identifiability". Namely, positivedefiniteness is required of the Fisher matrix, for the Maximum-Likelihood loss to have a unique minimum: we further assume a lower-bound on the smallest eigenvalue, which can be viewed as a strong identifiability condition. Eq. 12 can be interpreted as a bound on the second-order moments of the distribution p ( x ; θ ) , which is equivalent to the variance in every direction being bounded, which will be the case for parameters in a bounded domain θ ∈ Θ . An example along with the proofs of the following Theorems 2 and 3, are provided in Appendix B.

Theorem 2 (Exponential error of binary NCE) Assume the proposal p 0 is from the normalized exponential family, while the (unnormalized) target f 1 is from the simply unnormalized exponential family (Eq. 10). The log-partition function log Z ( θ ) is assumed to be strongly convex (Eq. 11).

Then in the binary case K = 1 , the estimation error of NCE is (at least) exponential in the parameterdistance between the proposal and the target:

where M is the strong convexity constant of log Z ( θ ) .

Annealing the importance sampling estimator (increasing K ) was proposed in the hope that we can trade the statistical cost in the dimension for a computational cost (number of classification tasks) which could be considered more acceptable. Yet, there is no definitive theory on the ability of annealing to reduce the statistical cost in a general setup [18, 21]. For both importance sampling and noise-contrastive estimation, we next prove that annealing with the near-universally used geometric path brings down the estimation error, from exponential to polynomial in the parameter-distance between the proposal and target. Given that we expect ∥ θ 1 -θ 0 ∥ 2 to scale as √ D with the dimension, using these paths effectively makes annealed estimation amenable to high-dimensional problems. This corroborates empirical [35] and theoretical [17] results which suggested in simple cases that annealing with an appropriate path can reduce the estimator error up to several orders of magnitude.

Theorem 3 (Polynomial error of annealed NCE with a geometric path) Assume the proposal p 0 is from the normalized exponential family, while the (unnormalized) target f 1 is from the simply unnormalized exponential family (Eq. 10). The log-partition function log Z ( θ ) is assumed to be strongly convex and smooth (Eq. 11, Eq. 12).

Then in the annealing limit of a continuous path K →∞ , the estimation error of annealed NCE with the geometric path is (at most) polynomial in the parameter-distance between the proposal and the target:

where M and L are respectively the strong convexity and smoothness constants of log Z ( θ ) .

To our knowledge, this is the first result building on Gelman and Meng [17, Table 1] and Grosse et al. [18] which showcases the benefits of annealed estimation for a general target distribution.

We conclude that annealing with the near-universally used geometric path provably benefits noisecontrastive estimation, as well as importance sampling and reverse importance sampling, when the proposal and target distributions have little overlap.

4.3 Path parameter, q -geometric vs. arithmetic

Despite the near-universal popularity of the geometric path ( q → 0 ), it is worth asking if there are other simple paths that are more optimal. Interpolating moments of exponential families was shown to outperform the geometric path by Grosse et al. [18], yet building such a path requires knowing the exponential family of the target. Other paths based on the arithmetic mean (and generalizations) of the target and proposal, were proposed in Masrani et al. [21], without a definitive theory of the estimation error.

Next, we analyze the error of the arithmetic path. We prove that the arithmetic path ( q = 1 ) does not exhibit the same benefits as the geometric path: in general, its estimation error grows exponentially in the parameter-distance between the target and proposal distributions. However, in the case where an oracle gives us the normalization Z 1 to be used only in the construction of the path (we will discuss what this means in practice below), the arithmetic path can be reparameterized so as to bring down the estimation error to polynomial, even constant, in the parameter-distance. We start by the negative result.

Theorem 4 (Exponential error of annealed NCE with an arithmetic path) Assume the proposal p 0 is from the normalized exponential family, while the (unnormalized) target f 1 is from the simply unnormalized exponential family (Eq. 10). The log-partition function log Z ( θ ) is assumed to be strongly convex (Eq. 11).

Consider the annealing limit of a continuous path K →∞ path and a far-away target with large enough ∥ θ 1 -θ 0 ∥ > 0 . For estimating the log normalization of the (unnormalized) target density f 1 , the estimation error of annealed NCE with the arithmetic path is (at least) exponential in the parameter-distance between the proposal and the target:

where C is constant defined in Appendix B.3.

We suggest an intuitive explanation for this negative result. We begin with the observation that the estimation error (Eq. 8) depends on the normalized path of densities. Suppose the target model is rescaled by a constant 100 , so that the new unnormalized target density is f 1 ( x ) × 100 and its new normalization is Z 1 × 100 . Looking at table 2, this rescaling does not modify the geometric path of normalized densities, while it does the arithmetic path of normalized densities. Because the estimation error depends on path of normalized densities, this makes the arithmetic choice sensitive to target normalization, even more so as the parameter distance grows and the log-normalization with it, as a strongly convex function of it (Appendix, Eq. 94). This suggests making the arithmetic path of normalized distributions "robust" to the choice of Z 1 . We will show this can be achieved by re-parameterizing the path in terms of Z 1 .

We next prove that certain reparameterizations can bring down the error to a polynomial and even constant function of the parameter-distance between the target and proposal. The following theorems may seem purely theoretical, as if necessitating an oracle for Z 1 , but they will actually lead to an efficient estimation algorithm later.

Theorem 5 (Polynomial error of annealed NCE with an arithmetic path and oracle) Assume the same as in Theorem 4, replacing the strong convexity of the log-partition by smoothness (Eq. 12). Additionally, suppose an oracle gives the normalization constant Z 1 to be used only in the reparameterization of the arithmetic path with t → t t + Z 1 (1 -t ) (see Table 2). This brings down the estimation error of annealed NCE to (at most) polynomial in the parameter-distance:

where L is the smoothness constant of log Z ( θ ) .

In fact, supposing we have (oracle) access to the normalizing constant Z 1 , the arithmetic path can even be reparameterized such that it is the optimal path in a certain limit. We next prove such optimality in the limits of a continuous path K →∞ and "far-away" target and proposal [17]:

Theorem 6 (Constant error of annealed NCE with an arithmetic path and oracle) Consider the limits of a continuous annealing path K →∞ , and of a target distribution whose overlap with the proposal

goes to zero. Namely, consider the quantities:

Assume sup x ∈ R D ϵ ′ ( x ) → 0 , sup x ∈ R D ϵ ′′ ( x ) → 0 , ϵ → 0 , ϵ ′′′ → 0 , and consider the distributions p arith -trig t and p oracle t as defined in Table 2.

Then, the optimal annealing path convergences pointwise to an arithmetic path reparameterized trigonometrically with t → t sin 2 ( πt 2 )+ Z 1 (1 -sin 2 ( πt 2 )) and the estimation error tends to the optimal estimation error (which is constant with respect to the parameter-distance):

By way of remarks, we note that the assumptions that the quantities sup ϵ ′ ( x ) , sup ϵ ′′ ( x ) , ϵ, ϵ ′′′ go to 0 are mutually incomparable (i.e. none of them implies the others). For example, sup x ∈ R D ϵ ′ ( x ) and ϵ going 0 require that the function √ p 0 ( x ) p 1 ( x ) goes to 0 in the L ∞ and L 1 sense, respectively and these two norms are not equivalent in the Lebesgue measure.

Two-step estimation Thus, we see that, perhaps unsurprisingly, the "optimal" mixture weights in the space of unnormalized densities depends on the true Z 1 : however, this dependency is simple. We propose a two-step estimation method: first, Z 1 is pre-estimated, for example using the geometric path; second, the estimate of Z 1 is plugged into the "oracle" or "oracle-trig" weight of the arithmetic path (table 2, column 2), and which is used to obtain a second estimation of Z 1 . Note that preestimating a problematic (hyper)parameter, here Z 1 , has proved beneficial to reduce the estimation error of NCE in a related context [36].

5 Numerical results

We now present numerical evidence for our theory and validate our two-step estimators. Importantly, we do not claim to achieve state of the art in terms of practical evaluation of the normalization constants; our goal is to support our theoretical analysis. We follow the evaluation methods of importance sampling literature [18] and evaluate our methods on synthetic Gaussians. This setup is specially convenient for validating our theory: the optimal estimation error can conveniently be computed in closed-form, so too can the geometric and arithmetic paths which avoids a sampling error from MCMC algorithms. These derivations are included in the Appendix B. We specifically consider the high-dimensional setting, where the computation of the determinant of a high-dimensional (covariance) matrix which appears in the normalization of a Gaussian, can in fact be challenging [37].

Numerical Methods The proposal distribution is always a standard Gaussian, while the target differs by the second moment: p 1 = N ( 0 , 2 Id ) in Figure 2, p 1 = N ( 0 , 0 . 25 Id ) in Figure 3b and p 1 = N ( 0 , σ 2 Id ) in Figure 3a, where the target variance decreases as σ ( i ) = i -1 so that the (natural) parameter distance grows linearly [34, Part II-4]. We use a sample size of N = 50000 points, and, unless otherwise mentioned, K +1 = 10 distributions from the annealing paths and the dimensionality is 50 . To compute an estimator of the normalization constant using the non-convex NCE loss, we used a non-linear conjugate gradient scheme implemented in Scipy [38]. We chose the conjugate-gradient algorithm as it is deterministic (no residual variance like in SGD). The empirical Mean-Squared Error was computed over 100 random seeds, parallelized over 100 CPUs. The longest experiment was for Figure 3b and took 7 wall-clock hours to complete. For the two-step estimators ("two-step" and "two-step (trig)"), a pre-estimate of the normalization was first computed using the geometric path with 10 distributions. Then, this estimate was used to to re-parameterize an arithmetic path with 10 distributions which produced the second estimate.

Results Figure 2 numerically supports the optimality of the NCE loss for a finite K (here, K = 2 so three distributions are used) proven in Theorem 1; note that the proposal distribution was normalized for this experiment, so that the log normalizing constant is zero. Figure 3 validates our main results for annealing paths. It shows how the estimation error scales with the proposal and target distributions growing apart, either with the parameter-distance in Figure 3a or with the dimensionality in Figure 3b. Using no annealing path ( K = 1 ) produces an estimation error which grows linearly in log space; this numerically supports the exponential growth predicted by Theorem 2. Meanwhile, annealing ( K →∞ ) sets the estimation error on different trends, depending on the choice of path. Choosing the geometric path brings the growth down to sub-exponential, as predicted by Theorem 3, while choosing the (basic) arithmetic path does not as in Theorem 4. To alleviate this, our two-step estimation methods consist in reparameterizing the arithmetic path so that it actually does bring down the estimation error. In fact, our two-step estimators in table 2 empirically approach the optimal estimation error in Figure 3. While this requires more computation, it has the appeal of making the estimation error constant with respect to the parameter-distance between the target and proposal distributions. Practically, this means that in Figure 3a, regular Noise-Contrastive Estimation (black, full line) fails when the parameter-distance between the target and proposal distributions is higher than 20 , while our two-step estimators remain optimal.

We next explain interesting observations in Figure 3 which are actually coherent with our theory. First, in Figure 3a, the "two-step (trig)" estimator is only optimal when the parameter-distance between the target and proposal distributions is larger than 10 . This is because the optimality of this two-step estimator was derived in Theorem 6 conditionally on non-overlapping distributions, here achieved by a large parameter-distance. Second, in both Figures 3a and 3b, the "two-step" estimator empirically achieves the optimal estimation error that was predicted for the "two-step (trig)" estimator. This suggests our polynomial upper bound from Theorem 5 may be loose in certain cases. This further explains why, in Figure 3a, the arithmetic path is near-optimal for a single value of the parameter-distance. At this value of 20 , the partition function happens to be equal to one Z ( θ 1 ) = 1 , so that the arithmetic path is effectively the same as the "two-step" estimator.

0.2

0.0

0.2

Estimates

IS

RevIS

NCE

6 Related work and Limitations

Previous work has mainly focused on annealed importance sampling [18, 39], which is a special case of our annealed Bregman estimator. They have evaluated the merits of different paths empirically, using an approximation of the estimation error called Effective Sample Size (ESS) and the consistency-

0

10

20

30

Parameter distance

5.0

2.5

0.0

MSE (log10)

no path

optimal

geometric

arithmetic

two-step

two-step (trig)

(a) MSE against parameter distance.

0

50

100

Dimension

5.0

2.5

0.0

MSE (log10)

no path

optimal

geometric

arithmetic

two-step

two-step (trig)

(b) MSE against dimensionality.

gap. In our analysis, we consider consistent estimators and derive and optimize the exact estimation error of the optimal Noise-Contrastive Estimation. Liu et al. [27] considered the NCE estimate for Z (not log Z ) with the name "discriminance sampling", and annealed the estimator using an extended state-space construction similar to Neal [10]. Their analysis of the estimation error is relevant but does not deal with hyperparameters other than the classification loss. Other annealed estimators also include annealed NCE as a special case but with little theory on its estimation error [17, 40, 29].

We made the common assumption of perfect sampling [17, 18, 41, 42] in order to make the estimation error tractable and obtain practical guidelines to reduce it. We note however, that this leaves a gap to bridge with a practical setup where the sampling error cannot be ignored; in fact, annealed importance sampling [10] was originally proposed such that the samples can be obtained from a Markov Chain that has not converged. In this original formulation, AIS is a special case of a larger framework called Sequential Monte Carlo (SMC) [43] in which the path of distributions is implicitly defined (by Markov transitions), sometimes even "on the go" [42]. Yet even within that theory, it seems that analyzing the estimation error for an inexplicit path of distributions is challenging [44, Eq. 38]. In particular, samples from MCMC will typically follow marginal distributions that are not analytically tractable, thus the stronger assumption of "perfect sampling" is often used to make estimation error depend explicitly on the path of distributions [45, Eq. 3.2]. Even then, results on how the estimation error scales with dimensionality are heuristic [10] or limited by assumptions such as an essentially log-concave path of distributions or a factorial target distribution [45].

It might also be argued that the limit of almost no overlap between proposal and target, which we use a lot, is unrealistic. To see why it can be realistic, consider the case of natural image data. A typical proposal is Gaussian, since nothing much more sophisticated is tractable in high dimensions. However, there is almost no overlap between Gaussian data and natural images, which is seen in the fact that a human observer can effortlessly discriminate between the two.

More generally, many methods based on "annealing" were developed to deal with sampling issues. In fact, the path costs for two such methods, parallel tempering [46, eq. 17] and tempered transitions [47, eq. 18], can be written with a sum of f-divergences which in the limit of a continuous path is the same cost function as in our Theorem 1. This suggests our results may apply to these methods as well.

7 Conclusion

We defined a class of estimators of the normalization constant, annealed Bregman estimation (ABE), which relies on a sampling phase from a path of distributions, and an estimation phase where these samples are used to estimate the log-normalization of the target distribution. Our results suggest a number of simple recommendations regarding hyperparameter choices in annealing. First, if the path has very few intermediate distributions, it is better to choose NCE due to its statistical optimality (Theorem 1). If however, the path has many intermediate distributions and approaches the annealing limit, then IS enjoys the same statistical optimality as NCE but has the advantage of its computational simplicity. Annealing can always provide substantial benefits (Theorem 2). Moreover, if we have a reasonable a priori estimate of Z 1 , the arithmetic path achieves very low error (Theorem 5) sometimes even approaching optimality (Theorem 6). On the other hand, even absent an initial estimate of Z 1 , the geometric path can exponentially reduce the estimation error compared with no annealing (Theorems 2 and 3).

Acknowledgements This work was supported by the French ANR-20-CHIA-0016. Aapo Hyvärinen was supported by funding from the Academy of Finland and a Fellowship from CIFAR. Andrej Risteski was supported in part by NSF awards IIS-2211907, CCF-2238523, and an Amazon Research Award.

References

Y. Song, S. Sohl-Dickstein, D.P. Kingma, A. Kumar, S. Ermon, and B. Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations (ICLR) , 2021.

A. Hyvärinen. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research , 6(24):695-709, 2005.

M. Gutmann and A. Hyvärinen. Noise-contrastive estimation of unnormalized statistical models, with applications to natural image statistics. Journal of Machine Learning Research , 13(11): 307-361, 2012.

G.E. Hinton. Training products of experts by minimizing contrastive divergence. Neural Computation , 14(8):1771-1800, 2002.

R. Gao, E. Nijkamp, D.P. Kingma, Z. Xu, A.M. Dai, and Y. Nian Wu. Flow contrastive estimation of energy-based models. 2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) , pages 7515-7525, 2020.

A.B. Owen. Monte Carlo theory, methods and examples . https://artowen.su.domains/ mc/ , 2013.

B. Liu, E. Rosenfeld, P. Ravikumar, and A. Risteski. Analyzing and improving the optimization landscape of noise-contrastive estimation. In International Conference on Learning Representations (ICLR) , 2022.

H. Lee, C. Pabbaraju, A.P. Sevekari, and A. Risteski. Pitfalls of gaussians as a noise distribution in NCE. In International Conference on Learning Representations (ICLR) , 2023.

O. Chehab, A. Gramfort, and A. Hyvärinen. The optimal noise in noise-contrastive learning is not what you think. In Conference on Uncertainty in Artificial Intelligence (UAI) , volume 180, pages 307-316. PMLR, 2022.

R.M. Neal. Annealed importance sampling. Statistics and Computing , 11:125-139, 1998.

C. Jarzynski. Nonequilibrium equality for free energy differences. Physical Review Letters , 78: 2690-2693, 1996.

C. Jarzynski. Equilibrium free-energy differences from nonequilibrium measurements: A master-equation approach. Physical Review E , 56:5018-5035, 1997.

R. Salakhutdinov and G. Hinton. Replicated softmax: an undirected topic model. In Neural Information Processing Systems (NIPS) , 2009.

Y. Dauphin and Y. Bengio. Stochastic ratio matching of rbms for sparse high-dimensional inputs. In Neural Information Processing Systems (NIPS) , volume 26. Curran Associates, Inc., 2013.

V. Masrani, T.A. Le, and F.D. Wood. The thermodynamic variational objective. In Neural Information Processing Systems (NeurIPS) , 2019.

R. Brekelmans, S. Huang, M. Ghassemi, G. Ver Steeg, R.B. Grosse, and A. Makhzani. Improving mutual information estimation with annealed and energy-based bounds. In International Conference on Learning Representations (ICLR) , 2022.

A. Gelman and X.-L. Meng. Simulating normalizing constants: From importance sampling to bridge sampling to path sampling. Statistical Science , 13:163-185, 1998.

R. Grosse, C. Maddison, and R. Salakhutdinov. Annealing between distributions by averaging moments. In Advances in Neural Information Processing Systems (NIPS) , volume 26. Curran Associates, Inc., 2013.

G.M. Torrie and J.P. Valleau. Nonphysical sampling distributions in monte carlo free-energy estimation: Umbrella sampling. Journal of Computational Physics , 23(2):187-199, 1977.

X.-L. Meng and W.H. Wong. Simulating ratios of normalizing constants via a simple identity: a theoretical exploration. Statistica Sinica , pages 831-860, 1996.

V. Masrani, R. Brekelmans, T. Bui, F. Nielsen, A. Galstyan, G. Ver Steeg, and F. Wood. q-paths: Generalizing the geometric annealing path using power means. In Conference on Uncertainty in Artificial Intelligence (UAI) , volume 161, pages 1938-1947. PMLR, 27-30 Jul 2021.

O. Chehab, A. Gramfort, and A. Hyvärinen. Optimizing the Noise in Self-Supervised Learning: from Importance Sampling to Noise-Contrastive Estimation. ArXiv , 2023. doi: 10.48550/arXiv. 2301.09696.

M.A. Newton and A.E. Raftery. Approximate Bayesian Inference with the Weighted Likelihood Bootstrap. Journal of the Royal Statistical Society. Series B (Methodological) , 56(1):3-48, 1994.

A.W. van der Vaart. Asymptotic Statistics . Asymptotic Statistics. Cambridge University Press, 2000.

M.-H. Chen and Q.-M. Shao. On monte carlo methods for estimating ratios of normalizing constants. The Annals of Statistics , 25(4):1563-1594, 1997.

M. Uehara, T. Matsuda, and F. Komaki. Analysis of noise contrastive estimation from the perspective of asymptotic variance. ArXiv , 2018. doi: 10.48550/ARXIV.1808.07983.

Q. Liu, J. Peng, A.T. Ihler, and J.W. Fisher III. Estimating the partition function by discriminance sampling. In Uncertainty in Artificial Intelligence (UAI) , 2015.

A. Hyvärinen, J. Hurri, and P. O. Hoyer. Natural Image Statistics . Springer-Verlag, 2009.

O. Krause, A. Fischer, and C. Igel. Algorithms for estimating the partition function of restricted boltzmann machines. Artificial Intelligence , 278:103195, 2020.

S. Chatterjee and P. Diaconis. The sample size required in importance sampling. arXiv: Probability , 2015.

B. Efron. Exponential Families in Theory and Practice . Institute of Mathematical Statistics Textbooks. Cambridge University Press, 2022. ISBN 9781108488907.

B. Sriperumbudur, K. Fukumizu, A. Gretton, A. Hyvärinen, and R. Kumar. Density estimation in infinite dimensional exponential families. Journal of Machine Learning Research , 18(57): 1-59, 2017.

O. Krause, A. Fischer, T. Glasmachers, and C. Igel. Approximation properties of dbns with binary hidden units and real-valued visible units. In International Conference on Machine Learning (ICML) , 2013.

F. Nielsen and V. Garcia. Statistical exponential families: A digest with flash cards. ArXiv , 2011. doi: 10.48550/arXiv.0911.4863.

X.-L. Meng and S. Schilling. Fitting full-information item factor models and an empirical investigation of bridge sampling. Journal of the American Statistical Association , 91(435): 1254-1267, 1996.

M. Uehara, T. Kanamori, T. Takenouchi, and T. Matsuda. A unified statistically efficient estimation framework for unnormalized models. In International Conference on Artificial Intelligence and Statistics (AISTATS) , volume 108, pages 809-819. PMLR, 2020.

L. Ellam, H. Strathmann, M.A. Girolami, and I. Murray. A determinant-free method to simulate the parameters of large gaussian fields. Stat , 6:271-281, 2017.

P. Virtanen, R. Gommers, T.E. Oliphant, M. Haberland, T. Reddy, D. Cournapeau, E. Burovski, P. Peterson, W. Weckesser, J. Bright, S.J. van der Walt, M. Brett, J. Wilson, J.K. Millman, N. Mayorov, A.R.J. Nelson, E. Jones, R. Kern, E. Larson, C.J. Carey, I. Polat, Y. Feng, E.W. Moore, J. VanderPlas, D. Laxalde, J. Perktold, R. Cimrman, I. Henriksen, E.A. Quintero, C.R. Harris, A.M. Archibald, A.H. Ribeiro, F. Pedregosa, P. van Mulbregt, and SciPy 1.0 Contributors. SciPy 1.0: Fundamental Algorithms for Scientific Computing in Python. Nature Methods , 17: 261-272, 2020.

T. Kiwaki. Variational optimization of annealing schedules. ArXiv , 2015. doi: 10.48550/arXiv. 1502.05313.

R. M. Neal. Estimating ratios of normalizing constants using linked importance sampling. ArXiv , 2005. doi: 10.48550/arXiv.math/0511216.

R. Brekelmans, V. Masrani, F.D. Wood, G. Ver Steeg, and A.G. Galstyan. All in the exponential family: Bregman duality in thermodynamic variational inference. In International Conference on Machine Learning (ICML) , 2020.

S. Goshtasbpour, V. Cohen, and F. Perez-Cruz. Adaptive annealed importance sampling with constant rate progress. In International Conference on Machine Learning (ICML) , volume 202 of Proceedings of Machine Learning Research , pages 11642-11658. PMLR, 2023.

N. Chopin and O. Papaspiliopoulos. An Introduction to Sequential Monte Carlo . Springer Series in Statistics. Springer International Publishing, 2020. ISBN 9783030478452.

P. Del Moral and A. Doucet. Sequential monte carlo samplers. Journal of the Royal Statistical Society: Series B (Statistical Methodology) , 68, 2002.

C. Dai, J. Heng, P.E. Jacob, and N. Whiteley. An invitation to sequential monte carlo samplers. Journal of the American Statistical Association , 117:1587-1600, 2020.

S. Syed, V. Romaniello, T. Campbell, and A. Bouchard-Cote. Parallel tempering on optimized paths. In International Conference on Machine Learning (ICML) , volume 139 of Proceedings of Machine Learning Research , pages 10033-10042. PMLR, 2021.

G. Behrens, N. Friel, and M. Hurn. Tuning tempered transitions. Statistics and Computing , 22: 65-78, 2010.

M. Pihlaja, M. Gutmann, and A. Hyvärinen. A family of computationally efficient and simple estimators for unnormalized statistical models. In Uncertainty in Artificial Intelligence (UAI) , 2010.

L. Wang, D.E. Jones, and X.-L. Meng. Warp bridge sampling: The next generation. Journal of the American Statistical Association , 117(538):835-851, 2022.

H. Xing. Improving bridge estimators via f-GAN. Statistics and Computing , 32, 2022.

Y. Polyanskiy and Y. Wu. Information Theory: From Coding to Learning . Cambridge University Press, 2022.

S.-I. Amari and H. Nagaoka. Methods of information geometry . AMS, 2000.

Appendices

In the following, we will study the estimation error of annealed Bregman estimation (ABE) in two important setups: the log-normalization is computed using two distributions ( K = 1 ), the proposal and the target, or else using a path of distributions ( K →∞ ).

The code used for the experiments is available at https://github.com/l-omar-chehab/ annealing-normalizing-constants .

A No annealing, K = 1

We use [22, Eq.21] for the estimation error of any suitably parameterized 3 classifier F ( x ; β ) between two distributions p 1 and p 0 . The estimation error is measured by the asymptotic Mean-Squared Error (MSE)

which depends on the sample sizes N = N 1 + N 0 , their ratio ν = N 1 /N 0 , the Bregman classification loss indexed by the convex function ϕ ( x ) , and the asymptotic variance matrix

We suppose the standard technical conditions of van der Vaart [24, Th. 5.23] apply so that the remainder term ∥ ˆ β -β ∗ ∥ 2 can indeed be written independently of the parameterization, as o ( N -1 ) . Here, m w ( β ∗ ) , I w ( β ∗ ) and I v ( β ∗ ) are the reweighted mean and covariances of the parametegradient of the classifier, also known as the "relative" Fisher score ∇ β F ( x ; β ∗ ) ,

where the reweighting of data points is by w ( x ) := p 1 νp 0 ( x ) ϕ ′′ ( p 1 νp 0 ( x ) ) and by v ( x ) = w ( x ) 2 νp 0 ( x )+ p 1 ( x ) νp 0 ( x ) , which are all evaluated at the true parameter value β ∗ .

Scalar parameterization We now consider a specific parameterization of the classifier:

where the optimal parameter is the log-ratio of normalizations β ∗ = log( Z 1 /Z 0 ) . Consequently, we have ∇ β F ( x ; β ∗ ) = -1 and plugging this into the above quantities yields

which matches the formula found in [20, Eq 3.2]. For different choices of the Bregman classification loss, the estimation error is written using a divergence between the two distributions

where

Proof of Theorem 2 Exponential error of binary NCE

In the following, we will drop the remainder term in o ( N -1 ) given that no other limits are taken and that we will study the dominant term only. The estimation error of binary NCE is expressed in terms of the harmonic divergence

which is intractable for general exponential families. Instead, we can lower-bound the estimation error. To do so, we lower-bound the harmonic divergence using the inequality of means (harmonic vs. geometric)

and therefore

This lower bound is expressed in terms of the squared Hellinger distance, that is tractable for exponential families:

We now wish to lower bound MSE LB , and therefore D H 2 ( p 1 , p 0 ) , by an expression which is exponential in the parameter distance ∥ θ 1 -θ 0 ∥ . To do so, we note that for exponential families, the log-normalization is convex in the parameters. Here, we further assume strong convexity, so that

where M is the strong convexity constant. Plugging this back into the squared Hellinger distance, we obtain

so that the MSE

grows exponentially with the euclidean distance between the parameters.

B Annealing limit, K →∞

We now consider annealing paths ( p t ) t ∈ [0 , 1] that interpolate between between the proposal p 0 and the target p 1 .

B.1 Estimation error

We first show the optimality of the NCE loss within the family of Annealed Bregman Estimators, in the sense that it produces the smallest estimation error. We then study the estimation error of different annealed Bregman estimators in the annealing limit of a continuous path ( K →∞ ).

Proof of Theorem 1 Optimality of the NCE loss and the estimation error in the annealing limit K →∞

Optimality of the NCE loss Because the annealed Bregman estimator is built by adding independent estimators

the total Mean Squared Error (MSE) is the sum of each MSEs for each estimator (indexed by k ∈ J 0 , K -1 K )

where we highlighted the dependency on the classification losses identified by ( ϕ k ) k ∈ J 0 ,K K . The MSEs follow Eq. 26. It was shown by Meng and Wong [20] that for any of these MSEs, the optimal loss is identified by ϕ k ( x ) = x log x -(1 + x ) log( 1+ x 2 ) and is in fact the NCE loss [22]. Thus the sum of MSEs is minimized for the same loss.

Annealed Noise-Contrastive Estimation (NCE)

We are interested in the estimation error (asymptotic MSE) obtained for the NCE loss. Based off table 1, it is written as

where a sample budget of N/K is used for each estimator that is summed. The estimation error of balanced ( ν = 1 ) NCE-JS between two "close" distributions p t and p t + h , is

up to the remainder term. The estimation error can be simplified using a Taylor expansion. To do so, we recall that D HM is an f-divergence generated by ϕ ( x ) = 1 -x π +(1 -π ) x [49, 50] ( π = 1 2 here) and its expansion is therefore [51, Eq.7.64]

It follows that

Summing these estimation errors along the path of distributions with h = 1 /K ,

In the case of a parametric path p ( x | θ ( t )) t ∈ [0 , 1] , the proof is the same. Simply, the secondorder term in the Taylor expansion of Eq. 42 is computed using the chain rule

Annealed importance sampling (IS)

Similarly, for the choice of the importance sampling base estimator,

given that ϕ ( x ) = -log( x ) and therefore ϕ ′′ (1) = 1 for the reverse χ 2 divergence.

Annealed reverse importance sampling (RevIS)

Similarly, for the choice of the reverse importance sampling base estimator,

given that ϕ ( x ) = x log( x ) and therefore ϕ ′′ (1) = 1 for the χ 2 divergence.

B.2 Examples of paths

Geometric path The geometric path is defined in the space of unnormalized densities by

so in the space of normalized densities, the path is

where the normalization is

Arithmetic path The arithmetic path is defined in the space of unnormalized densities by

so in the space of normalized densities, the path is actually a mixture between the target and the proposal, where the weight of the mixture is a nonlinear function of the target normalization

Optimal path We know ( e.g. from Gelman and Meng [17, Eq. 49] that the optimal path is

where the coefficients a ( t ) and b ( t )

are simple functions of the squared Hellinger distance D H 2 ( p 0 , p 1 ) between the proposal and the target 4

The estimation error produced by that optimal path is [17, Eq. 48]

the squared Hellinger distance can be written in closed-form

and plugs into the optimal path formula, which is also obtained in closed-form.

For two Gaussians

B.3 Estimation error from taking different paths

Proof of Theorem 3 Polynomial error of annealed NCE with the geometric path

We next study the estimation error produced by the geometric path (Figure 1). In the annealing limit K →∞ , the MSE is written as

We recall from Grosse et al. [18] that the geometric path is closed for distributions in the exponential family: all distributions along the path remain in the exponential family. Furthermore, their Fisher information can be written in terms of the terms parameters [18, Eq. 17]; this is based off a a result of exponential families from [52, Section 3.3]

where µ ( t ) are the generalized moments, defined as µ ( t ) = E x ∼ p t ( x ) [ t ( x )] = ∇ θ log Z t ( θ ) . It follows,

The geometric path is defined in parameter space by θ t = t θ 1 +(1 -t ) θ 0 , therefore

as in [18, Eq. 17]. For exponential families, log Z ( θ ) is convex in θ . Here, we further assume strong convexity (with constant M ) and smoothness (with constant L ) so that

is polynomial in the euclidean distance between the parameters.

Proof of Theorem 4 Exponential error of annealed NCE with the arithmetic path and "vanilla" schedule

We now study the estimation error produced by the arithmetic path with "vanilla" schedule (table 2, line 3). Similarly, we start with the formula of the estimation error of NCE in the limit of a continuous path

where I ( t ) = E x ∼ p ( x ,t ) [( d dt log p ( x , t )) 2 ] is the Fisher information over the path, using time t as the parameter. The arithmetic path is a Gaussian mixture (see table 2) so we will conveniently use the parametric form of the path to compute the Fisher information

where the parameter here is the weight of the Gaussian mixture ˜ w t = tZ 1 / ( tZ 1 +1 -t ) . We will need to compute two quantities: the Fisher information to that mixture parameter (not the time), and the parameter speed ˙ ˜ t w .

so that the MSE

which is an f-divergence with generator ϕ ( x ) = (1 -x ) 2 / (1 + x ) that provides a t -independent lower bound. This will allow us to factor this quantity out of the integral defining the MSE, and simplify computations. We also have

where we choose to keep the dependency on t . The intuition is that integrating this quantity will yield a function of Z 1 , which will drive the MSE toward high values. We next show this rigorously and finally compute the estimation error.

We would like to write Z 1 in terms of the parameters. To do so, we now suppose the unnormalized target is in a simply unnormalized exponential family. Consequently,

using the strong convexity of the log-partition function. For large enough ∥ θ 1 -θ 0 ∥ > 0 , the quadratic term in the exponential is larger than the linear term, and the divergence D ϕ ( p 1 , p 0 ) is larger than a constant. It follows that for large enough ∥ θ 1 -θ 0 ∥ > 0 , there exists a constant C > 0 such that the MSE grows (at least) exponentially with the parameter-distance

Proof of Theorem 5 Polynomial error of annealed NCE with the arithmetic path and "oracle" schedule

We now study the estimation error produced by the arithmetic path with "oracle" schedule (table 2, line 4). Similarly, we start with the formula of the estimation error of NCE annealed over a continuous path

where I ( t ) = E x ∼ p ( x ,t ) [( d dt log p ( x , t )) 2 ] is the Fisher information over the path, using time t as the parameter. The arithmetic path is the Gaussian mixture p t ( x ) = tp 1 ( x ) + (1 -t ) p 0 ( x ) (see table 2). The Fisher information is therefore

where we choose to keep the dependency on t in the bound.

We briefly justify this choice. We had first tried a t -independent bound, which led to an upper bound of the MSE that was too loose. We share insight as to why: first, recognize that the fraction can be broken in two terms, each of them a chi-square divergence between an endpoint of the path ( p 0 or p 1 ) and the mixture p t . Each of them admits a t -independent upper bound given by the chi-square divergence between the endpoints p 0 and p 1 , using lemma 1. However, the chi-square divergence between two Gaussians, for example, is exponential (not polynomial) in the natural parameters [51, eq 7.41]. In fact, plotting I ( t ) for a univariate Gaussian model revealed that it took high values at the endpoints t = 0 and t = 1 , and was near zero almost everywhere else in the interval t ∈ [0 , 1] , which again suggested that dropping the dependency on t was unreasonable.

Now we can compute the estimation error, as

Let us try to solve one of these integrals, say J 1 .

which we rewrote using an f-divergence defined by ϕ ( x ) = -log( x ) 1 -x -1 . Similarly, we obtain

Putting this together, we get

How does this divergence depend on the parameter-distance ∥ θ 1 -θ 0 ∥ ? Does it bring down the dependency from exponential to something lower? We next analyze this:

which looks like a Kullback-Leibler divergence, where the integrand is reweighted by 1 1 -p 0 ( x ) p 1 ( x ) . Note that

which motivates separating the integral over both domains

Hence we get

We now suppose the proposal and target are distributions in an exponential family. The KL divergence between exponential distributions with parameters θ 0 and θ 1 , is given by the Bregman divergence of the log-partition on the swapped parameters [34, Eq. 29]

Hence

using the L -smoothness of the log-partition function log Z ( θ ) .

Discussion on the assumptions for theorems 2, 3, 4, 5 For these theorems, we have supposed that the target and proposal distributions are in an exponential family with a log partition that verifies

We now look at the validity of this assumption for a simple example: the univariate Gaussian, which is in an exponential family. The canonical parameters are its mean and variance ( µ, v ) . Written as an exponential family,

the natural parameters are θ = ( µ/v, -1 / (2 v )) , associated with the sufficient statistics t ( x ) = ( x, x 2 ) [34]. The log-partition function and its derivatives are

When the mean is zero, the eigenvalues of the hessian are in fact the diagonal values ( v, 2 v 2 ) , and they are bounded if and only if the variance v is bounded.

Proof of Theorem 6 Constant error of annealed NCE with the arithmetic path and "oracle-trig" schedule

We now study the estimation error produced by the arithmetic path with "oracle-trig" schedule (table 2, line 5). We write the optimal path of Eq. 66 in the limit when the distributions have little overlap, pointwise (with error ϵ ′ ) and on average (with error ϵ ):

We briefly note that there exist certain conditions, given by the dominated convergence theorem, where the first error (pointwise) going to zero implies the second (on average) going to zero as well, but this is outside the scope of this proof. Now, many relevant quantities involved in the optimal distribution and its estimation error simplify as ϵ ′ ( x ) → 0 pointwise and ϵ → 0 . The notation O ( · ) will hide dependencies on absolute constants only.

This leads to the following simplification of the optimal path

where we denoted by

the arithmetic path with "oracle-trig" schedule defined in table 2 (line 5); the trigonometric weights in evolve slowly at the end points t = 0 and t = 1 . We also define g 1 ( x ) = sin( πt )( p 0 ( x )( t -1) -p 1 ( x ) t ) which is an integrable function. This proves the first part of this theorem, which is that the optimal path p opt is close to a certain arithmetic path (with trigonometric weights) p arith -trig , and that closeness is controlled by how little overlap there is between the endpoint distributions p 0 and p 1 , on average and pointwise.

Similarly, we can control how close these two paths are in terms of estimation errors. The estimation error in Eq. 70 produced by a path is

We denote by MSE optimal and MSE arith -trig the estimation errors produced respectively by the optimal path and the arithmetic path with trigonometric weights. The estimation error of the optimal path can be Taylor-expanded in terms of the overlap ( ϵ and ϵ ′ ( x ) ) as well. We next compute the intermediate quantities that are required to obtain the estimation error.

where g 2 ( x ) = p 0 ( x ) ( sin( πt )+ π ( t -1) cos( πt ) ) + p 1 ( x ) ( -sin( πt ) -πt cos( πt ) ) is an integrable function. It follows,

by expanding and using that ∂ t p arith -trig t ( x ) = π sin( πt ) ( p 1 ( x ) -p 0 ( x ) 2 ) is bounded so that ∂ t p arith -trig t ( x ) × ϵ ′ ( x ) = O ( ϵ ′ ( x )) . Moreover,

Denoting by ϵ ′′ ( x ) := ( O ( ϵ ′ ( x )) + O ( ϵg 1 ( x ))) /p arith -trig t ( x ) a third quantity which we assume goes to zero, we get

We can now write

by expanding. We can then write

Additionally, we assume that the remainder term denoted by

is integrable and also goes to zero. On the other hand, the estimation error produced by the optimal path is known [17, Eq. 48] and the result can be Taylor-expanded as well

This means that

C Useful Lemma

Lemma 1 Chi-square divergence of between a density and a mixture We wish to upper bound the chi-square divergence between a distribution p ( x ) and a mixture wp ( x ) + (1 -w ) q ( x ) , where 0 < w < 1 .

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe article theoretically bounds the estimation error of NCE depending on the path. The paper first shows that NCE has minimum variance among all estimators using K bridging distributions and then continues by showing that the standard polynomial path can achieve polynomial error in the limit of infinitely many bridging distributions. Then the authors derive first a negative result for the arithmetic path and a positive result on the oracle estimate of the path that uses the unknown normalisation constant. The results are empirically verified.\n\nSTRENGTHS:\nThe authors responses clarified the main weaknesses and I am raising my score accordingly.\n--------------------------\n\nThe paper appears to be well written and some of the results appear to be novel, especially the error bounds. I have not checked the math in the appendix.\n\nThe presentation is rather clear and to the point and the experiments nicely validate the theorems.\n\nWEAKNESSES:\n- Some related work is missing. importantly, NCE is a method that got discovered and rediscovered several times, see for example this review that calls the same method BAR (Bennets acceptance ratio): \n\n[1] Krause, Oswin, Asja Fischer, and Christian Igel. \"Algorithms for estimating the partition function of restricted Boltzmann machines.\" Artificial Intelligence 278 (2020): 103195.\n\n- As a result of this, Theorem 1 was already known, as [1] reproduces an earlier proof that shows that NCE is a maximum likelihood estimator. Since MLEs are efficient, their variance is bounded by the variance of the Fisher matrix. In that light, the novelty of Theorem 1 is to produce a bound on the Fisher matrix via the path integral. Similarly, [1] already proposed a framework that included the generalisation of several estimators, similar to the Bregman methodology proposed in this work.\n\n- The practical applicability of the results is questionable. While the work shows that the two step estimator works well given perfect samples, it might be very challenging to obtain good samples for the arithmetic path. This is because the modes that appear are disconnected from each other from the very beginning on, which makes it difficult to discover them using MCMC, unlike for example parallel tempering with the polynomial path (see [1]).\n\nQUESTIONS:\nI would suggest to discuss [1] in the article due to its relevance and clear overlap. A few of the references therein might be important as well.\n\nQ1: you stress in the article that your proposed two-step method is not efficient. Do you foresee a way to make it efficient? Is it possible to create a larger experiment on a real estimation task that shows that it can improve on previous estimates?\n\n\n\nFinally, on page 5 line 165 where you quote [29], I would propose to also mention Deep Belief networks as an important model class with universal approximation properties for which estimation the normalisation constant is relevant:\n\nKrause, Oswin, et al. \"Approximation properties of DBNs with binary hidden units and real-valued visible units.\" International conference on machine learning. ICML, 2013.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper investigates the benefits of using bridge distributions to estimate the unknown normalizing constant $Z_1$ of a given density $p_1 = f_1 / Z_1$, with $f_1$ known. There are many ways for estimating a normalizing constant, e.g. importance sampling, bridge sampling, umbrella sampling and noise contrastive estimation. For example, in importance sampling an instrumental density $p_0$ from which it is easy to sample is introduced and the normalizing constant is estimated by weighting the samples from $p_0$ according to the ratio $f_1 / p_0$. Naturally, the accuracy of the estimation depends on how far $p_0$ is from $p_1$ (in Rényi 2-divergence for example or KL), and in most practical cases $p_0$ will be in fact very far from $p_1$. A nice idea then consists in introducing bridge distributions that \"simplify\" the estimation; we introduce a sequence of distributions $p_k$ (which themselves have intractable normalizing constants) such that the discrepancy between $p_{k+1}$ and $p_k$ is small and use them to estimate the normalizing constant. There are many ways for choosing such a sequence of distributions including geometric path, arithmetic path.  \nThis paper investigates the effect of each design choice: the estimator (they challenge the default choice of IS) and the sequence of bridging distributions. More specifically, their contributions are the following:  \n- They show that the NCE loss is optimal, whatever the number $K$ of bridging distributions and thus showing that IS is in fact not optimal. This extends the result for $K = 1$.  \n- The performance of IS is known to suffer from exponential dependence on the dimension and this phenomenon is now sharply quantified. In this paper it is shown that this is also the case for NCE at least for the exponential family, which hints that it should be also the case for more complex densities.  \n- The hope of using the bridging distributions was to diminish the effect of the dimensionality. It is shown that in the limit of infinite bridge distributions and with geometric paths, the error is no longer exponential but only polynomial.  \n- Finally, in the presence of an oracle providing the target normalizing constant $Z_1$, they show that the arithmetic path can provide an error independent of the dimension.\n\nSTRENGTHS:\nI find the theoretical contributions of this paper original, interesting and significant. Indeed, I appreciate that the authors have deeply investigated some a priori knowledge that practitioners have, such as the optimality of IS and the effect of the annealed distributions. Furthermore, I very much appreciate that the theoretical results are validated empirically on toy examples.\n\nWEAKNESSES:\nThe assumptions made in the paper are very strong (perfect sampling for example) and some of the results are only asymptotic in the number of bridge distributions. This makes me doubt about the broad relevance of this paper to the NeurIPS community. Furthermore, although the proposed two step procedure was shown to be numerically optimal, it is not tested numerically in the realistic scenario where one has to resort to MCMC for sampling. I believe that there is no reason why this should not be tested and compared to more traditional methods. Otherwise it is practically useless.\n\nQUESTIONS:\nIf I am not mistaken, the estimator $\\hat{Z}_1$ in the first line of table 1 is not an estimator of $Z_1$ as $f_0$ is not normalized. Also, in eq. 4 and eq. 5 I believe that there should be $+ \\log Z_0$.\n\nLIMITATIONS:\nSee weaknesses", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-06-27"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies annealing schedules used in noise contrastive estimation (NCE). In particular, it conducts rigorous theoretical analyses of existing annealing schedules used in NCE. The paper's framework is general and yields importance sampling, umbrella sampling, and bridge sampling as special cases. The theory suggests that NCE is more efficient in certain scenario compared to annealed importance sampling (AIS).\n\nSTRENGTHS:\nDisclaimer: I am new to noise contrastive estimation, but familiar with the annealed importance sampling literature.\n\n* The paper provides a rigorous analysis of different annealing schedules used in noise contrastive estimation.\n* Previous methods for estimating normalizing constants are obtained as a special case, such as umbrella sampling, ration sampling, and bridge sampling.\n* The paper shows that the popularly used geometric annealing schedule is suboptimal.\n\nWEAKNESSES:\nGiven my lack of expertise in terms of NCE, I will focus my comments on the claims about AIS.\n\n* A major concern with the theoretical comparison against AIS is that the specific instance of AIS considered in this paper is known to be suboptimal. The analysis of AIS nowadays involves the \"backward sampling path\" formalism (first thoroughly explored in [1], but also tackled in Neal's paper [2]), where the paper's formulation of AIS corresponds to the \"suboptimal $L$-kernel\" (kernel implicitly used for backward sampling), where the suboptimality is already in the name. There exists attempts to obtain approximately optimal $L$-kernels [3,4].\n* This is a minor point, but the fact that the AIS with the suboptimal $L$-kernel scales poorly (which is what is referred to as AIS in the paper) is already known. See p. 335 in [5], p. 11 in [6].\n* As such, the claimed scope of the paper is way too broad compared to the actual contribution of the paper. I strongly recommend stating that this is a paper focusing exclusively on noise contrastive estimation in the title, abstract, and introduction.\n* Also, I think the paper should be more subtle when comparing performance against importance sampling; with the backward sampling \nformalism, you do not need to \"assume\" perfect sampling; you still have non-asymptotic performance guarantees through [7], and asymptotic guarantees through [1]. So not being able to capture this is theoretically quite critical.\n\nQUESTIONS:\n### Major Comments\n* Line 40: As mentioned above, we do have some understanding on how to tune the performance of annealing-based normalizing constant estimation. In fact, although less formal, Neal's paper [2] contains a lengthy discussion on this. Why is this not mentioned?\n* Line 189-190: \"... there is no definitive theory on the ability of annealing to reduce the statistical cost in a general setup.\" Again, I believe [2] provides some theory, and with the backward sampling formalism, further analysis is possible as in [3], which is based on the results of [7].\n* Line 136-137: This is true for importance sampling, but not necessarily for AIS. In fact, this point in the cited reference [27] is shown with only importance sampling. (Thus, the quotation needs to be corrected.) In AIS, with an optimal $L$-kernel, one can achieve the variance of $i.i.d$ samples regardless of the tails, although whether one can obtain an optimal $L$-kernel is something else.\n* Line 65-79: On a similar note, this is a description specific to thermodynamic integration/bridge sampling and is not how newer AIS approaches do things nowadays. In fact, in AIS/SMC, you do not need to choose a path before sampling. You can adaptively determine the sequence based on the estimated quality of samples [6].\n\n### References\nDisclaimer: I am not the author or affiliated with the authors of the papers below.\n* [1] Del Moral, Pierre, Arnaud Doucet, and Ajay Jasra. \"Sequential monte carlo samplers.\" Journal of the Royal Statistical Society Series B: Statistical Methodology 68.3 (2006): 411-436.\n* [2] Neal, Radford M. \"Annealed importance sampling.\" Statistics and computing 11 (2001): 125-139.\n* [3] Bernton, Espen, et al. \"Schr\\\" odinger Bridge Samplers.\" arXiv preprint arXiv:1912.13170 (2019).\n* [4] Doucet, Arnaud, et al. \"Score-based diffusion meets annealed importance sampling.\" Advances in Neural Information Processing Systems 35 (2022): 21482-21494.\n* [5] Chopin, Nicolas, and Omiros Papaspiliopoulos. An introduction to sequential Monte Carlo. Vol. 4. New York: Springer, 2020.\n* [6] Dai, Chenguang, et al. \"An invitation to sequential Monte Carlo samplers.\" Journal of the American Statistical Association 117.539 (2022): 1587-1600.\n* [7] Agapiou, Sergios, et al. \"Importance sampling: Intrinsic dimension and computational cost.\" Statistical Science (2017): 405-431.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this work, the authors make a number of contributions to the area of estimating normalization factors using annealing from some \"proposal\" $p_0$ to a \"target\" $p_1$.  \n\nThe authors start off by nicely extending recent works by Chehab et al. (2023) on the relation between importance sampling (IS) and noise-constrastive estimation (NCE) to the scenario involving annealing, introducing the notion of *annealed Bregman estimators (ABE)* in the proces, computing hte intermediate log-ratios by solving a classificiation task between samples drawn from neighboring densities.  \n\nMaking in particular use of the ABE defined by the NCE loss, the authors then obtain a number of asymptotic results (in the number of temperatures $K \\to \\infty$) for the mean-squared error of the estimator of $log Z_1$, i.e. the normalization constant of the target $p_1$.  \n\nIn Theorem 3, they show that in the case of the commonly used geometric path $p_t = p_0^{1 - \\beta} p_1^{\\beta}$, both annealed IS and annealed NCE result in an MSE estimate that polynomial in the parameter distance between $p_0$ and $p_1$. This result is established under assumptions of the exponential family for both the proposal and the target, which is a more general setting than seen in previous works. Moreover, the authors also show in Theorem 2, under almost identical assumptions to Theorem 3, that even with the NCE loss instead of IS, the MSE will still scales exponentially in the distance between the parameters of the proposal and target, as is a well-known fact for the IS-based estimator.  \nCombining these two results, the authors have demonstrated theoretically that *annealing* allows us to move from exponential to polynomial dependence on the parameter distance, effectively bridging the gap between many empirical results from the literature where annealing has been observed to be of much help even in higher dimensions vs. standard importance sampling.  \n\nFurthermore, in Theorem 4, 5, and 6, the authors show that the MSE when using the recently introduced arithmetic path $p_t = (1 - \\beta) p_0 + \\beta p_1$, has exponential, polynomial, and constant dependence on the parameter distance, respectively, under different path-parameterizations. Of course, this path requires the normalization constant in its computation, but the authors use the observation to propose a two-step estimation procedure which they then demonstrate to be useful on a simple toy-problem.\n\nSTRENGTHS:\nI believe this work to be of excellent standard.  \n\nThe authors are very nicely combining the previous works of Gelman & Meng (1998) and Chehab et al. (2023) to provide both additional theoretical and empirical insights. The theoretical results are, as far as I can tell, both very novel and very relevant. In particular, I find Theorems 2 & 3 to be very pleasing as they give us a clear theoretical motivation for the usage of annealing in higher dimensions, which is a result long sought after. In addition to this, they also bring the work of Chehab et al. (2023) into the annealing setting, and demonstrates its utility both theoretically and empirically (though on toy-problems).  \n\nThe presentation is of very high quality; the text is well-written, the results are presented in a clear and nice manner, illustrations are used to get the intuition across, and tables are used when appropriate to give the reader a quick overview of results.  \n\nAs far as I can tell, the authors also do a good job mentioning previous works and the limitations of their own work.\n\nWEAKNESSES:\nThere are honestly very few weaknesses to point out when it comes to the paper itself; the only weaknesses I can find are related to limitations of the analysis and the resulting proposed practical methods, which I will leave to the \"Limitations\" section. The one aspect I'd like to raise, though this is slightly related to the aforementioned limitations, is that I would have liked to see slightly more extensive empirical results for the proposed methods. I do agree with the authors in that the empirical results are not the main focus of the work, but still, have a slightly more complex problem would have been nice to see how, for example, the two-step (trig) approach would do in this.\n\nQUESTIONS:\n-   Theorem 6: In Gelman & Meng (1998) in Section 4.3, they make the following comment \"we doubt (51) is achievable as it is bounded above by $\\pi^2 / n$ even if $p_0$ and $p_1$ are infinitely apart\" referring to the continuous-time limit.  \n    This is then slightly sharper than the $2 \\pi^2 / N$ result presented in this work. Can the authors comment on this? Was this just a slightly overly confident statement by Gelman & Meng, i.e. they really meant including the factor of 2?\n\nLIMITATIONS:\nThe main limitations of the theoretical analysis are of course the assumptions.  \n\nThe authors assume that both the proposal and the target are in the exponential family, in addition varying assumptions on the normalization constant of the target as a function of the parameters. This is often not the case in practice, but, as mentioned by the authors, the exponential family some universality properties, which one could argue makes this slightly less of an issue.  \n\nAs most other works on annealing, this work also assumes perfect sampling from the intermediate distributions, which is usually not the case in practice. But, as mentioned, this is a very common assumption to make. This is also mentioned by the authors.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}