Free Energy Evaluation Using Marginalized Annealed Importance Sampling

For the sequence of the distributions in equation (3), we introduce the sequence of their marginal distributions as

where $P_{k}(\bm{v})$ is the marginal distribution of $P_{k}(\bm{x})$, i.e.,

where

is the energy function of the marginal distribution. From the definition, the partition functions of $P_{k}(\bm{x})$ and $P_{k}(\bm{v})$ are the same.

Based on the sequence of the marginal distributions in equation (19), in a similar manner to equation (7), the annealing process from the initial state $\bm{v}^{(1)}$ to the final state $\bm{v}^{(K)}$ can be defined as

where $\bm{V}:=\{\bm{v}^{(k)}\mid k=1,2,\ldots,K\}$ and $\tau_{k}(\bm{v}^{\prime}\mid\bm{v})$ is the transition probability on the marginal distribution in equation (20). The transition probability satisfies the following condition:

Using almost the same derivation to obtain equation (13), we derive

where $\Lambda(\bm{V})$ is the importance weight of mAIS defined as

where

In mAIS, the free energy can be evaluated using a technique similar to the derivation of equation (16). The sequence of the sample points, $\mathbf{V}:=\{\mathbf{v}^{(1)},\mathbf{v}^{(2)},\ldots,\mathbf{v}^{(K)}\}$, is generated based on the ancestral sampling on $\tau(\bm{V})$:

By conducting this ancestral sampling process $N$ times, $N$ independent sample sequences $S_{\tau}:=\{\mathbf{V}_{\mu}\mid\mu=1,2,\ldots,N\}$ are obtained, and subsequently, using the $N$ sample sequences, equation (24) is approximated as

Therefore, the free energy is evaluated as

Similar to $Z_{\mathrm{AIS}}(S_{T})$, $Z_{\mathrm{mAIS}}(S_{\tau})$ acts as an unbiased estimator for the true partition function because $\mathbb{E}_{\tau}[Z_{\mathrm{mAIS}}(S_{\tau})]=Z$, where $\mathbb{E}_{\tau}[\cdots]$ denotes the expectation of the distribution over $S_{\tau}=\{\mathbf{V}_{\mu}\}$, i.e.,

Similar to equation (18), based on Jensen’s inequality,

can be obtained; therefore, $F_{\mathrm{mAIS}}(S_{\tau})$ is not an unbiased estimator for the true free energy as is the case with $F_{\mathrm{AIS}}(S_{T})$.

In the aforementioned discussions, we have considered mAIS based on the sequence of $P_{k}(\bm{v})$ (we refer to this mAIS as “mAIS-V”). An opposite mAIS (referred to as “mAIS-H”), which is based on the sequence of

where

can be constructed in the same manner. However, we do not need to consider mAIS-H separately because the difference between the methods only lies in the way the variables are labelled, i.e., mAIS-V is identified with mAIS-H by exchanging the labels of variable sets. Therefore, we consider mAIS-V as mAIS in this paper.

Theoretically, mAIS can be applied to any model. However, whether it is practical or not strongly depend on the structures of $\{P_{k}(\bm{x})\}$ because mAIS requires the marginal operation, as expressed in equation (20). mAIS can be efficiently applied to a bipartite-type MRF (i.e., an MRF defined on a bipartite undirected graph), which will be discussed in Section III.3.

Suppose that we have a model that mAIS can be efficiently applied to. Therefore, the free energy of the model can be evaluated based on the two methods, i.e., $F_{\mathrm{AIS}}(S_{T})$ and $F_{\mathrm{mAIS}}(S_{\tau})$. Intuitively, mAIS seems to be better because the entire sample space of the variables is reduced through the marginalization, which is similar to the concept of Rao–Blackwellization Liu (2001).

Here, we briefly explain Rao–Blackwellization. Suppose that we wish to evaluate the expectation of $f(\bm{x}_{t})$ over a distribution $Q(\bm{x})$: $m_{t}:=\sum_{\bm{x}}f(\bm{x}_{t})Q(\bm{x})$, where $\bm{x}_{t}\subseteq\bm{x}$. Based on a simple sampling approximation (or Monte Carlo integration), it is approximated by

