# Building up to DIM (Notes)
In this note, I trace the evolution of research into using mutual information as a representation learning objective.

Here I cover the following papers:
1. Mutual Information Neural Estimation (MINE). ([Belghazi et al, 2018](https://arxiv.org/pdf/1801.04062.pdf))
2. Learning Deep Representations By Mutual Information Estimation And Maximization. ([Hjelm et al, 2019](https://arxiv.org/pdf/1808.06670.pdf))
3. Learning Representations by Maximizing Mutual
Information Across Views. ([Bachman et al, 2019](https://arxiv.org/pdf/1906.00910.pdf))

---
# Mutual Information Neural Estimation (MINE). ([Belghazi et al, 2018](https://arxiv.org/pdf/1801.04062.pdf))

MINE introduces a mutual information (MI) estimator that is:
1. Differentiable.
2. Scalable in sample size and dimensionality.
3. Consistent.

The estimator operates using two key ideas:

1. MI can also be written as the KL-divergence between a joint and product of marginals.
2. Use the dual formulation of the KL-divergence to make estimation tractable.

## Method Overview
#### Generating Joint and Product Of Marginals
Let $X, Y$ represent two random variables. Define the dataset of possible pairings as $D=\{(X_1, Y_1), ..., (X_n, Y_n)\}$ which also defines the joint probability function $P(X, Y)$.

If $X, Y$ are independent, then we can write the joint as the product of marginals $P(X, Y) = P(A)P(B)$.

MINE proposes to obtain the product of marginals by shuffling a batch of data along the batch axis such that $(X_{i \ne j}, Y_{j})$. By shuffling we enforce independence and thus generate the valid product of marginals.

In summary:

\begin{align}
\text{joint = batch} &= P(X, Y) \\
\text{shuffle(batch(X))} &= P(X) \\
\text{shuffle(batch(Y))} &= P(Y) \\
\text{product of marginals} &= P(X)P(Y)
\end{align}

#### Mutual Information as KL divergence.

Mutual information (MI) measures the reduction in uncertainty given the other variable.

\begin{align}
I(X, Y) = H(Y) - H(Y|X) = H(X) - H(X|Y)
\end{align}

Mutual information can be written as the KL divergence between the joint $P_{XY}$ and product of marginals $P_X P_Y$:

\begin{align}
I(X, Y) = D_{KL}(P_{XY} || P_X P_Y)
\end{align}

where $D_{KL}$ is defined by:
\begin{align}
D_KL(P || Q) = \mathbb{E}_{P} \left [log \frac{P}{Q} \right]
\end{align}

#### Dual formulations
MIME proposes a dual representation of the KL divergence for two main reasons:
1. Estimating MI is difficult because we have to integrate over a log density.
2. Formulating it this way also allows us to use adversarial training.

Let $T$ be a function. MINE proposes to use two known dual formulations:

1. *The Donsker-Vardhan representation*.

\begin{align}
D_{KL}(P||Q) = \sup_{T: \Omega -\rightarrow \mathbb{R}} \mathbb{E_P}[T] - log(\mathbb{E_Q[e^T]})
\tag{1}
\end{align}

- where the supremum is taken over all functions T such that the two expectations are finite.

2. *The f-divergence representation*.
\begin{align}
D_{KL}(\mathbb{P} || \mathbb{Q}) \ge \sup_{T \in F} \mathbb{E_P}[T] - \mathbb{E_Q}[e^{T-1}]
\tag{2}
\end{align}

Both of these duals offer tight bounds. However, the DV representation is stronger because for any fixed T, the right hand-side of (1) is larger than the right hand-side of (2).

#### MINE Formulation.
To obtain the proposed MINE formulation we do three things:
1. Use the DV dual form (1).
2. Parametrize T with a neural network $T_{\theta}$.
3. Use the dataset pairs $(x^{(1)}, z^{(1)})$ in place of the random variables $(X, Y)$.

\begin{align}
I(X, Z) = \sup_{\theta \in \Theta}\frac{1}{N} \sum_{i=1}^{N}T_{\theta}(x^{(i)}, z^{(j)}) - log \left[ \frac{1}{N} \sum_{i=1}^{N} \exp{{T_{\theta}(x^{(i)}, \tilde{z}^{(i)})}} \right]
\tag{3}
\end{align}

Where $x^{(i)}, y^{(j)} \sim \mathbb{P}_{XZ}$ (joint) and $\tilde{z}^{(i)} \sim \mathbb{P}_Z$ (marginal, ie: shuffled labels).

**Intuition**: We train a classifier to discriminate between samples from the joint vs samples from the marginal.

#### Biased gradients.
In a mini-batch setting, the SGD gradients of MINE are biased.

Naive SGD application to (3) leads to the biased gradient:

\begin{align}
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}}]}
\end{align}

because in the second term, the expectations are over the samples of a minibatch, which is a biased estimate of the full batch gradient.

(ie: we use only one 1 y to estimate the denominator when we shuffle the batch. Instead, take a moving average over a number of steps)

**Solution**: Replace denominator estimate by using an exponential moving average.

## Results
In the rest of the paper they do the following experiments:
1. Compare how accurate their MI estimate is against calculating it directly.
- Recover multivariate gaussians.
3. Use it to reduce mode collapse in GANs.
- Regularize generator loss with estimate of MI.
4. Improve reconstruction in autoencoders and adversarial architectures.
5. Information bottleneck (ie: Build representation of data that only holds some information).

---
# Learning Deep Representations By Mutual Information Estimation And Maximization ([Hjelm et al, 2019](https://arxiv.org/pdf/1808.06670.pdf)).

In this work, DIM, the authors build on MINE to build representations which have high MI with the input. Their goal is to learn useful representations for downstream tasks.

#### Representations
Representation is posterior over explanatory factors for an observed input.

Properties the authors suggest representations should have:
- Separability.
- Independence in the dimensions.

### Method Setup
Let $f$ represent an encoder (CNN in this case). For two images $I_1, I_2$ we generate two sets of respective feature maps $M_1 = f(I_1), M_2=f(I_2)$.

Let $Y$ represent a global feature vector generated by applying an MLP $h$ to $M_1$

$$
Y_1 = h(f(I_1))
$$


### DIM Objective
The DIM objective uses the DV dual representation of the KL divergence from eq (1) to maximize the MI between the representation vector $Y$ and a feature map $M$.

For convenience, we rewrite the DV representation here:

\begin{align}
D_{KL}(P||Q) = \sup_{T: \Omega -\rightarrow \mathbb{R}} \mathbb{E_P}[T] - log(\mathbb{E_Q[e^T]})
\tag{1}
\end{align}

For the purposes of DIM, we let the function $T$ be defined by a discriminator CNN $T_{\omega}$ that discriminates between samples of the joint $P_{XY}$ and product of marginals $P_XP_Y$.

The discriminator is trained to produce a high score when a positive pair $(M_1, Y_1)$ is seen and a low score when a negative pair $(M_2, Y_1)$ is seen.

We obtain the full discriminator loss by replacing $T=T_\omega$ in eq (1),

\begin{align}
D_{KL}(P||Q) = \sup_{T: \Omega -\rightarrow \mathbb{R}} \mathbb{E}_{P_{XY}}[T_\omega(M_1, Y_1)] - log(\mathbb{E}_{P_XP_Y}[exp(T_\omega(M_2, Y_1)])
\tag{4}
\end{align}

The above equation is the core DIM formulation

### Illustration


### Experiments

#### Linear Classification
This is a proxy for the MI of the representation with linear separability.

1. Pretrain (100% data unlabeled)
2. Freeze Weights.
3. Train an SVM on top of representations (100% data labeled).


#### Non-linear Classification:
Proxy for MI between representation and labels.

1. Pretrain (100% data unlabeled)
2. Freeze Weights.
3. Train 200 unit MLP on top of representations (100% data labeled).




#### Semi-Supervised Learning
How well it can leverage labeled + unlabeled data.

1. Pretrain (100% data unlabeled)
2. DO NOT FREEZE NETWORK.
3. Unclear what finetuning model used, likely a 200 unit MLP.
4. Finetune (unfreeze weights) on labeled training split.


---
# Learning Representations by Maximizing Mutual Information Across Views. ([Bachman et al, 2019](https://arxiv.org/pdf/1906.00910.pdf))

In this work (AMDIM), the authors maximimize mutual information between multiple augmented versions of the same image. The features they use come from 2 different feature maps taken from 2 random layers in the CNN.

## Method
### Core idea 1: Multiple Views Of Shared Context
AMDIM generates instances of a class by adding random augmentations to the ***same*** image.

Let $A$ define an augmentation such as sheer, rotation, etc... Given an input $x$ we can generate an augmented version of that input $x^1 = A(x)$.

Given two augmentations of the same image $x^1 \sim A(x), x^2 \sim A(x)$ this new objective maximizes mutual information $(MI)$ between these two samples.

#### Augmentations used:
First, randomly flip x. Then choose 2 of these (to get two samples from the same image):

1. random resized crop
2. random jitter in color space
3. random conversion to grayscale

### Core idea 2: Spatial Feature Feature Extraction.

For each augmentation we compute two sets of features. The first one is a global context vector computer as follows:

**Global Context Vector:**
Let $E$ represent an encoder (CNN in this case) which outputs a feature map. Let $f_1$ represent a global feature vector generated by applying an MLP $h$ to $E(x^1)$

$$
f_1 = h(E(x^1))
$$

**Local features:**
The second set of features (coming from a second augmentation $x^2$) is another feature map $f_m$ from any other layer $m$ in the CNN.

$$
f_m = E(x^2)
$$

### Core idea 3: Negative sampling
The objective for AMDIM is a NCE loss using 10,000 distractors.

The positive pair $f_1(x^1), f_7(x^2)$ comes from the global context features $f_1$ for augmentation $x^1$ and the local set of features for a second augmentation $x^2$.

The 10,000 negative pairs come from using other images $X'$ different from $x$, ie: $(f_1(x^1), f_n(x'))$.

For each of these pairs we calculate a score $\Phi(f_1, f_n)$ and use that score in the standard log-softmax function.

\begin{align}
\mathcal{L}_{\Phi}(f_1, f_7, X') = - log \frac{exp(\Phi(f_1(x^1), f_7(x^2)))}{\sum_{x' \in X'} exp(\Phi(f_1(x^1), f_7(x')))}
\end{align}

