Lost Relatives of the Gumbel Trick

Introduction

In this work we are concerned with the fundamental problem of sampling from a discrete probability distribution and evaluating its normalizing constant. A probability distribution $p$ on a discrete sample space $\mathcal{X}$ is provided in terms of its potential function $\phi:\mathcal{X}\to[-\infty,\infty)$, corresponding to log-unnormalized probabilities via $p(\mathbf{x})=e^{\phi(\mathbf{x})}/Z$, where the normalizing constant $Z$ is the partition function. In this context, $p$ is the Gibbs distribution on $\mathcal{X}$ associated with the potential function $\phi$. The challenges of sampling from such a discrete probability distribution and estimating the partition function are fundamental problems with ubiquitous applications in machine learning, classical statistics and statistical physics (see, e.g., Lauritzen, 1996).

Perturb-and-MAP methods (Papandreou & Yuille, 2010) constitute a class of randomized algorithms for estimating partition functions and sampling from Gibbs distributions, which operate by randomly perturbing the corresponding potential functions and employing maximum a posteriori (MAP) solvers on the perturbed models to find a maximum probability configuration. This MAP problem is NP-hard in general; however, substantial research effort has led to the development of solvers which can efficiently compute or estimate the MAP solution on many problems that occur in practice (e.g., Boykov et al., 2001; Kolmogorov, 2006; Darbon, 2009). Evaluating the partition function is a harder problem, containing for instance #P-hard counting problems. The general aim of perturb-and-MAP methods is to reduce the problem of partition function evaluation, or the problem of sampling from the Gibbs distribution, to repeated instances of the MAP problem (where each instance is on a different random perturbation of the original model).

The Gumbel trick (Papandreou & Yuille, 2011) relies on adding Gumbel-distributed noise to each configuration’s potential $\phi(\mathbf{x})$. We derive a wider family of perturb-and-MAP methods that can be seen as perturbing the model in different ways – in particular using the Weibull and Fréchet distributions alongside the Gumbel. We show that the new methods can be implemented with essentially no additional computational cost by simply averaging existing Gumbel MAP perturbations in different spaces, and that they can lead to more accurate estimators of the partition function.

Evaluating or perturbing each configuration’s potential with i.i.d. Gumbel noise can be computationally expensive. One way to mitigate this is to cleverly prune computation in regions where the maximum perturbed potential is unlikely to be found (Maddison et al., 2014; Chen & Ghahramani, 2016). Another approach exploits the product structure of the sample space in discrete graphical models, replacing i.i.d. Gumbel noise with a “low-rank” approximation. Hazan & Jaakkola (2012); Hazan et al. (2013) showed that from such an approximation, upper and lower bounds on the partition function and a sequential sampler for the Gibbs distribution can still be recovered. We show that a subfamily of our new methods, consisting of Fréchet, Exponential and Weibull tricks, can also be used with low-rank perturbations, and use these tricks to derive new upper and lower bounds on the partition function, and to construct new sequential samplers for the Gibbs distribution.

A family of tricks that can be implemented by simply averaging Gumbel perturbations in different spaces, and which can lead to more accurate or more sample efficient estimators of $Z$ (Section 2).

New upper and lower bounds on the partition function of a discrete graphical model computable using low-rank perturbations, and a corresponding family of sequential samplers for the Gibbs distribution (Section 3).

Discussion of advantages of the simpler analytical form of the Gumbel trick including new links between the errors of estimating $Z$, sampling, and entropy estimation using low-rank Gumbel perturbations (Section 4).

The idea of perturbing the potential function of a discrete graphical model in order to sample from its associated Gibbs distribution was introduced by Papandreou & Yuille (2011), inspired by their previous work on reducing the sampling problem for Gaussian Markov random fields to the problem of finding the mean, using independent local perturbations of each Gaussian factor (Papandreou & Yuille, 2010). Tarlow et al. (2012) extended this perturb-and-MAP approach to sampling, in particular by considering more general structured prediction problems. Hazan & Jaakkola (2012) pointed out that MAP perturbations are useful not only for sampling the Gibbs distribution (considering the argmax of the perturbed model), but also for bounding and approximating the partition function (by considering the value of the max).

Afterwards, Hazan et al. (2013) derived new lower bounds on the partition function and proposed a new sampler for the Gibbs distribution that samples variables of a discrete graphical model sequentially, using expected values of low-rank MAP perturbations to construct the conditional probabilities. Due to the low-rank approximation, this algorithm has the option to reject a sample. Orabona et al. (2014) and Hazan et al. (2016) subsequently derived measure concentration results for the Gumbel distribution that can be used to control the rejection probability. Maji et al. (2014) derived an uncertainty measure from random MAP perturbations, using it within a Bayesian active learning framework for interactive image boundary annotation.

Perturb-and-MAP was famously generalized to continuous spaces by Maddison et al. (2014), replacing the Gumbel distribution with a Gumbel process and calling the resulting algorithm A* sampling. Maddison (2016) cast this work into a unified framework together with adaptive rejection sampling techniques, based on the notion of exponential races. This recent view generally brings together perturb-and-MAP and accept-reject samplers, exploiting the connection between the Gumbel distribution and competing exponential clocks that we also discuss in Section 2.1.

Inspired by A* sampling, Kim et al. (2016) proposed an exact sampler for discrete graphical models based on lazily-instantiated random perturbations, which uses linear programming relaxations to prune the optimization space. Further recent applications of perturb-and-MAP include structured prediction in computer vision (Bertasius et al., 2017) and turning the discrete sampling problem into an optimization task that can be cast as a multi-armed bandit problem (Chen & Ghahramani, 2016), see Section 5.2 below.

In addition to perturb-and-MAP methods, we are aware of three other approaches to estimate the partition function of a discrete graphical model via MAP solver calls. The WISH method (weighted-integrals-and-sums-by-hashing, Ermon et al., 2013) relies on repeated MAP inference calls applied to the model after subjecting it to random hash constraints. The Frank-Wolfe method may be applied by iteratively updating marginals using a constrained MAP solver and line search (Belanger et al., 2013; Krishnan et al., 2015). Weller & Jebara (2014a) instead use just one MAP call over a discretized mesh of marginals to approximate the Bethe partition function, which itself is an estimate (which often performs well) of the true partition function.

