On Mutual Information Maximization for Representation Learning

Abstract: Many recent methods for unsupervised or self-supervised representation learning train feature extractors by maximizing an estimate of the mutual information (MI) between different views of the data. This comes with several immediate problems: For example, MI is notoriously hard to estimate, and using it as an objective for representation learning may lead to highly entangled representations due to its invariance under arbitrary invertible transformations. Nevertheless, these methods have been repeatedly shown to excel in practice. In this paper we argue, and provide empirical evidence, that the success of these methods cannot be attributed to the properties of MI alone, and that they strongly depend on the inductive bias in both the choice of feature extractor architectures and the parametrization of the employed MI estimators. Finally, we establish a connection to deep metric learning and argue that this interpretation may be a plausible explanation for the success of the recently introduced methods.

Reviews and author responses:

Thank you for your insightful comments and thorough review. While indeed our main focus was to start a discussion and point out widespread misconceptions with the current understanding of the methods motivated by information maximization, we believe that there are many exciting approaches to address some of these issues which we elaborate below.

[A] Measures of information which capture the geometry
On the theoretical side, we believe that the question of developing new notions of information suitable for representation learning should receive more attention. While mutual information has appealing theoretical properties,  it is clearly not sufficient for this task — it is hard to estimate, invariant to bijections and can result in suboptimal representations which do not correlate with downstream performance. Therefore, a new notion of information should account for both the amount of information stored in a representation and the geometry of the induced space necessary for good performance on downstream tasks. For example, the representation space (induced by the encoder) is typically endowed with some geometry, which is not taken into account when computing the mutual information. However, precisely this geometry is often critical for building good classifiers and it should be accounted for explicitly. Thus, building notions information which rely on statistical divergences with this property should be investigated. For example, using the Wasserstein divergence leads to promising results in representation learning [1].

[B] A holistic view
We believe that any theory on measuring information for representation learning built on critics should explicitly take into account the function families one uses (e.g. that of the critic and estimator). Most importantly, we would expect some natural trade-offs between the amount of information that can be stored against how hard it is to extract it in the downstream tasks as a function of the architectural choices. While the distribution of downstream tasks is typically assumed unknown in representation learning, it might be possible to rely on weaker assumptions such as a family of invariances relevant for the downstream tasks.  Moreover, it seems that in the literature (i) the critics that are used to measure the information, (ii) the encoders, and (iii) the downstream models, are all mostly chosen independently of each other. Our empirical results show that the downstream performance depends on the intricate balance between these choices and we believe that one should co-design them. This holistic view is currently underexplored and due to the lack of any theory or extensive studies to guide the practitioners.

[C] Part with the probabilistic interpretation, focus on systematic investigations into design decisions that matter (e.g. negative sampling)
On the practical side, we believe that the link to metric learning could lead to new methods, that break away from the goal of estimating MI and place more weight on the aspects that have a stronger effect on the performance such as the negative sampling strategy. An example where the metric learning perspective led to similar methods as the MI view is presented in [2]: They developed a multi-view representation learning approach for video similar to CMC [3], but seemingly without relying on the MI mental model to motivate their design choices. 

[1] “Wasserstein dependency measure for representation learning”. Sherjil Ozair, Corey Lynch, Yoshua Bengio, Aaron van den Oord, Sergey Levine, and Pierre Sermanet. (2019)
[2] “Time-Contrastive Networks: Self-Supervised Learning from Video”. Pierre Sermanet, Corey Lynch, Yevgen Chebotar, Jasmine Hsu, Eric Jang, Stefan Schaal, Sergey Levine. (2018)
[3] “Contrastive Multiview Coding”. Yonglong Tian, Dilip Krishnan, Phillip Isola. (2019)

Thank you for your detailed review, please find our reply below:

- That is indeed what we argue — for the same estimate of the MI some estimators result in more useful representations. While there are clear benefits of having tighter and computationally tractable bounds on MI, we find compelling evidence that representation learning would benefit more from a thorough investigation into inductive biases and alternative interpretations of the loss functions implied by these estimators.

- This is indeed the key question and the experiments in Sections 3.1 and 3.2 provide some insight: For a fixed inductive bias on the critic architecture, the InfoNCE loss (“triplet”) outperforms the other estimator. For a fixed estimator, the inductive bias on the critic architecture plays a role, although a counter-intuitive one: stronger critics can lead to worse downstream performance. Unfortunately, the bias in the encoder (e.g. using a CNN vs an MLP) can also be of paramount importance in practice, independently of the other design choices. Hence, the downstream performance hinges on a delicate balance between these components, which additionally could be data-dependant as one can see in Appendix G, where we on a different dataset (CIFAR-10) we observe a smaller improvement by replacing $I_{\textrm{NWJ}}$ with $I_{\textrm{NCE}}$. 


Thank you for your astute review, please find our reply below:

1. With lower-case letters ($x$, $y$) we denote the outcomes / realizations, while with upper case letters ($X$, $Y$) the variables themselves. We have updated the manuscript to clarify this point and avoid any confusion (footnote on page 1).

