Mutual Information (MI) plays an important role in representation learning. However, MI is unfortunately intractable in continuous and high-dimensional settings. Recent advances establish tractable and scalable MI estimators to discover useful representation. However, most of the existing methods are not capable of providing an accurate estimation of MI with low-variance when the MI is large. We argue that directly estimating the gradients of MI is more appealing for representation learning than estimating MI in itself. To this end, we propose the Mutual Information Gradient Estimator (MIGE) for representation learning based on the score estimation of implicit distributions. MIGE exhibits a tight and smooth gradient estimation of MI in the high-dimensional and large-MI settings. We expand the applications of MIGE in both unsupervised learning of deep representations based on InfoMax and the Information Bottleneck method. Experimental results have indicated the significant performance improvement in learning useful representation.
This paper proposes MIGE---a novel estimator of the mutual information (MI) gradient, based on estimating the score function of an implicit distribution. To this end, the authors employ the spectral Stein gradient estimator (SSGE) and propose its scalable version based on random projections of the original input. The theoretical advantages of the method are presented using a toy experiment with correlated Gaussian random variables, where both the mutual information and its gradient can be computed analytically. In this setting, MIGE provides gradient estimates that are less biased and smoother than baselines. The method is also evaluated on two more complicated tasks: unsupervised representation learning on Cifar-10 and CIfar-100 via DeepInfoMax (DIM) and classification on MNIST with Information Bottleneck (IB), where MIGE outperforms all baselines by a significant margin.
I recommend ACCEPTing this paper. This work discusses a vital problem and proposes a novel, well-motivated, principled and very performant solution; additionally, it demonstrates the broad applicability of the introduced method. While the proposed technique consists of previously known building blocks (spectral Stein gradient estimator and random projections), it is cleverly applied in a novel context of estimating MI gradients.
While the paper is solid, I believe that it could be improved in the following ways. Firstly, I would like to see 1) more extensive and 2) larger-scale evaluations. In the DIM experiment, 1) would correspond to trying the DIM(L) approach, which maximises patch-wise MI. In fact, I strongly recommend including this experiment as it corresponds to and could improve current state-of-the-art. If it turns out that MIGE does not work well on DIM(L), then this would correspond to a serious issue with the method. In this experiment, 1) would also include providing other metrics for learned representations. It would be much more convincing to include estimates of true mutual information (e.g. InfoNCE bound evaluated with a large number of samples [1]) and showing that MIGE can attain higher values than baselines. 2) would correspond to evaluation on bigger datasets: (tiny) ImageNet and STL-10 dataset. Also, the toy experiment would benefit from a higher-dimensional setting (e.g. d=256 to d=1024), since these are often used in practice. 
Secondly, the paper is sloppily-written, which quite a few grammar and stylistic mistakes (e.g. sentence in sec 3, paragraph 2: “we assume obtain to…”, which starts with a lower-case letter and doesn’t make sense). Finally, the paper would benefit from the following clarifications: 1) explain what is the Nystr\:om method, 2) provide either a proof or a citation for eq (19); also the error bound for SSGE should be provided for the paper to be self-contained, 3) explain the difference between q_\psi and p_\psi, which seem to be used interchangeably.
Additional remarks:
Sec 2.2, 2), “streamlining” is unclear
Circumstances 2 and 3 can be quite easily derived from circumstance 1; also they are not evaluated empirically; it would be nice to have experiments for them, and they can be moved to the appendix in case of lack of space
Eq (19) while nice, seem to bear no significance for the proposed method and the rest of the paper; consider removing it
Section 4.2 paragraph 2: “shrinking” for different layers wasn’t mentioned before, and is not immediately clear what it means; the reader needs to be intimately familiar with the DIM paper to understand.
Section 4.3 mentions “threshold” for stein gradient estimator, which was not mentioned before. Please explain what it is.
Equations (8-10) are just simple derivations and are not necessary; it would be enough to provide Eq (10).
The authors talk about MINE, which optimizes the InfoNCE bound [1], which is also used in DIM and CPC [2]. I strongly encourage the authors to cite [1] and [2] and mention them in the related works. Additionally, it would be clear if Figure 1 and related references and description used “InfoNCE” instead of “MINE” as the name of the method since InfoNCE is an estimator and MINE is just a particular implementation of the method.
 
[1] Poole et. al., “On variational bounds of mutual information”, ICML 2019.
[2] van den Oord et. al, “Representation Learning with Contrastive Predictive Coding“, arXiv 2018.