Relatives of the Gumbel Trick

In this section, we review the Gumbel trick and state the mechanism by which it can be generalized into an entire family of tricks. We show how these tricks can equivalently be viewed as averaging standard Gumbel perturbations in different spaces, instantiate several examples, and compare the various tricks’ properties.

We write $\operatorname{Exp}(\lambda)$ for the exponential distribution with rate (inverse mean) $\lambda$ and $\operatorname{Gumbel}(\mu)$ for the Gumbel distribution with location $\mu$ and scale $1$. The latter has mean $\mu+c$, where $c\approx 0.5772$ is the Euler-Mascheroni constant.

1 The Gumbel Trick

Similarly to the connection between the Gumbel trick and the Poisson process established by Maddison (2016), we introduce the Gumbel trick for discrete probability distributions using a simple and elegant construction via competing exponential clocks. Consider $N$ independent clocks, started simultaneously, such that the $j$-th clock rings after a random time $T_{j}\sim\operatorname{Exp}(\lambda_{j})$. Then it is easy to show that (1) the time until some clock rings has $\operatorname{Exp}(\sum_{j=1}^{N}\lambda_{j})$ distribution, and (2) the probability of the $j$-th clock ringing first is proportional to its rate $\lambda_{j}$. These properties are also widely used in survival analysis (Cox & Oakes, 1984).

Property (2) says that the random variable $\operatorname*{argmin}_{x}T_{x}$, taking values in $\mathcal{X}$, is distributed according to $p$:

The Gumbel trick is obtained by applying the function $g(x)=-\ln x-c$ to the equalities in distribution (1) and (2). When $g$ is applied to an $\operatorname{Exp}(\lambda)$ random variable, the result follows the $\operatorname{Gumbel}(-c+\ln\lambda)$ distribution, which can also be represented as $\ln\lambda+\gamma$, where $\gamma\sim\operatorname{Gumbel}(-c)$. Defining $\{\gamma(x)\}_{x\in\mathcal{X}}\stackrel{{\scriptstyle\text{i.i.d.}}}{{\sim}}\operatorname{Gumbel}(-c)$ and noting that $g$ is strictly decreasing, applying the function $g$ to equalities in distribution (1) and (2), we obtain:

2 Constructing New Tricks

Given the equality in distribution (1), we can treat the problem of estimating the partition function $Z$ as a parameter estimation problem for the exponential distribution. Applying the function $g(x)=-\ln x-c$ as in the Gumbel trick to obtain a $\operatorname{Gumbel}(-c+\ln Z)$ random variable, and estimating its mean to obtain an unbiased estimator of $\ln Z$, is just one way of inferring information about $Z$.

We consider applying different functions $g$ to (1); particularly those functions $g$ that transform the exponential distribution to another distribution with known mean. As the original exponential distribution has rate $Z$, the transformed distribution will have mean $f(Z)$, where $f$ will in general no longer be the logarithm function. Since we often are interested in estimating various transformations $f(Z)$ of $Z$, this provides us a with a collection of unbiased estimators from which to choose. Moreover, further transforming these estimators yields a collection of (biased) estimators for other transformations of $Z$, including $Z$ itself.

For any $\alpha>0$, applying the function $g(x)=x^{\alpha}$ to an $\operatorname{Exp}(\lambda)$ random variable yields a random variable with the $\operatorname{Weibull}(\lambda^{-\alpha},\alpha^{-1})$ distribution with scale $\lambda^{-\alpha}$ and shape $\alpha^{-1}$, which has mean $\lambda^{-\alpha}\Gamma(1+\alpha)$ and can be also represented as $\lambda^{-\alpha}W$, where $W\sim\operatorname{Weibull}(1,\alpha^{-1})$. Defining $\{W(x)\}_{x\in\mathcal{X}}\stackrel{{\scriptstyle\text{i.i.d.}}}{{\sim}}\operatorname{Weibull}(1,\alpha^{-1})$ and noting that $g$ is increasing, applying $g$ to the equality in distribution (1) gives

Estimating the mean of $\operatorname{Weibull}(Z^{-\alpha},\alpha^{-1})$ yields an unbiased estimator of $Z^{-\alpha}\Gamma(1+\alpha)$. The special case $\alpha=1$ corresponds to the identity function $g(x)=x$; we call the resulting trick the Exponential trick. ∎

where $\{\gamma(x)\}_{x\in\mathcal{X}}\stackrel{{\scriptstyle\text{i.i.d.}}}{{\sim}}\operatorname{Gumbel}(-c)$.

As $\max_{x}\{\phi(x)+\gamma(x)\}\sim\operatorname{Gumbel}(-c+\ln Z)$, we have $e^{-c}\exp(\max_{x}\{\phi(x)+\gamma(x)\})\sim\operatorname{Exp}(Z)$ and the result follows by the assumption relating $f$ and $g$. ∎

Proposition 2 shows that the new tricks can be implemented by solving the same MAP problems $\max_{x}\{\phi(x)+\gamma(x)\}$ as in the Gumbel trick, and then merely passing the solutions through the function $x\mapsto g(e^{-c}\exp(x))$ before averaging them to approximate the expectation.

3 Comparing Tricks

The Delta method (Casella & Berger, 2002) is a simple technique for assessing the asymptotic variance of estimators that are obtained by a differentiable transformation of an estimator with known variance. The last column in Table 1 lists asymptotic variances of corresponding tricks when unbiased estimators of $f(Z)$ are passed through the function $f^{-1}$ to yield (biased, but consistent and non-negative) estimators of $Z$ itself. It is interesting to examine the constants that multiply $Z^{2}$ in some of the obtained asymptotic variance expressions for the different tricks. For example, it can be shown using Gurland’s ratio (Gurland, 1956) that this constant is at least $1$ for the Weibull and Fréchet tricks, which is precisely the value achieved by the Exponential trick (which corresponds to $\alpha=1$). Moreover, the Gumbel trick constant $\pi^{2}/6$ can be shown to be the limit as $\alpha\to 0$ of the Weibull and Fréchet trick constants. In particular, the constant of the Exponential trick is strictly better than that of the standard Gumbel trick: $1<\pi^{2}/6\approx 1.65$. This motivates us to compare the Gumbel and Exponential tricks in more detail.

