SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Omar Chehab
Université Paris-Saclay, Inria, CEA Palaiseau, France
l-emir-omar.chehab@inria.fr
Aapo Hyvarinen
aapo.hyvarinen@helsinki.fi
Andrej Risteski
Department of Machine Learning
Carnegie-Mellon University
Pittsburgh, USA
aristesk@andrew.cmu.edu
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.
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:
1. 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.
2. 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.
3. 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.
Importance sampling and NCE The problem considered here is computing the normalization constant  $^1$ , i.e. the integral of some unnormalized density  $f_{1}(\pmb {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  $\phi(\pmb{x})$ . For example, minimizing the classification loss identified by  $\phi_{\mathrm{IS}}(x) = x \log x$  yields the importance sampling estimator. Similarly,  $\phi_{\mathrm{RevIS}}(x) = -\log x$  leads to the reverse importance sampling estimator [23], and  $\phi_{\mathrm{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}(\pmb{x})$ .
These intermediate distributions are obtained from a path  $(f_t)_{t \in [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
$^1$ in this paper we also say we "estimate" the normalization constant, though in classical statistics it is more traditional to use "estimation" when referring to parameters of a statistical model
Name	Loss identified by φ(x)	Estimator Z1	MSE
IS	x log x	Ep0f1p0	1+ν/νN Dχ2(p1,p0)
RevIS	- log x	(Ep1p0f1)-1	1+ν/N Dχ2(p0,p1)
NCE	x log x - (1+x) log(1+x/2)	implicit	(1+ν)2/νN D_HM(p1,p0)/1-D_HM(p1,p0)
Table 1: Some estimators of the normalization obtained by minimizing a classification loss, and their estimation error in terms of well-known divergences [22]. For details and definitions, see Appendix A.
General (q ∈]0,1])	ft(x) = ((1 - t)p0(x)q + tf1(x)q)1/q
Geometric (q → 0)	ft(x) = p0(x)1-t × f1(x)t
Arithmetic (q = 1)	ft(x) = (1 - t)p0(x) + tf1(s)
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 \to 0$  defines 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<sup>2</sup> discretized into a sequence of  $K + 1$  unnormalized densities, denoted by  $(f_{k / K})_{k\in [0,K]}$  with corresponding normalizations  $(Z_{k / K})_{k\in [0,K]}$ . In practice, samples are drawn from the corresponding normalized densities  $(p_{k / K})_{k\in [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 minimizing 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
${}^{2}$  other discretization schemes can be equivalently achieved by re-parameterizing the path [17]
normalization  $\log Z_{1}$  is estimated additively along the sequence of distributions
$$
\widehat {\log Z _ {1}} = \sum_ {k = 0} ^ {K - 1} \log \left(\frac {Z _ {(k + 1) / K}}{Z _ {k / K}}\right) + \log Z _ {0}. \tag {1}
$$
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  $\phi_k(\pmb{x})$  and defined as
$$
\mathcal {L} _ {k} \left(\beta_ {k}\right) := \mathbb {E} _ {\boldsymbol {x} \sim p _ {k / K}} \left[ \phi_ {k} ^ {\prime} \left(r _ {k} (\boldsymbol {x}; \beta_ {k})\right) \times r _ {k} (\boldsymbol {x}; \beta_ {k}) - \phi_ {k} \left(r _ {k} (\boldsymbol {x}; \beta_ {k})\right) \right] - \mathbb {E} _ {\boldsymbol {x} \sim p _ {(k + 1) / K}} \left[ \phi_ {k} ^ {\prime} \left(r _ {k} (\boldsymbol {x}; \beta_ {k})\right) \right], \tag {2}
$$
where the regression function  $r_k(\pmb{x};\beta_k)$  is parameterized by the unknown log-ratio  $\beta_{k}$
$$
r _ {k} (\boldsymbol {x}; \beta_ {k}) = \exp \left(- \beta_ {k}\right) \times f _ {(k + 1) / K} (\boldsymbol {x}) / f _ {k / K} (\boldsymbol {x}). \tag {3}
$$
For the true  $\beta_{k}^{*}$ , it holds  $\beta_{k}^{*} = \log (Z_{(k + 1) / K} / Z_{k / K})$  and  $r_k(\pmb {x};\pmb {\beta}_k^*) = p_{(k + 1) / K}(\pmb {x}) / p_{k / K}(\pmb {x})$ . The convex functions  $(\phi_k)_{k\in [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  $(\phi_k)_{k\in [0,K - 1]}$ .
Different combinations of these hyperparameters recover several common estimators of the log-partition 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  $(\phi_0 = x\log x)$  for the first ratio, and reverse importance sampling  $(\phi_1 = -\log x)$  for the second ratio, recovers the bridge sampling estimator as a special case [20]
$$
\widehat {\log Z _ {1}} = - \log \mathbb {E} _ {p _ {1}} \frac {f _ {1 / 2}}{f _ {1}} + \log \mathbb {E} _ {p _ {0}} \frac {f _ {1 / 2}}{f _ {0}} + \log Z _ {0}. \tag {4}
$$
Alternatively, we can use these classification losses in reverse order: reverse importance sampling  $(\phi_0 = -\log x)$  for the first ratio, and importance sampling  $(\phi_1 = x\log x)$  for the second ratio, and recover the umbrella sampling estimator [19] also known as the ratio sampling estimator [25]
$$
\widehat {\log Z _ {1}} = \log \mathbb {E} _ {p _ {1 / 2}} \frac {f _ {1}}{f _ {1 / 2}} - \log \mathbb {E} _ {p _ {1 / 2}} \frac {f _ {0}}{f _ {1 / 2}} + \log Z _ {0}. \tag {5}
$$
Another option yet, is to use the same classification loss for all ratios. With importance sampling  $(\phi_{k} = x\log x,\forall k\in [[0,K - 1]])$ , we recover the annealed importance sampling estimator [10-12]
$$
\widehat {\log Z _ {1}} = \sum_ {k = 1} ^ {K} \log \mathbb {E} _ {\boldsymbol {x} \sim p _ {k - 1}} \left[ \frac {f _ {k}}{f _ {k - 1}} (\boldsymbol {x}) \right] + \log Z _ {0}. \tag {6}
$$
The family of annealed Bregman estimators is visibly large enough to include many existing estimators, obtained for different hyperparameter 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,  $\phi_{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 = \mathcal{N}(\mathbf{0},\mathbf{Id})$  and  $p_1 = \mathcal{N}(\mathbf{0},2\mathbf{Id})$ , the geometric path produces Gaussians with increasing variances  $\Sigma_t = (1 - t / 2)^{-1}\mathbf{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 super-Gaussian [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:
$$
\operatorname {M S E} \left(p _ {0}, p _ {1}; q, K, N, \phi_ {\mathrm {N C E}}\right) \leq \operatorname {M S E} \left(p _ {0}, p _ {1}; q, K, N, \phi\right), \quad \forall q, K, N, \phi . \tag {7}
$$
In the limit of  $K \to \infty$  (such that  $K^2 / N \to 0$ ), NCE, IS, and revIS converge to the same estimation error, given by the Fisher-Rao path length from the proposal to the target:
$$
\operatorname {M S E} \left(p _ {0}, p _ {1}; q, K, N, \phi\right) = \frac {1}{N} \int_ {0} ^ {1} I (t) d t + o \left(\frac {1}{N}\right) + o \left(\frac {K ^ {2}}{N}\right), \forall \phi \in \left\{\phi_ {\mathrm {N C E}}, \phi_ {\mathrm {I S}}, \phi_ {\mathrm {R e v I S}} \right\}, \tag {8}
$$
where the Fisher-Rao metric  $I(t) \coloneqq \mathbb{E}_{\boldsymbol{x} \sim p(\boldsymbol{x}, t)}[(\frac{d}{dt} \log p_t(\boldsymbol{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 \to \infty$  then in  $K \to \infty$ , 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  $\phi_{\mathrm{NCE}}$  with a subscript, instead. We highlight that our theorems in this paper apply to the MSE in the limit of  $K \to \infty$ : 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  $\pmb{t}(\pmb{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
$$
p (\boldsymbol {x}; \boldsymbol {\theta}) := \exp (\langle \boldsymbol {\theta}, \boldsymbol {t} (\boldsymbol {x}) \rangle - \log Z (\boldsymbol {\theta})) \tag {9}
$$
where  $Z(\pmb{\theta}) = \int \exp(\langle \pmb{\theta}_1, \pmb{t}(\pmb{x}) \rangle)$  is the partition function. We will consider that the (unnormized) target density  $f_1$  is what we call a simply unnormalized model defined as
$$
f _ {1} (\boldsymbol {x}) = \exp \left(\left\langle \boldsymbol {\theta} _ {1}, \boldsymbol {t} (\boldsymbol {x}) \right\rangle\right) \tag {10}
$$
Note that in general, a pdf can be unnormalized in many ways: one can multiply an unnormalized density by any positive function of  $\theta$  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(\pmb{\theta})$  is a convex function ("log-sum-exp") of the parameter  $\pmb{\theta}$  [34], which implies  $0 \preccurlyeq \nabla_{\pmb{\theta}}^{2} \log Z(\pmb{\theta})$ . In our theorems, we use the further assumptions of strong convexity with constant  $M$ , and/or smoothness with constant  $L$  (gradient is  $L$ -Lipschitz):
$$
\nabla_ {\boldsymbol {\theta}} ^ {2} \log Z (\boldsymbol {\theta}) \succcurlyeq M \mathbf {I d} \tag {11}
$$
$$
\nabla_ {\boldsymbol {\theta}} ^ {2} \log Z (\boldsymbol {\theta}) \preccurlyeq L \mathbf {I d} \tag {12}
$$
For exponential families, the derivatives of the log partition function yield moments of the sufficient statistics, and the hessian  $\nabla_{\theta}^{2}\log Z(\theta) = \mathrm{Cov}_{\boldsymbol{x}\sim p}[\boldsymbol {t}(\boldsymbol {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, positive-definiteness 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(\boldsymbol {x};\boldsymbol {\theta})$  , which is equivalent to the variance in every direction being bounded, which will be the case for parameters in a bounded domain  $\pmb {\theta}\in \Theta$  . 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 (unnormized) target  $f_1$  is from the simply unnormalized exponential family (Eq. 10). The log-partition function  $\log Z(\pmb{\theta})$  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 parameter-distance between the proposal and the target:
$$
\mathrm {M S E} _ {\mathrm {N C E}} \left(p _ {0}, p _ {1}; q, K, N\right) \geq \frac {4}{N} \exp \left(\frac {1}{8} M \| \boldsymbol {\theta} _ {1} - \boldsymbol {\theta} _ {0} \| ^ {2}\right) - 1 + o \left(\frac {1}{N}\right), \quad \text {w h e n} K = 1 \tag {13}
$$
where  $M$  is the strong convexity constant of  $\log Z(\pmb{\theta})$ .
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  $\| \pmb{\theta}_1 - \pmb{\theta}_0\|_2$  to scale as  $\sqrt{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 (unnormized) target  $f_1$  is from the simply unnormalized exponential family (Eq. 10). The log-partition function  $\log Z(\pmb{\theta})$  is assumed to be strongly convex and smooth (Eq. 11, Eq. 12).