where $\{\mathbf{x}^{(\nu)}\mid\nu=1,2,\ldots,M\}$ is the sample set generated from $Q(\bm{x})$, and $\mathbf{x}_{t}^{(\nu)}$ is the corresponding subset of $\mathbf{x}^{(\nu)}$. Using the decomposition $Q(\bm{x})=Q(\bm{x}_{a}\mid\bm{x}_{b})Q(\bm{x}_{b})$, $m_{t}$ can be transformed as

where $m_{t}^{\dagger}(\bm{x}_{b})$ is the conditional expectation expressed as

Therefore, based on equation (35), $m_{t}$ is approximated by

where $\mathbf{x}_{b}^{(\nu)}$ is the corresponding subset of $\mathbf{x}^{(\nu)}$. The Rao–Blackwell theorem guarantees that the estimator in equation (36) is more effective than the estimator in equation (34) from the perspective of the variance. The transformation of equation (35) is called Rao–Blackwellization. The applications of Rao–Blackwellization are widely investigated; for example, its application to MRFs, spatial Monte Carlo integration, has been developed Yasuda (2015); Yasuda and Uchizawa (2021); Yasuda and Sekimoto (2021).

The relation between equations (34) and (36) looks similar to that between AIS and mAIS. Thus, we can expect mAIS to be more effective. In fact this expectation can be justified under a certain condition, which is discussed in the next section.

First, we compare the statistical efficiencies of the two estimators for the partition function, i.e., $Z_{\mathrm{AIS}}(S_{T})$ and $Z_{\mathrm{mAIS}}(S_{\tau})$. The variance of $Z_{\mathrm{AIS}}(S_{T})$ is

where $\mathbb{V}_{T}[A]:=\mathbb{E}_{T}[A^{2}]-\mathbb{E}_{T}[A]^{2}$. The variance of $Z_{\mathrm{mAIS}}(S_{\tau})$ is

where $\mathbb{V}_{\tau}[A]:=\mathbb{E}_{\tau}[A^{2}]-\mathbb{E}_{\tau}[A]^{2}$. As discussed in Sections II and III.1, both estimators are unbiased; therefore, the estimator with smaller variance is more effective.

Assume that the transition probability of AIS is expressed as

where $\tau_{k}(\bm{v}^{\prime}\mid\bm{v})$ is the transition probability of mAIS satisfying equation (23), and

is the conditional distribution on $P_{k}(\bm{x})$. According to equation (39), the state transition from $\bm{x}$ to $\bm{x}^{\prime}$ is performed as follows: first, the state of $\bm{v}$ is updated to $\bm{v}^{\prime}$ according to $\tau_{k}(\bm{v}^{\prime}\mid\bm{v})$; subsequently, the state of $\bm{h}$ is updated to $\bm{h}^{\prime}$ according to $P_{k}(\bm{h}^{\prime}\mid\bm{v}^{\prime})$ (see figure 1). The use of the transition probability of equation (39) is accepted in AIS because it satisfies the condition of equation (6) as follows:

Based on this assumption, the following theorem can be obtained.

Based on this theorem, it is ensured that $Z_{\mathrm{mAIS}}(S_{\tau})$ is statistically more efficient. The proof of this theorem is described in Appendix A.1. As mentioned in the final part of Appendix A.1, the claim of this theorem is essentially the same as that of the Rao–Blackwell theorem.

Next, the statistical efficiencies of the two estimators for the free energy, i.e., $F_{\mathrm{AIS}}(S_{T})$ and $F_{\mathrm{mAIS}}(S_{\tau})$, are compared. As expressed in equations (18) and (31), $\mathbb{E}_{T}[F_{\mathrm{AIS}}(S_{T})]$ and $\mathbb{E}_{\tau}[F_{\mathrm{mAIS}}(S_{\tau})]$ are upper bounds of the true free energy. Based on the assumption of equation (39), the following theorem can be obtained.

Based on this theorem, although $F_{\mathrm{AIS}}(S_{T})$ and $F_{\mathrm{mAIS}}(S_{\tau})$ are biased, the bias of $F_{\mathrm{mAIS}}(S_{\tau})$ is ensured to be smaller. The proof of this theorem is presented in Appendix A.2. The aforementioned two theorems do not depend on the details of the design of $E_{k}(\bm{x})$ and the values of $N$ and $K$; only the condition in equation (39) is required.