3.2 Mean squared error (MSE)

For estimators $Y$, their $\operatorname{MSE}(Y)=\operatorname{var}(Y)+\operatorname{bias}(Y)^{2}$ is a commonly used comparison metric. When the Gumbel or Exponential tricks are used to estimate either $Z$ or $\ln Z$, the biases, variances, and MSEs of the estimators can be computed analytically using standard methods (Appendix A).

For example, the unbiased estimator of $\ln Z$ from the Gumbel trick can be turned into a consistent non-negative estimator of $Z$ by exponentiation: $Y=\exp(\frac{1}{M}\sum_{m=1}^{M}X_{m})$, where $X_{1},\ldots,X_{M}\stackrel{{\scriptstyle\text{i.i.d.}}}{{\sim}}\operatorname{Gumbel}(-c+\ln Z)$ are obtained using equation (1’). The bias and variance of $Y$ can be computed using independence and the moment generating functions of the $X_{m}$’s, see Appendix A for details.

Perhaps surprisingly, all estimator properties only depend on the true value of $Z$ and not on the structure of the model (distribution $p$), since the estimators rely only on i.i.d. samples of a $\operatorname{Gumbel}(-c+\ln Z)$ random variable. Figure 1 shows the analytically computed estimator variances and MSEs. For estimating $Z$ itself (left), the Exponential trick outperforms the Gumbel trick in terms of MSE for all sample sizes $M\geq 3$ (for $M\in\{1,2\}$, both estimators have infinite variance and MSE). The ratio of MSEs quickly approaches $\pi^{2}/6$, and in this regime the Exponential trick requires $1-6/\pi^{2}\approx 39\%$ fewer samples than the Gumbel trick to reach the same MSE. Also, for estimating $\ln Z$, (Figure 1, right), the Exponential trick provides a lower MSE estimator for sample sizes $M\geq 2$; only for $M=1$ the Gumbel trick provides a better estimator.

Note that as biases are available analytically, the estimators can be easily debiased (by subtracting their bias). One then obtain estimators with MSEs equal to the variances of the original estimators, shown dashed in Figure 1. The Exponential trick would then always outperform the Gumbel trick when estimating $\ln Z$, even with sample size $M=1$.

For Weibull tricks with $\alpha\not=1$ and Fréchet tricks, we estimated the biases and variances of estimators of $Z$ and $\ln Z$ by constructing $K=100,000$ estimators in each case and evaluating their bias and variance. Figure 2 shows the results for varying $\alpha$ and several sample sizes $M$. We plot the analytically computed value for the Gumbel trick at $\alpha=0$, as we observe that the Weibull trick interpolates between the Gumbel trick and the Exponential trick as $\alpha$ increases from to $1$. We note that the minimum MSE estimator is obtained by choosing a value of $\alpha$ that is close to $1$, i.e. the Exponential trick. This agrees with the finding from Section 2.3.1 that $\alpha=1$ is optimal as $M\to\infty$.

4 Bayesian Perspective

A Bayesian approach exposes two choices when constructing estimators of $Z$, or of its transformations $f(Z)$:

A choice of prior distribution $p_{0}(Z)$, encoding prior beliefs about the value of $Z$ before any observations.

A choice of how to summarize the posterior distribution $p_{M}(Z|X_{1},\ldots,X_{M})$ given $M$ samples.

Taking the Jeffrey’s prior $p_{0}(Z)\propto Z^{-1}$, an improper prior that it is invariant under reparametrization, observing $M$ samples $X_{1},\ldots,X_{M}\stackrel{{\scriptstyle\text{i.i.d.}}}{{\sim}}\operatorname{Exp}(Z)$ yields the posterior:

Recognizing the density of a $\operatorname{Gamma}(M,\sum_{m=1}^{M}X_{m})$ random variable, the posterior mean is

coinciding with the Exponential trick estimator of $Z$.

Low-rank Perturbations

One way of exploiting perturb-and-MAP to yield computational savings is to replace independent perturbations of each configuration’s potential with an approximation. Such approximations are available e.g. in discrete graphical models, where the sampling space $\mathcal{X}$ has a product space structure $\mathcal{X}=\mathcal{X}_{1}\times\cdots\times\mathcal{X}_{n}$, with $\mathcal{X}_{i}$ the state space of the $i$-th variable.

The sum-unary perturbation MAP value is the random variable

where $\{\gamma_{i}(x_{i})\mid x_{i}\in\mathcal{X}_{i},1\leq i\leq n\}\stackrel{{\scriptstyle\text{i.i.d}}}{{\sim}}\operatorname{Gumbel}(-c)$.

This definition involves $|\mathcal{X}_{1}|+\cdots+|\mathcal{X}_{n}|$ i.i.d. Gumbel random variables, rather than $|\mathcal{X}|$. (With $n=1$ this coincides with full-rank perturbations and $U\sim\operatorname{Gumbel}(-c+\ln Z)$.) For $n>2$ the distribution of $U$ is not available analytically. One can similarly define the pairwise (or higher-order) perturbations, where independent Gumbel noise is added to each pairwise (or higher-order) potential.

The following family of upper bounds on $\ln Z$ can be derived from the Fréchet and Weibull tricks.

For any $\alpha\in(-1,0)\cup(0,\infty)$, the upper bound $\ln Z\leq\mathcal{U}(\alpha)$ holds with

(Sketch.) By induction on $n$, with the induction step provided by our Clamping Lemma (Lemma 7) below. ∎

The question remains: When is it advantageous to use a value $\alpha\not=0$ to obtain a tighter bound on $\ln Z$ than the Gumbel trick bound? The next result can provide guidance:

The function $\mathcal{U}(\alpha)$ is differentiable at $\alpha=0$ and the derivative equals

