SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Then in the annealing limit of a continuous path  $K \to \infty$ , the estimation error of annealed NCE with the geometric path is (at most) polynomial in the parameter-distance between the proposal and the target:
$$
\operatorname {M S E} _ {\mathrm {N C E}} \left(p _ {0}, p _ {1}; q, K, N\right) \leq \frac {L ^ {2}}{M N} \| \boldsymbol {\theta} _ {1} - \boldsymbol {\theta} _ {0} \| ^ {2} + o \left(\frac {1}{N}\right) + o \left(\frac {K ^ {2}}{N}\right), \quad \text {w h e n} q = 0 \tag {14}
$$
where  $M$  and  $L$  are respectively the strong convexity and smoothness constants of  $\log Z(\pmb{\theta})$ .
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 noise-contrastive 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\to 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 (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).
Consider the annealing limit of a continuous path  $K \to \infty$  path and a far-away target with large enough  $\| \pmb{\theta}_1 - \pmb{\theta}_0 \| > 0$ . For estimating the log normalization of the (unnormized) 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:
$$
\mathrm {M S E} _ {\mathrm {N C E}} \left(p _ {0}, p _ {1}; q, K, N\right) > \frac {C}{N} \exp \left(\frac {M}{2} \| \boldsymbol {\theta} _ {1} - \boldsymbol {\theta} _ {0} \| ^ {2}\right) + o \left(\frac {1}{N}\right) + o \left(\frac {K ^ {2}}{N}\right), \text {w h e n} q = 1 \tag {15}
$$
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}(\pmb{x}) \times 100$  and its new normalization is  $Z_{1} \times 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 \to \frac{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:
$$
\operatorname {M S E} _ {\mathrm {N C E}} \left(p _ {0}, p _ {1}; q, K, N\right) \leq \frac {1}{N} \left(2 + L \| \boldsymbol {\theta} _ {1} - \boldsymbol {\theta} _ {0} \| ^ {2}\right) + o \left(\frac {1}{N}\right) + o \left(\frac {K ^ {2}}{N}\right), \quad \text {w h e n} q = 1 \tag {16}
$$
where  $L$  is the smoothness constant of  $\log Z(\pmb{\theta})$ .
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\to \infty$  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 \to \infty$ , and of a target distribution whose overlap with the proposal
Path name	Unnormalized density	Normalized density	Error
Geometric	f_t(x) = p_0(x)^1-t f_1(x)^t	p_t(x) ∝ p_0(x)^1-t p_1(x)^t	poly
Arithmetic	f_t(x) = (1 - w_t)p_0(x) + w_t f_1(x)	p_t(x) = (1 - \tilde{w}_t)p_0(x) + \tilde{w}_t p_1(x)
vanilla	w_t = t	\tilde{w}_t = \frac{tZ_1}{(1 - t) + tZ_1}	exp
oracle	w_t = \frac{t}{t + Z_1(1 - t)}	\tilde{w}_t = t	poly
oracle-trig	w_t = \frac{\sin^2\left(\frac{\pi t}{2}\right)}{\sin^2\left(\frac{\pi t}{2}\right) + Z_1\cos^2\left(\frac{\pi t}{2}\right)}	\tilde{w}_t = \sin^2\left(\frac{\pi t}{2}\right)	const
Table 2: Geometric and arithmetic paths, defined in the space of unnormalized densities (second column); "oracle" and "oracle-trig" are reparameterizations of the arithmetic path which depend on the true normalization  $Z_{1}$ . The corresponding normalized densities (third column) produce an estimation error (fourth column) which we quantify.
goes to zero. Namely, consider the quantities:
$$
\epsilon^ {\prime} (\boldsymbol {x}) := \sqrt {p _ {0} (\boldsymbol {x}) p _ {1} (\boldsymbol {x})} \quad \epsilon := \int_ {\mathbb {R} ^ {D}} \epsilon^ {\prime} (\boldsymbol {x}) d \boldsymbol {x} \tag {17}
$$
$$
\epsilon^ {\prime \prime} (\boldsymbol {x}) := \frac {\epsilon^ {\prime} (\boldsymbol {x}) - \epsilon \sin (\pi t) p _ {t} ^ {\text {o r a c l e}} (\boldsymbol {x})}{p _ {t} ^ {\text {o r a c l e - t r i g}} (\boldsymbol {x})} \quad \epsilon^ {\prime \prime} := \int_ {\mathbb {R} ^ {D}} \int_ {0} ^ {1} \epsilon^ {\prime \prime} (\boldsymbol {x}) \left(1 + \frac {\left(\partial_ {t} p _ {t} ^ {\text {o r a c l e - t r i g}} (\boldsymbol {x})\right) ^ {2}}{p _ {t} ^ {\text {o r a c l e - t r i g}} (\boldsymbol {x})}\right) d t d \boldsymbol {x}. \tag {18}
$$
Assume  $\sup_{\pmb{x} \in \mathbb{R}^D} \epsilon'(x) \to 0, \sup_{\pmb{x} \in \mathbb{R}^D} \epsilon''(x) \to 0, \epsilon \to 0, \epsilon''' \to 0,$  and consider the distributions  $p_t^{\text{arith-trig}}$  and  $p_t^{\text{oracle}}$  as defined in Table 2.
Then, the optimal annealing path convergences pointwise to an arithmetic path reparameterized trigonometrically with  $t \to \frac{t}{\sin^2(\frac{\pi t}{2}) + Z_1(1 - \sin^2(\frac{\pi t}{2}))}$  and the estimation error tends to the optimal estimation error (which is constant with respect to the parameter-distance):
$$
\mathrm {M S E} _ {\mathrm {N C E}} \left(p _ {0}, p _ {1}; q, K, N\right) = \frac {1}{N} \pi^ {2} + O \left(\frac {\epsilon + \epsilon^ {\prime \prime \prime}}{N}\right) + o \left(\frac {1}{N}\right) + o \left(\frac {K ^ {2}}{N}\right) \quad \text {w h e n} q = 1. \tag {19}
$$
By way of remarks, we note that the assumptions that the quantities  $\sup \epsilon'(x), \sup \epsilon''(x), \epsilon, \epsilon''$  go to 0 are mutually incomparable (i.e. none of them implies the others). For example,  $\sup_{\boldsymbol{x} \in \mathbb{R}^D} \epsilon'(\boldsymbol{x})$  and  $\epsilon$  going 0 require that the function  $\sqrt{p_0(x)p_1(x)}$  goes to 0 in the  $L_{\infty}$  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 pre-estimating a problematic (hyper)parameter, here  $Z_{1}$ , has proved beneficial to reduce the estimation error of NCE in a related context [36].
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 = \mathcal{N}(\mathbf{0}, 2\mathbf{Id})$  in Figure 2,  $p_1 = \mathcal{N}(\mathbf{0}, 0.25\mathbf{Id})$  in Figure 3b and  $p_1 = \mathcal{N}(\mathbf{0}, \sigma^2\mathbf{Id})$  in Figure 3a, where the target variance decreases as  $\sigma(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 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 \to \infty$ ) 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(\theta_1) = 1$ , so that the arithmetic path is effectively the same as the "two-step" estimator.
Figure 2: Optimality of the NCE loss, using the geometric path with  $K = 2$ . NCE has the smallest deviation from zero, the true value of the log normalizing constant.
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
(a) MSE against parameter distance.
(b) MSE against dimensionality.
Figure 3: Estimation error as the target and proposal distributions grow apart. Without annealing, the error is exponential in the parameter distance (diagonal in log-scale). Annealing with the geometric path and our two-step methods brings down the error to slower growth, as predicted by our theorems.
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.
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.
[1] 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.
[2] A. Hyvarinen. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research, 6(24):695-709, 2005.
[3] M. Gutmann and A. Hyvarinen. Noise-contrastive estimation of unnormalized statistical models, with applications to natural image statistics. Journal of Machine Learning Research, 13(11): 307–361, 2012.
[4] G.E. Hinton. Training products of experts by minimizing contrastive divergence. Neural Computation, 14(8):1771-1800, 2002.
[5] 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.
[6] A.B. Owen. Monte Carlo theory, methods and examples. https://artowen.sudomains/ mc/, 2013.
[7] 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.
[8] 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.
[9] O. Chehab, A. Gramfort, and A. Hyvarinen. 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.
[10] R.M. Neal. Annealed importance sampling. Statistics and Computing, 11:125-139, 1998.
[11] C. Jarzynski. Nonequilibrium equality for free energy differences. Physical Review Letters, 78: 2690-2693, 1996.
[12] C. Jarzynski. Equilibrium free-energy differences from nonequilibrium measurements: A master-equation approach. Physical Review E, 56:5018-5035, 1997.
[13] R. Salakhutdinov and G. Hinton. Replicated softmax: an undirected topic model. In Neural Information Processing Systems (NIPS), 2009.
[14] 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.
[15] V. Masrani, T.A. Le, and F.D. Wood. The thermodynamic variational objective. In Neural Information Processing Systems (NeurIPS), 2019.
[16] 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.
[17] A. Gelman and X.-L. Meng. Simulating normalizing constants: From importance sampling to bridge sampling to path sampling. Statistical Science, 13:163-185, 1998.