2. You are indeed correct that the quantity inside the bracket in (3) is estimated using Monte Carlo estimation, and in practice we estimate (3) by averaging over multiple batches of samples. The expectation is necessary for the lower bound to hold. Please see Appendix D for a detailed derivation.

3. In the proof of Proposition 1 it is stated that they are “Markov equivalent”, meaning that they have the same conditional independencies. In particular, $X_1$ and $X_2$ are conditionally independent given $X$, which suffices for the proof. We updated the manuscript with a brief clarification of “Markov equivalence” (page 13, second paragraph of the proof).


In this paper, the authors studied the usage of the mutual information maximization principle in representation learning. They argue that the bias in the estimation of lower bound of MI may loosen the connection between InfoMax principle and representation learning. Specifically, they figure out the following phenomenon by experiments.
    1. Large MI is not predictive of downstream performance.
    2. Higher capacity critics can lead to worse downstream performance.
    3. Encoder architecture can be more important than the specific estimator.
Finally, the authors make a connection to deep metric learning.

Recently, using the principles of mutual information in deep learning has been very hot. However, a lot of papers just use the concept of mutual information without deep thinking (why use MI? why not other metrics? what is the estimation bias and variance for MI?). Also, most paper do not have theoretical guarantees. In this paper, the authors think deeper about the InfoMax principle, and point out some weakness about the connection between InfoMax principle and the quality of representation learning. I think this paper can trigger some deep and calm thinking in this area.

However, I think this paper is more critical then constructive. It is good to point out some weakness of the InfoMax principle, but what should we do next? How should we fix the InfoMax principle, or mutual information estimator, to make it more robustly useful for representation learning? Or does it mean that InfoMax principle is not a good tool for representation learning? I agree that this questions are out of the scope of this paper, but I would like to see the authors share more constructive thoughts in the discussion area, to make this paper more constructive.

Overall, I would like to weakly accept this paper. I think this paper worth a strong accept if some constructive ideas are provided in the rebuttal. 

The paper addresses a question on whether mutual information (MI) based models for representation learning succeed primarily thanks to the MI maximization. The motivation of the work comes from the fact that although MI is known to be problematic in treatment, it has been successfully applied in a number of recent works in computer vision and natural language processing. The paper conducts a series of experiments that constitute a convincing evidence for a weak connection between the InfoMax principle and these practical successes by showing that maximizing established lower bounds on MI are not predictive of the downstream performance and that contrary to the theory higher capacity instantiations of the critics of MI may result in worse downstream performance of learned representations. The paper concludes that there is a considerable inductive bias in the architectural choices inside MI models that are beneficial for downstream tasks and note that at least one of the lower bounds on MI can be interpreted as a triplet loss connecting it with a metric learning approach.
I consider this paper to be a considerable contribution to the understanding of what underlies the performance of unsupervised representation learning with MI maximization and provides a good discussion and analogous insights from parallel works and a number of possible directions to explore. Although I still have a couple of questions addressing which will help to understand the paper.
-  In experiment 3.1, when using RealNVP and maximizing the lower bound of MI may it be the case that the representations are learned to benefit the downstream task because of the form of the lower bounds? In other words, the lower bound is possibly not very tight and its maximization has a side effect of weights adjusting to yield simple representations useful for a linear classifier?
-  It would be interesting to see which factor contributes more to the performance: I_NCE being a triplet loss or an inductive bias in the design choice?

This paper gives a nice interpretation why recent works that are based on variational lower bounds of mutual information can demonstrate promising empirical results, where they argue that the success depends on "the inductive biasin both the choice of feature extractor architectures and the parametrization of theemployed MI estimators." To support this argument, they carefully design a series convincing experiments which are stated in full in Section 3. Moreover they show some connection to metric learning. 

I have confusions thought about the writing. These should be addressed before the acceptance.

1. For many equations, if the X, Y, are random variables, they should be capitalized. For example, in equation (1), should p(x,y) be written as p(X,Y)? If I'm right, please correct it every where in the paper.

2. In equation (3), should the symbol E be there? I think it shouldn't. Since in real implement, Monte Carlo estimation is used and the mini-batch samples are selected in accordance to one positive sample and the rest negative samples, E shouldn't be there. But I see E in all-most every papers such as Poole and Oord's papers. To equation (5), the connection to metric learning, there is no E. If I'm correct, please correct in the paper including the appendix.

3. For the proof of proposition 1, it is stated X_1 <-- X --> X_2 is equivalent to X_1 --> X --> X_2, why? I don't have Cover's book on hand, so I couldn't make sure.


Meta-review: This paper exams the role of mutual information (MI) estimation in representation learning. Through experiments, they show that the large MI is not predictive of downstream performance, and the empirical success of  methods like InfoMax may be more attributed to the inductive bias in  the choice of architectures of discriminators, rather than accurate MI estimation. The work is well appreciated by the reviewers. It forms a strong contribution and may motivate subsequent works in the field. 


Decision: accept-poster