While the variance of $U$ is generally not tractable, in practice one obtains samples from $U$ to estimate the expectation in $\mathcal{U}(\alpha)$ and these samples can be reused to assess $\operatorname{var}(U)$. Interestingly, $\operatorname{var}(U)$ equals $n\pi^{2}/6$ for both the uniform distribution and the distribution concentrated on a single configuration, and in our empirical investigations always $\operatorname{var}(U)\leq n\pi^{2}/6$. Then the derivative at is non-negative and Fréchet tricks provide tighter bounds on $\ln Z$. However, as $\mathcal{U}(\alpha)$ is estimated with samples, the question of estimator variance arises. We investigate the trade-off between tightness of the bound $\ln Z\leq\mathcal{U}(\alpha)$ and the variance incurred in estimating $\mathcal{U}(\alpha)$ empirically in Section 5.3.

2 Clamping

Consider the partial sum-unary perturbation MAP values, where the values of the first $j-1$ variables have been fixed, and only the rest are perturbed:

The following lemma involving the $U_{j}$’s serves three purposes: (I.) it provides the induction step for Proposition 4, (II.) it shows that clamping never hurts partition function estimation with Fréchet and Weibull tricks, and (III.) it will be used to show that a sequential sampler constructed in Section 3.3 below is well-defined.

For any $j\in\{1,\ldots,n\}$ and $(x_{1},\ldots,x_{j-1})\in\mathcal{X}_{1}\times\cdots\times\mathcal{X}_{j-1}$, the following inequality holds with any $\alpha\in(-1,0)\cup(0,\infty)$:

This follows directly from the Fréchet trick ($\alpha\in(-1,0)$) or the Weibull trick ($\alpha>0$) and representing the Fréchet resp. Weibull random variables in terms of Gumbel random variables. See Appendix B.1 for more details. ∎

Clamping never hurts $\ln Z$ estimation using any of the Fréchet or Weibull upper bounds $\mathcal{U}(\alpha)$.

Applying the function $x\mapsto\ln(x)$ to both sides of the Clamping Lemma 7 with $j=1$, the right-hand side equals $\mathcal{U}(\alpha)$, while the left-hand side is the estimate of $\ln Z$ after clamping variable $x_{1}$. ∎

This was shown previously in restricted settings (Hazan et al., 2013; Zhao et al., 2016). Similar results showing that clamping improves partition function estimation have been obtained for the mean field and TRW approximations (Weller & Domke, 2016), and in certain settings for the Bethe approximation (Weller & Jebara, 2014b) and L-Field (Zhao et al., 2016).

3 Sequential Sampling

Hazan et al. (2013) derived a sequential sampling procedure for the Gibbs distribution by exploiting the $\mathcal{U}(0)$ Gumbel trick upper bound on $\ln Z$. In the same spirit, one can derive sequential sampling procedures from the Fréchet and Weibull tricks, leading to the following algorithm.

This algorithm is well-defined if $p_{j}(\text{reject})\geq 0$ for all $j$, which can be shown by canceling terms in the Clamping Lemma 7. We discuss correctness in Appendix B.2. As for the Gumbel sequential sampler of Hazan et al. (2013), the expected number of restarts (and hence the running time) only depend on the quality of the upper bound $(\mathcal{U}(\alpha)-\ln Z)$, and not on the ordering of variables.

4 Lower Bounds on the Partition Function

Similarly as in the Gumbel trick case (Hazan et al., 2013), one can derive lower bounds on $\ln Z$ by perturbing an arbitrary subset $S$ of variables.

Let $\mathcal{X}=\mathcal{X}_{1}\times\cdots\mathcal{X}_{n}$ be a product space and $\phi$ a potential function on $\mathcal{X}$. Let $\alpha\in(-1,0)\cup(0,\infty)$. For any subset $S\subseteq\{1,\ldots,n\}$ of the variables $x_{1},\ldots,x_{n}$ we have $\ln Z\geq$

where $\mathbf{x}_{S}:=\{x_{i}:i\in S\}$ and $\gamma_{S}(\mathbf{x}_{S})\sim\operatorname{Gumbel}(-c)$ independently for each setting of $\mathbf{x}_{S}$.

By averaging $n$ such lower bounds corresponding to singleton sets $S=\{i\}$ together, we obtain a lower bound on $\ln Z$ that involves the average-unary perturbation MAP value

For any $\alpha\in(-1,0)\cup(0,\infty)$, we have the lower bound $\ln Z\geq\mathcal{L}(\alpha)$, where

Advantages of the Gumbel Trick

We have seen how the Gumbel trick can be embedded into a continuous family of tricks, consisting of Fréchet, Exponential, and Weibull tricks. We showed that the new tricks can provide more efficient estimators of the partition function in the full-rank perturbation setting (Section 2), and in the low-rank perturbation setting lead to sequential samplers and new bounds on $\ln Z$, which can be also more efficient, as we investigate in Section 5.3. To balance the discussion of merits of different tricks, in this section we briefly highlight advantages of the Gumbel trick that stem from its simpler analytical form.

First, by consulting Table 1 we see that the function $g(x)=-\ln x-c$ has the property that the variance of the resulting estimator (of $\ln Z$) does not depend on the value of $Z$; the function $g$ is a variance stabilizing transformation for the Exponential distribution.

Third, the additive character of the Gumbel perturbations can also be used to derive a new result relating the error of the lower bound $\mathcal{L}(0)$ and of sampling $x^{**}$ as the configuration achieving the maximum average-unary perturbation value $L$, instead of sampling from the Gibbs distribution $p$:

While we knew from Hazan et al. (2013) that $\ln Z-\mathcal{L}(0)\geq 0$, this is a stronger result showing that the size of the gap is an upper bound on the KL divergence between the approximate sampling distribution of $x^{**}$ and the Gibbs distribution $p$.

Proofs of the new results appear in Appendix B.3 and C.2.

Experiments

We conducted experiments with the following aims:

To show that the higher efficiency of the Exponential trick in the full-rank perturbation setting is useful in practice, we compared it to the Gumbel trick in A* sampling (Maddison et al., 2014) (Section 5.1) and in the large-scale discrete sampling setting of Chen & Ghahramani (2016) (Section 5.2).