We have proved that mAIS is statistically more efficient than AIS with regard to the partition function and free energy when the condition of equation (39) is satisfied. The remaining issue is determining whether this condition is satisfied in practical situations. This condition significantly limits the modeling of $T_{k}(\bm{x}^{\prime}\mid\bm{x})$; e.g., the standard (synchronous or asynchronous) Gibbs sampling on $P_{k}(\bm{x})$ does not satisfy the condition because it uses the states of all variables to obtain the subsequent states. As explained in Section III.3, the state transition according to equation (39) can be natural and practical when $P_{k}(\bm{x})$ is a bipartite-type MRF.

In Theorems 1 and 2, the condition of equation (39) is not a necessary condition but a sufficient condition. This implies the possibility of the relaxation of the condition. However, it is an open problem.

The proposed mAIS is applicable when the marginal operation in equation (20) is feasible. Owing to this restriction, the range of application of mAIS is limited. As discussed below, mAIS is practical on bipartite-type MRFs, which include important applications such as RBM and DBM (a DBM can be observed as a bipartite-type MRF Gao and Duan (2017)). Moreover, a square-grid-type MRF is regarded as a bipartite-type MRF (see Fig. 2).

Suppose that $P_{k}(\bm{x})$ is a bipartite-type MRF and that $\bm{v}$ and $\bm{h}$ are the variables of each layers. In this case, the energy function $E_{k}(\bm{x})$ can be expressed as

where $\phi_{i}^{(k)}$ and $\psi_{i,j}^{(k)}$ are one- and two-variable functions determined based on the design of $\{P_{k}(\bm{x})\}$; the marginal operation in equation (20) is feasible and leads to

On the bipartite-type MRF, the layer-wise blocked Gibbs sampling based on the conditional distributions $P_{k}(\bm{h}\mid\bm{v})$ and $P_{k}(\bm{v}\mid\bm{h})$ can be easily implemented, where $P_{k}(\bm{h}\mid\bm{v})$ is the conditional distribution presented in equation (40) and

because $v_{i}$s are conditionally independent from each other in $P_{k}(\bm{v}\mid\bm{h})$ and $h_{j}$s are also conditionally independent from each other in $P_{k}(\bm{h}\mid\bm{v})$, i.e.,

Based on the layer-wise blocked Gibbs sampling, the transition probability of mAIS can be modeled as

which satisfies the condition in equation (23) as follows:

This transition probability can be considered as a collapsed Gibbs sampling Liu (1994). From equations (39) and (44), when the transition probability of AIS is modeled by

the theoretical results mentioned in Section III.2 (Theorems 1 and 2) are applicable. Approximating the expectations in equations (44) and (45) using the sampling approximation with one sample point leads to the widely used sampling procedure, i.e., the blocked Gibbs sampling, on bipartite-type MRFs (see figure 3).

The theoretical results obtained in section III.2 revealed the qualitatively effectiveness of mAIS. We investigate the quantitatively effectiveness of mAIS using numerical experiments. In the experiments, $b_{i}$s and $c_{j}$s were independently drawn from a uniform distribution in $[-0.001,+0.001]$, and $w_{i,j}$s are independently drawn from a normal distribution whose mean is zero and variance is $1/n$; the sizes of the sample set and $\bm{v}$ were fixed to $N=1000$ and $|V|=20$. The blocked Gibbs sampling shown in figure 3 was used for the state transitions of AIS and mAIS.

Figure 4 shows the absolute percentage error (APE) between the true free energy, $f:=F/n$, and its approximation, $f_{\mathrm{app}}:=F_{\mathrm{app}}/n$, obtained using AIS or mAIS, against the inverse temperature $1/T$, where the APE is defined as

In this experiment, the size of $\bm{h}$ is fixed to $|H|=40$. We observe that the APEs obtained using mAIS are always lower than those obtained using AIS, as supported by Theorem 1. mAIS is particularly effective in the high-temperature region.

Next, we compare the true free energy $f$ with the trial averages of its approximations, $\mathbb{E}_{\mathrm{trial}}[f_{\mathrm{app}}]$, where $\mathbb{E}_{\mathrm{trial}}[\cdots]$ was estimated based on the average of 30 trials. The trial average can be regarded as the approximation of $\mathbb{E}_{T}[\cdots]$ in AIS or $\mathbb{E}_{\tau}[\cdots]$ in mAIS. The results are listed in table 1. In this experiment, the size of $\bm{h}$ was fixed to $|H|=40$. The estimates of AIS and mAIS are higher than the true free energy and the estimates of mAIS are lower than those of AIS, as supported by Theorem 2.

