Estimate Mutual Information

Estimate the mutual information between two continuous variables by training a neural critic with the InfoNCE loss.

For continuous data you cannot count frequencies, but you can train a classifier. Take a batch of B pairs (x_i, y_i). A critic gives a score s[i, j] for every pairing of x_i with y_j. The InfoNCE loss asks the critic to pick the true partner y_i for each x_i, out of all B candidates:

loss  = cross_entropy(s, labels = [0, 1, ..., B - 1])
I(X; Y) >= log(B) - loss

The better the critic separates true pairs from shuffled ones, the tighter the bound. A simple and fast critic is separable: two small MLPs g and h, with s = g(x) @ h(y).T, one B x B matrix product for the whole batch.

Train the critic with Adam on random mini batches. Then, under torch.no_grad(), average log(B) - loss over several batches, and return the estimate in nats as a Python float.

It is graded against correlated Gaussians, where the true value is known: I = -0.5 * d * log(1 - rho ** 2). Your estimate must be within 0.1 nats.

Function signature

def estimate_mi(x, y):

Example

Input

x = torch.randn(4096, 1)
y = 0.8 * x + 0.6 * torch.randn(4096, 1)
estimate_mi(x, y)

Output

0.51    # about -0.5 * log(1 - 0.64) = 0.511

Reasoning

Each y holds most of the information in its x, so the critic can tell the true pair from the others. The bound reaches close to the true value of 0.511 nats.

Constraints and notes

• x and y are float32 of shape (n, d), with n = 4096 and d <= 2.
• The bound can never be more than log(B). Use a batch of a few hundred.
• Training must take a few seconds on a CPU. Set a seed so it is repeatable.