To show that non-zero values of $\alpha$ can lead to better estimators of $\ln Z$ in the low-rank perturbation setting as well, we compare the Fréchet and Weibull trick bounds $\mathcal{U}(\alpha)$ to the Gumbel trick bound $\mathcal{U}(0)$ on a common discrete graphical model with different coupling strengths; see Section 5.3.

A* sampling (Maddison et al., 2014) is a sampling algorithm for continuous distributions that perturbs the log-unnormalized density $\phi$ with a continuous generalization of the Gumbel trick, called the Gumbel process, and uses a variant of A* search to find the location of the maximum of the perturbed $\phi$. Returning the location yields an exact sample from the original distribution, as in the discrete Gumbel trick. Moreover, the corresponding maximum value also has the $\operatorname{Gumbel}(-c+\ln Z)$ distribution (Maddison et al., 2014). Our analysis in Section 2.3 tells us that the Exponential trick yields an estimator with lower MSE than the Gumbel trick; we briefly verified this on the Robust Bayesian Regression experiment of Maddison et al. (2014). We constructed estimators of $\ln Z$ from the Gumbel and Exponential tricks (debiased version, see Section 2.3.2), and assessed their variances by constructing each estimator $K=1000$ times and looking at the sample variance. Figure 3(a) shows that the Exponential trick requires up to 40% fewer samples to reach a given MSE.

2 Scalable Partition Function Estimation

Chen & Ghahramani (2016) considered sampling from a discrete distribution of the form $p(x)\propto f_{0}(x)\prod_{s=1}^{S}f_{s}(x)$ when the number of factors $S$ is large relative to the sample space size $|\mathcal{X}|$. Computing i.i.d. Gumbel perturbations $\gamma(x)$ for each $x\in\mathcal{X}$ is then relatively cheap compared to evaluating all potentials $\phi(x)=f_{0}(x)+\sum_{s=1}^{S}\ln f_{s}(x)$. Chen & Ghahramani (2016) observed that each (perturbed) potential can be estimated by subsampling the factors, and potentials that appear unlikely to yield the MAP value can be pruned off from the search early on. The authors formalized the problem as a Multi-armed bandit problem with a finite reward population and derived approximate algorithms for efficiently finding the maximum perturbed potential with a probabilistic guarantee.

While Chen & Ghahramani (2016) considered sampling, by modifying their procedure to return the value of the maximum perturbed potential rather than the argmax (cf equations (1) and (2)), we can estimate the partition function instead. However, the approximate algorithm only guarantees to find the MAP configuration with a probability $1-\delta$. Figure 3(b) shows the results of running the Racing-Normal algorithm of Chen & Ghahramani (2016) on the synthetic dataset considered by the authors with the “very hard” noise setting $\sigma=0.1$. For low error bounds $\delta$ the Exponential trick remained close to optimal, but for a larger error bound the Weibull trick interpolation between the Gumbel and Exponential tricks proved useful to provide an estimator with lower MSE.

3 Low-rank Perturbation Bounds on ln⁡Z𝑍\ln Z

Hazan & Jaakkola (2012) evaluated tightness of the Gumbel trick upper bound $\mathcal{U}(0)\geq\ln Z$ on $10\times 10$ binary spin glass models. We show one can obtain more accurate estimates of $\ln Z$ on such models by choosing $\alpha\not=0$. To account for the fact that in practice an expectation in $\mathcal{U}(\alpha)$ is replaced with a sample average, we treat $\mathcal{U}(\alpha)$ as an estimator of $\ln Z$ with asymptotic bias equal to the bound gap $(\mathcal{U}(\alpha)-\ln Z)$, and estimate its MSE.

Figure 4 shows the MSEs of $\mathcal{U}(\alpha)$ as estimators of $\ln Z$ on $10\times 10$ ($n=100$) binary pairwise grid models with unary potentials sampled uniformly from $$and pairwise potentials from$[0,C]$(attractive models) or from$[-C,C]$(mixed models), for varying coupling strengths$C$. We replaced the expectations in$U(\alpha)$’s with sample averages of size$M=100$, using libDAI (Mooij, 2010) to solve the MAP problems yielding these samples. We constructed each estimator$1000$ times to assess its variance.

Discussion

By casting partition function evaluation as a parameter estimation problem for the exponential distribution, we derived a family of methods of which the Gumbel trick is a special case. These methods can be equivalently seen as (1) perturbing models using different distributions, or as (2) averaging standard Gumbel perturbations in different spaces, allowing implementations with little additional cost.

We showed that in the full-rank perturbation setting, the new Exponential trick provides an estimator with lower MSE, or instead allows using up to 40% fewer samples than the Gumbel trick estimator to reach the same MSE.

In the low-rank perturbation setting, we used our Fréchet, Exponential and Weibull tricks to derive new bounds on $\ln Z$ and sequential samplers for the Gibbs distribution, and showed that these can also behave better than the corresponding Gumbel trick results. However, the optimal trick to use (as specified by $\alpha$) depends on the model, sample size, and MAP solver used (if approximate). Since in practice the dominant computational cost is carried by solving repeated instances of the MAP problem, one can try and assess different values of $\alpha$ on the problem at hand. That said, we believe that investigating when different tricks yield better results is an interesting avenue for future work.

Finally, we balanced the discussion by pointing out that the Gumbel trick has a simpler analytical form which can be exploited to derive more interesting theoretical statements in the low-rank perturbation setting. Beyond existing results, we derived new connections between errors of different procedures using low-rank Gumbel perturbations.

Acknowledgements

The authors thank Tamir Hazan for helpful discussions, and Mark Rowland, Maria Lomeli, and the anonymous reviewers for helpful comments. AW acknowledges support by the Alan Turing Institute under EPSRC grant EP/N510129/1, and by the Leverhulme Trust via the CFI.

References

APPENDIX: Lost Relatives of the Gumbel Trick

Here we provide proofs for the results stated in the main text, together with additional supporting lemmas required for these proofs.

Appendix A Comparison of Gumbel and Exponential tricks