We investigate the relative accuracy of mAIS compared with AIS for fraction $\alpha:=|H|/|V|$. The experiments in section IV.1 correspond to the cases of $\alpha=2$. Intuitively, mAIS is more efficient for a larger $\alpha$ because, as $\alpha$ increases, the fraction of the size of remaining variables through marginalization to $n$ reduces (in other words, the dimension of space evaluated by the sampling approximation relatively shrinks). However, the theoretical results obtained in section III.2 do not directly support this assumption. Thus, we check the relative accuracy for several $\alpha$ values using numerical experiments on the RBM. In the experiments, the relative accuracy is measured by the ratio of the APEs obtained using AIS and mAIS, i.e.,

The relative accuracy increases as $r$ increases. In the experiment, the sizes of the sample set and $\bm{v}$ were fixed to $N=1000$ and $|V|=20$; the parameter setting of the RBM was the same as that in section IV.1, and the blocked Gibbs sampling was used for the state transitions of AIS and mAIS. The distributions of $\ln r$ for $1/T=0.5$, $1$, and $2$ are shown in figures 5–7, respectively. Each distribution in these figures was created based on the kernel density estimation Bishop (2006) using the results obtained from 4000 experiments, in which the Gaussian kernel with a bandwidth (or smoothing parameter) of 0.25 was used. In all cases, the distributions transit to the positive direction with an increase in $\alpha$, which implies that the relative accuracy monotonically improves as $\alpha$ increases. This result implies that we should use mAIS-V when $\alpha>1$ and mAIS-H when $\alpha<1$.

The experiments in sections IV.1 and IV.2 were conducted for relatively small $K$ values to emphasize the performance difference between AIS and mAIS under less-than-ideal conditions. However, in practice, $K$ is set to a sufficient large value to obtain precise estimations. To obtain precise estimations, the value of $K$ should be set to a value larger than the mixing time (or relaxation time). The mixing time tends to increase as the complexity of distribution increases. In this section, the convergence property of APE in equation (47) on a wider range of $K$ is demonstrated using numerical experiments.

The detailed discussions for the mixing time in RBM have been provided from both inference and learning perspectives Roussel et al. (2021); Decelle et al. (2021). Roussel et al. proved that the standard blocked Gibbs sampling shown in figure 3 is not efficient in RBMs having high clustered distributions from the perspective of the mixing time, and they proposed a more efficient sampling method by combining the blocked Gibbs sampling and Metropolis–Hastings (MH) method illustrated in figure 8 Roussel et al. (2021). Roussel’s sampling method can be employed as the transition $\tau_{k}(\bm{v}^{\prime}\mid\bm{v})$ in our framework. For the experiments, we used two different RBMs: the RBM with Gaussian interactions which is the same model used in sections IV.1 and IV.2, and the RBM with Hopfield-type interactions in which $w_{i,j}=\xi_{i}^{(j)}/\sqrt{|V|}$, where $\xi_{i}^{(j)}\in\{-1,+1\}$ was determined uniformly at random. The Hopfield-type interactions can exhibit a clustered distribution Roussel et al. (2021). In both RBMs, the bias parameters, $b_{i}$s and $c_{j}$s, were independently drawn from a uniform distribution in $[-0.001,+0.001]$, and $|V|=20$ and $|H|=40$ were fixed. The size of the sample set was fixed to $N=10000$.

Figures 9 and 10 show the behavior of the APE for the increase in $K$. In the figures, “AIS” and “mAIS” indicate the results obtained using the standard blocked Gibbs sampling, while “AIS+MH” and “mAIS+MH” indicate the results obtained using Roussel’s sampling method. mAIS exhibits faster convergence than AIS in all experiments, and Roussel’s sampling method improves the convergence properties of both AIS and mAIS, although the improvement is very small in figure 10(b). mAIS+MH achieved the best performance the best in all experiments. The convergence speed of mAIS seems to be higher than that of AIS+MH; this implies that the improvement by mAIS is more effective than that by Roussel’s sampling method. Finally, we comment on the computational efficiency of the four methods: AIS, AIS+MH, mAIS, and mAIS+MH. The computational cost of the four methods have the same order; they are $O(NK|V||H|)$. However, mAIS is remarkably faster than AIS+MH with regard to the CPU time (around 10 times faster in our implementation). The MH procedure is the bottleneck in the Roussel’s sampling method. Thus, mAIS+MH is as slow as AIS+MH.