In the MINE paper, why is $\hat{G}_B$ biased, and how does the exponential moving average reduce the bias?

Question

While reading the Mutual Information Neural Estimation (MINE) paper [1] I came across section 3.2 Correcting the bias from the stochastic gradients. The proposed method requires the computation of the gradient

$$\hat{G}_B = \mathbb{E}_B[\nabla_{\theta}T_{\theta}] - \frac{\mathbb{E_B}[\nabla_{\theta}T_{\theta}e^{T_{\theta}}]}{\mathbb{E}_B[e^{T_{\theta}}]},$$

where $\mathbb{E}_B$ denotes the expectation operation w.r.t. a minibatch $B$, and $T_{\theta}$ is a neural network parameterized by $\theta$. The authors claim that this gradient estimation is biased and that can be reduced by simply performing an exponential moving average filtering.

Can someone give me a hint to understand these two points:

1. Why is $\hat{G}_B$ biased, and
2. How does the exponential moving average reduce the bias?

Answer 1

The lower bound in MINE is as follows:

$$\widehat{I(X;Z)}_n = \sup_{\theta\in\Theta} \mathbb{E}_{\mathbb{P}_{XZ}^{(n)}}[T_\theta] - \log{\mathbb{E}_{\mathbb{P}_X^{(n)} \otimes \hat{\mathbb{P}}_Z^{(n)}}[e^{T_\theta}]}$$

Here $\mathbb{\hat{P}^{(n)}}$ denotes the empirical distribution that we get from n i.i.d samples of $\mathbb{P}.$

Note that in the above equation, the first term is calculated from the joint distribution while the second term from the marginals of $X$ and $Z$. In the implementation of MINE, these statistics are calculated over the data from a minibatch. The marginal distribution is obtained by shuffling the values of Z (or X) along the batch dimension. Hence, in this case, the gradient is as follows.

$$\hat{G}_B = \mathbb{E}_B[\nabla_{\theta}T_{\theta}] - \frac{\mathbb{E_B}[\nabla_{\theta}T_{\theta}e^{T_{\theta}}]}{\mathbb{E}_B[e^{T_{\theta}}]},$$

1. As mentioned, the expectation over the marginals is not calculated over the true marginal distribution (i.e. over the entire dataset) but from the shuffled samples in the minibatch. Hence, the above gradient $G_B$ is biased.
2. When we maintain an exponential moving average of $\mathbb{E}_B[e^{T_{\theta}}]$, we kind of incorporate statistics from outside the current minibatch also (i.e. over the entire dataset). This is an attempt to get an approximation of the true marginal estimate. The denominator term in the gradient allows for such a computationally inexpensive bias reduction trick.