The Gumbel trick yields an unbiased estimator for $\ln Z$, and we can turn it into a consistent estimator of $Z$ by exponentiating it:

Recalling that the moment generating function of a $\operatorname{Gumbel}(\mu)$ distribution is $G(t)=\Gamma(1-t)e^{\mu t}$, we can obtain by using independence of the samples:

Therefore the squared bias, variance and MSE of the estimator $\hat{Z}$ are, respectively:

These formulas hold for $M>2$ where the moment generating functions are defined. For $M=1$ the estimator has infinite bias (and infinite variance), and for $M=2$ it has infinite variance. Figure 1 (left) shows the functional dependence of $\operatorname{MSE}(\hat{Z})$ on the number of samples $M\geq 3$, in units of $Z^{2}$.

The Exponential trick yields an unbiased estimator of $1/Z$, and we can turn it into a consistent estimator of $Z$ by inverting it:

As $X_{1},\ldots,X_{M}$ are independent and exponentially distributed with identical rates $Z$, their sum follows the Gamma distribution with shape $M$ and rate $Z$. Therefore the estimator $\hat{Z}$ can be written as $\hat{Z}=MY$, where $Y\sim\operatorname{InvGamma}(M,Z)$. Recalling the mean and variance of the Inverse-Gamma distribution, we obtain:

Again these formulas hold for $M>2$ where the relevant expectations are defined: for $M=1$ the estimator has infinite bias, and for $M\in\{1,2\}$ it has infinite variance. Figure 1 (left) shows the functional dependence of $\operatorname{MSE}(\hat{Z})$ on the number of samples $M\geq 3$, in units of $Z^{2}$. By inspecting the curves we observe that the Gumbel trick estimator requires roughly 45% more samples to yield the same MSE as the Exponential trick estimator.

A.2 Estimating ln⁡Z𝑍\ln Z

A similar analysis can be performed for estimating $\ln Z$ rather than $Z$. In that case the Gumbel trick estimator of $\ln Z$ is unbiased and has variance (and thus MSE) equal to $\frac{1}{M}\frac{\pi^{2}}{6}$. On the other hand, the Exponential trick estimator is

Again $\sum_{m=1}^{M}X_{m}\sim\operatorname{Gamma}(M,Z)$ and by reference to properties of the Gamma distribution,

where $\psi(\cdot)$ is the digamma function and $\psi_{1}(\cdot)$ is the trigamma function. Note that the estimator can be debiased by subtracting its bias $(\ln M-\psi(M))$. Figure 1 (right) compares the MSE of the Gumbel and Exponential trick estimators of $\ln Z$. We observe that the Gumbel trick estimator performs better only for $M=1$, and even in that case the Exponential trick estimator is better when debiased.

Appendix B Sum-unary perturbations

Recall that sum-unary perturbations refer to the setting where each variable’s unary potentials are perturbed with Gumbel noise, and the perturbed potential of a configuration sums the perturbations from all variables (see Definition 3 in the main text). Using sum-unary perturbations we can derive a family $\mathcal{U}(\alpha)$ of upper bounds on the log partition function (Proposition 4) and construct sequential samplers for the Gibbs distribution (Algorithm 1). Here we provide proofs for the related results stated in Sections 3.1 and 3.2.

For any finite set $\mathcal{Y}$ and any function $h$, we have

This follows from setting up competing exponential clocks with rates $\lambda_{y}=h(y)^{-1/\alpha}$ and then applying the function $g(x)=x^{\alpha}$ as in Example 1” for the case of the Weibull trick. The case of the Fréchet trick is similar, except that $g$ is strictly decreasing for $\alpha\in(-1,0)$, hence the maximization in place of the minimization. ∎

B.1 Upper bounds on the partition function

For any $\alpha\in(-1,0)\cup(0,\infty)$, the upper bound $\ln Z\leq\mathcal{U}(\alpha)$ holds with

We show the result for $\alpha\in(0,\infty)$ using the Weibull trick; the case of $\alpha\in(-1,0)$ can be proved similarly using the Fréchet trick. The idea is to prove by induction on $n$ that $Z^{-\alpha}\geq e^{-\alpha\mathcal{U}(\alpha)}$, so that the claimed result follows by applying the monotonically decreasing function $x\mapsto-\ln(x)/\alpha$.

The base case $n=1$ is the Clamping Lemma 7 below with $j=n=1$. Now assume the claim for $n-1\geq 1$ and for $x_{n}\in\mathcal{X}_{n}$ define

With this definition, the Clamping Lemma with $j=1$ states that $\sum_{x_{1}}\operatorname*{pow}_{-1/\alpha}e^{-\alpha\mathcal{U}_{n-1}(\alpha,x_{1})}\leq\operatorname*{pow}_{-1/\alpha}e^{-\alpha\mathcal{U}(\alpha)}$, so:

as required to complete the inductive step. ∎

The claimed result then follows by the Algebra of Limits, as the contributions of the first two terms cancel. ∎

The function $\mathcal{U}(\alpha)$ is differentiable at $\alpha=0$ and the derivative equals

First we show that $\mathcal{U}(\alpha)$ is differentiable on $(-1,0)\cup(0,\infty)$, and that the limit of the derivative as $\alpha\to 0$ exists and equals $n\pi^{2}/12-\operatorname{var}(U)/2$.

The first term of $\mathcal{U}(\alpha)$ is $n\frac{\ln\Gamma(1+\alpha)}{\alpha}$, which is differentiable for $\alpha\in(-1,0)\cup(0,\infty)$ by the Quotient Rule, and its derivative equals

where $\psi$ is the digamma function (logarithmic derivative of the gamma function). Applying L’Hôpital’s rule we note that

where $\psi^{(1)}$ is the trigamma function (derivative of the digamma function), whose value at $1$ is known to be $\zeta(2)=\pi^{2}/6$, the Riemann zeta function evaluated at $2$.

The second term of $\mathcal{U}(\alpha)$ is constant in $\alpha$. The last term can be written as $K(-\alpha)/(-\alpha)$, where $K$ is the cumulant generating function (logarithm of the moment generating function) of the random variable $U$. The cumulant generating function is differentiable, and by the Quotient rule