The $\Phi$ function in this case is a dot product.

#### Core idea 4: Coordinate prediction
The learning problem is set up so that the network actually predicts the spatial coordinate of a feature vector from a convolutional feature map.

It appears that $f_7(x^2)_{ij}$ is actually a single vector at spatial location $ij$.

## Technical details
The two feature maps can come from different parts in the CNN. In AMDIM they use the following pairs: 1-5, 1-7, 5-5.

#### Two tricks to stabilize training
1. Add weighted regularization term that penalizes loss. (?).
2. Clip $[-c, c]$, after calculating score but before softmax ($c = tanh(s/c)$).

## Illustration

## Experimental setup
#### Model
Resnet (unclear how many layers, but likely 50).

#### Hardware
Trained on 4-8 v100 gpus. Mention that other SOTA models can't be trained on smaller compute.

#### Datasets
1. Imagenet
- 128 x 128
3. Places
- 128 x 128
4. Cifar-10
5. CIFAR-100.

## Tasks
### Linear separability
They use the standard task of finetuning a linear model to test for separability.

1. Pretrain on dataset using no labels
2. Freeze encoder
3. Train linear and MLP classifier on top.
4. Report accuracy of these classifiers (looks like they use all training examples labeled).

**MLP:** Looks like the MLP is 200 units, but doesn't say anywhere in the paper.

#### CIFAR results (linear + non-linear)
Results for linear and MLP finetuning on pretrained model.

#### Imagenet results (linear only)
Results finetuning only using a linear classifier:

***Note: The latest CPC paper has 61.0 AMDIM hasn't beat it yet.***

#### Hyperparam details
ndf = encoder feature dimension
nrkhs = output dim of embedding function (the MLP context vector).

CF10, 100, STL ndf=256, nrkhs=2048, batch=300.

# General takeaways

## Good ideas we can try
1. Deriving the positive example from different transformations of an image.
2. Comparing features across layers. Our current loss uses features from the same layer. Positive and negative features could come from different layers.

## General questions
1. Why is metric learning a better approach than mutual information maximization?
###### - side note, why is it called metric learning?
3. Can we maximize MI for positive term and minimize $MI$ for negative term? ie:

$$
\mathcal{L} = max(0, MI(x, x^+) - MI(x, x^-) + m )
$$

## Advantages of DDT over DIM
1. No need for discriminator.
- but in this later version there isn't a discriminator.

2. No need to force the representation into a single context vector.
3. Allows mixed-label training (semi-sup learning). Although it doesn't quite work right now.