where we have used that the second derivative of the cumulant generating function is the variance.

As $\mathcal{U}(\alpha)$ is continuous at by construction, the above implies that it has left and right derivatives at . As the values of these derivatives coincide, the function is differentiable at and the derivative has the stated value. ∎

Recall that for a variable index $j\in\{1,\ldots,n\}$ we also defined partial sum-unary perturbations

which fix the variables $x_{1},\ldots,x_{j-1}$ and perturb the remaining ones.

For any $j\in\{1,\ldots,n\}$ and any fixed partial variable assignment $(x_{1},\ldots,x_{j-1})\in\mathcal{X}_{1}\times\cdots\times\mathcal{X}_{j-1}$, the following inequality holds with any trick parameter $\alpha\in(-1,0)\cup(0,\infty)$:

For $\alpha>0$, from the Weibull trick (Lemma 13), using independence of the perturbations and Jensen’s inequality,

Representing the Weibull random variables in terms of Gumbel random variables using the transformation $W=e^{-(\gamma+c)\alpha}$, where $\gamma\sim\operatorname{Gumbel}(-c)$, and manipulating the obtained expressions yields the claimed result. ∎

B.2 Sequential samplers for the Gibbs distribution

The family of sequential samplers for the Gibbs distribution presented in the main text as Algorithm 1 has the same overall structure as the sequential sampler derived by Hazan et al. (2013) from the Gumbel trick upper bound $\mathcal{U}(0)$, and hence correctness can be argued similarly. Conditioned on accepting the sample, the probability that $\mathbf{x}=(x_{1},\ldots,x_{n})$ is returned is

as required to show that the produced samples follow the Gibbs distribution $p$. Note, however, that in practice one introduces an approximation by replacing expectations with sample averages.

B.3 Relationship between errors of sum-unary Gumbel perturbations

We write $\mathbf{x}^{*}$ for the (random) MAP configuration after sum-unary perturbation of the potential function, i.e.,

The following results links together the errors acquired when using summed unary perturbations to upper bound the log partition function $\ln Z\leq\mathcal{U}(0)$ using the Gumbel trick upper bound by Hazan & Jaakkola (2012), to approximately sample from the Gibbs distribution by using $q_{\text{sum}}$ instead, and to upper bound the entropy of the approximate distribution $q_{\text{sum}}$ using the bound due to Maji et al. (2014).

Writing $p$ for the Gibbs distribution, we have

By conditioning on the maximizing configuration $\mathbf{x}^{*}$, we can rewrite the Gumbel trick upper bound $\mathcal{U}(0)$ as follows:

At the same time, the KL divergence between $q_{\text{sum}}$ and the Gibbs distribution $p$ generally expands as

Adding the two equations together and rearranging yields the claimed result. ∎

Appendix C Averaged unary perturbations

In the main text we stated the following two lower bounds on the log partition function $\ln Z$.

Let $\alpha\in(-1,0)\cup(0,\infty)$. For any subset $S\subseteq\{1,\ldots,n\}$ of the variables $x_{1},\ldots,x_{n}$ we have $\ln Z\geq$

where $\mathbf{x}_{S}:=\{x_{i}:i\in S\}$ and $\gamma_{S}(\mathbf{x}_{S})\sim\operatorname{Gumbel}(-c)$ independently for each setting of $\mathbf{x}_{S}$.

Let $\bar{S}:=\{1,\ldots,n\}\setminus S$. First we handle the case $\alpha>0$. We have trivially that

Expressing the Weibull random variable $W(\mathbf{x}_{S})$ as $e^{-\alpha(\gamma_{S}(\mathbf{x}_{S})+c)}$ with $\gamma_{S}(\mathbf{x}_{S})\sim\operatorname{Gumbel}(-c)$, the right-hand side can be simplified as follows:

Taking the logarithm and dividing by $-\alpha<0$ yields the claimed result for positive $\alpha$. For $\alpha\in(-1,0)$ we proceed similarly, obtaining that

where $F(\mathbf{x}(S))\sim\operatorname{Fr\acute{e}chet}(1,-\alpha^{-1})$. Representing these random variables as $e^{-\alpha(\gamma_{S}(\mathbf{x}_{S})+c)}$ with $\gamma_{S}(\mathbf{x}_{S})\sim\operatorname{Gumbel}(-c)$, simplifying as in the previous case and finally dividing the inequality by $-\alpha>0$ yields the claimed result for $\alpha\in(-1,0)$. ∎

For any $\alpha\in(-1,0)\cup(0,\infty)$, we have the lower bound $\ln Z\geq\mathcal{L}(\alpha)$, where

Applying Proposition 9 $n$ times with all singleton sets $S=\{i\}$ and averaging the obtained lower bounds yields

where the first equality used the fact that the perturbations $\gamma_{i}(x_{i})$ are mutually independent for different indices $i$ to replace the product of expectations with the expectation of the product. The claimed result follows by applying Jensen’s inequality to swap the summation and the convex $\max_{\mathbf{x}}$ function, noting that the inequality works out the right way for both positive and negative $\alpha$. ∎

Jensen’s inequality can be used to relate the general lower bound $\mathcal{L}(\alpha)$ to the Gumbel trick lower bound $\mathcal{L}(0)$, showing that the former cannot be arbitrarily worse than the latter:

For all $\alpha\in(-1,0)$, the lower bound $\mathcal{L}(\alpha)$ on $\ln Z$ satisfies

Apply Jensen’s inequality with the convex function $x\mapsto e^{-n\alpha}$ to the last term in the definition of $\mathcal{L}(\alpha)$, noting that the inequality works out the stated way for $\alpha<0$. ∎

Note that $\frac{\ln\Gamma(1+\alpha)}{\alpha}+c\leq 0$ for $\alpha\in(-1,0)$ so this result does not imply that the Fréchet lower bounds are tighter than the Gumbel lower bound $\mathcal{L}(0)$; it merely says that they cannot be arbitrarily worse than $\mathcal{L}(0)$.

C.2 Relationship between errors of averaged-unary Gumbel perturbations

In this section we write $\mathbf{x}^{*}$ for the (random) MAP configuration after average-unary perturbation of the potential function, i.e.,

We show that the gap of this Gumbel trick lower bound on $\ln Z$ upper bounds the KL divergence between the approximate distribution $q_{\text{avg}}$ and the Gibbs distribution $p$. To this end, we first need an entropy bound for $q_{\text{avg}}$ analogous to Theorem 1 of (Maji et al., 2014).

The entropy of $q_{\text{avg}}$ can be lower bounded using expected values of max-perturbations as follows:

Theorem 1 of (Maji et al., 2014) and this Theorem 15 differ in three aspects: (1) the former is an upper bound and the latter is a lower bound, (2) the former sums the expectations while the latter averages them, and (3) the distributions $q_{\text{sum}}$ and $q_{\text{avg}}$ of $\mathbf{x}^{*}$ in the two theorems are different.

By the duality relation between negative entropy and the log partition function (Wainwright & Jordan, 2008), the entropy $H(q_{\text{avg}})$ of the unary-avg perturb-max distribution $q_{\text{avg}}$ can be expressed as

where the variable $\varphi$ ranges over all potential functions on $\mathcal{X}$, and $Z_{\varphi}=\sum_{\mathbf{x}\in\mathcal{X}}\exp\varphi(\mathbf{x})$. Applying the Gumbel trick lower bound on the log partition function gives

Proposition 16 in Appendix D shows that $\mathcal{L}_{\varphi}(0)$ is a convex function of $\varphi$. The expression $-\sum_{\mathbf{x}\in\mathcal{X}}q(\mathbf{x})\varphi(\mathbf{x})$ is a linear function of $\varphi$, so also convex, and thus as a sum of two convex functions, the quantity $\mathcal{L}_{\varphi}(0)-\sum_{\mathbf{x}\in\mathcal{X}}q(\mathbf{x})\varphi(\mathbf{x})$ within the infimum is a convex function of $\varphi$. Moreover, Proposition 17 in Appendix D tells us that the partial derivatives can be computed as

where $q_{\varphi}(\mathbf{x})$ is the unary-avg perturb-max distribution associated with the potential function $\varphi$. Proposition 18 in Appendix D confirms that these partial derivatives are continuous, so we observe that as a function of $\varphi$, the expression $\mathcal{L}_{\varphi}(0)-\sum_{\mathbf{x}\in\mathcal{X}}q_{\text{avg}}(\mathbf{x})\varphi(\mathbf{x})$ is a convex function with continuous partial derivatives, so it is a differentiable convex function. This is sufficient to establish that the point $\varphi=\phi$ is a global minimum of this function (Wright & Nocedal, 1999). Hence

where we conditioned on the maximizing configuration $\mathbf{x}^{*}$ when expanding $\mathcal{L}_{\phi}(0)$. ∎

This proof proceeded in the same way as the proof of Maji et al. (2014) for the upper bound, except that establishing the minimizing configuration of the infimum is a non-trivial step that is actually required in this case. The second revision of (Hazan et al., 2016) computes the derivative of $\mathcal{U}_{\varphi}(0)-\sum_{\mathbf{x}\in\mathcal{X}}q_{\text{sum}}(\mathbf{x})\varphi(\mathbf{x})$, which is similar to our $\mathcal{L}_{\varphi}(0)-\sum_{\mathbf{x}\in\mathcal{X}}q_{\text{avg}}(\mathbf{x})\varphi(\mathbf{x})$, by differentiating under the expectation.

Equipped with Theorem 15, we can now show a link between the approximation “errors” of the averaged-unary perturbation MAP configuration distribution $q_{\text{avg}}$ (to the Gibbs distribution $p$) and estimate $\mathcal{L}(0)$ (to $\ln Z$).

Let $p$ be the Gibbs distribution on $\mathcal{X}$. Then

While we knew from Hazan et al. (2013) that $\ln Z-\mathcal{L}(0)\geq 0$ (i.e. that $\mathcal{L}(0)$ is a lower bound on $\ln Z$), this is a stronger result showing that the size of the gap is an upper bound on the KL divergence between the average-unary perturbation MAP distribution $q_{\text{avg}}$ and the Gibbs distribution $p$.

The Kullback-Leibler divergence in question expands as

From the proof of Theorem 15 we know that $H(q_{\text{avg}})\geq\mathcal{L}(0)-\sum_{\mathbf{x}\in\mathcal{X}}q_{\text{avg}}(\mathbf{x})\phi(\mathbf{x})$, so

Appendix D Technical results

In this section we write $\mathcal{L}(\phi)$ instead of $\mathcal{L}_{\phi}(0)$ for the Gumbel trick lower bound on $\ln Z$ associated with the potential function $\phi$, see equation (3).

The Gumbel trick lower bound $\mathcal{L}(\phi)$, viewed as a function of the potentials $\phi$, is convex.

Convexity can be proved directly from definition. Let $\phi_{1}$ and $\phi_{2}$ be two arbitrary potential functions on a discrete product space $\mathcal{X}$, and let $\lambda\in$. Then

where we have used convexity of the $\max$ function to obtain the inequality, and linearity of expectation to arrive at the final equality. ∎

This convexity proof goes through for other (low-dimensional) perturbations as well, e.g. it also works for $\mathcal{U}_{\phi}(0)$.

The Gumbel trick lower bound $\mathcal{L}(\phi)$, viewed as a function of the potentials $\phi$, has partial derivatives

where $q_{\phi}$ is the probability mass function of the average-unary perturbation MAP configuration’s distribution associated with the potential function $\phi$.

Now let’s condition on the size of the gap $G$ between the maximum and the runner-up:

Let’s examine all four terms on the right-hand side one by one:

which proves the stated claim directly from definition of a partial derivative. ∎

The probability mass function $q_{\phi}$ of the average-unary perturbation MAP configuration’s distribution associated with a potential function $\phi$ is continuous in $\phi$.

For any $\mathbf{x}^{*}\in\mathcal{X}$ we have from definition

The results above show that the Gumbel trick lower bound $\mathcal{L}(\phi)$, viewed as a function of the potentials $\phi$, is convex and has continuous partial derivatives.