Calculating mutual information between layer outputs and targets in a neural network

Question

I've seen in several papers that it is possible to calculate the mutual information between a layer's outputs and the desired outputs. For example:

Source: https://www.ncbi.nlm.nih.gov/pmc/articles/PMC7517266/#B9-entropy-22-00727

How exactly is this calculated? I'm trying to calculate this quantity and track how it evolves during training for a simple feed-forward neural network used in regression. The mutual information should increase with training.

I am currently using this Python package for information theory calculations: https://github.com/gregversteeg/NPEET

This is my dataset:

X = torch.FloatTensor(500, 1).uniform_(-3, 3)
Y_no_noise = X**2 -X - 1
noise = torch.randn(X.size()) 
Y = Y_no_noise + 0.2*noise

And this is my network:

class MLP(nn.Module):
    '''
    Multi-layer perceptron for non-linear regression.
    '''
    def __init__(self, nInput, nHidden, nOutput):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(nInput, nHidden), # Input layer
            nn.ReLU(), 
            nn.Linear(nHidden, nHidden), # Hidden layer
            nn.ReLU(),
            nn.Linear(nHidden, nHidden), # Hidden layer 
            nn.ReLU(),
            nn.Linear(nHidden, nHidden), # Hidden layer
            nn.ReLU(), 
            nn.Linear(nHidden, nOutput) # Output layer
        )

    def forward(self, x):
        return(self.layers(x))

model = MLP(1, 8, 1)

I assemble the activations from a layer into a $n \times d$ array, where $n$ is the number of training examples (500) and $d$ is the number of neurons.

The targets/desired outputs are in a $n \times 1$ array.

I then use the mutual information function, mi() provided in the package, passing in the two arrays as arguments.

However, the mutual information graph I obtain doesn't seem to have any meaningful pattern/trend as I would expect:

The presence of negative mutual information values here is also concerning.

I would be very grateful for ideas of where I might be going wrong in my calculation and how I can get the correct trend of mutual information increasing.

EDIT: For reference, the entirety of my code is as follows:

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.models as models
import torchvision.datasets as datasets
import torch.optim as optim
import torchvision.transforms as transforms
import torch.utils.data as data
import matplotlib.pyplot as plt
import copy
plt.style.use("ggplot")
from pyhessian import hessian
import sys
import seaborn as sns
from scipy.stats import gaussian_kde
from scipy.integrate import quad
from scipy.integrate import nquad
from sklearn.decomposition import PCA
from sklearn.feature_selection import mutual_info_regression
from npeet import entropy_estimators as ee

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

X = torch.FloatTensor(500, 1).uniform_(-3, 3)
Y_no_noise = X**2 -X - 1
noise = torch.randn(X.size()) 
Y = Y_no_noise + 0.2*noise

X,Y= X.to(device), Y.to(device)
dataset = data.TensorDataset(X, Y) 

train_size = int(0.7*len(dataset))
test_size = len(dataset) - train_size

train_dataset, test_dataset = data.random_split(dataset, [train_size, test_size])

x_train = [x[0] for x in train_dataset]
x_train = torch.stack(x_train)
x_test = [x[0] for x in test_dataset]   
x_test = torch.stack(x_test)

y_train = [x[1] for x in train_dataset]
y_train = torch.stack(y_train)
y_test = [x[1] for x in test_dataset]
y_test = torch.stack(y_test)

x_train = x_train.to(device)
y_train = y_train.to(device)
x_test = x_test.to(device)
y_test = y_test.to(device)

y_train_for_mi = y_train.cpu().detach().numpy()
y_train_for_mi = y_train_for_mi.tolist()
#print(len(y_train_for_mi))

nInput  = 1
nHidden = 8
nOutput = 1

class MLPcondensed(nn.Module):
    '''
    Multi-layer perceptron for non-linear regression.
    '''
    def __init__(self, nInput, nHidden, nOutput):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(nInput, nHidden), # Input layer
            nn.ReLU(), 
            nn.Linear(nHidden, nHidden), # Hidden layer
            nn.ReLU(),
            nn.Linear(nHidden, nHidden), # Hidden layer 
            nn.ReLU(),
            nn.Linear(nHidden, nHidden), # Hidden layer
            nn.ReLU(), 
            nn.Linear(nHidden, nOutput) # Output layer
        )

    def forward(self, x):
        return(self.layers(x))

model = MLPcondensed(nInput, nHidden, nOutput)
model.to(device)

# Select layer for which mutual information is to be calculated
layer_index = 8 # Goes 0, 2, 4, 6, 8

loss_fn = nn.MSELoss().to(device)
optimiser = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
epochs = 500
batch_size = 32

train_loader = data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

tl = []
mutual_info_list = []

for epoch in range(epochs):
    train_loss = 0
    #print(f"Epoch {epoch+1} of {epochs}")

    activations_overall = []

    for batch, (X, y) in enumerate(train_loader):
        model.train() # Training mode
        X, y = X.to(device), y.to(device) # Move data to GPU
        predictions = model(X) 
        activations = model.layers[0:layer_index+2](X)
        #print(activations)
        individual_tensors = torch.split(activations, 1, dim=0)
        for tensor in individual_tensors:
            individual_array = tensor.squeeze().detach().cpu().numpy()
            activations_overall.append(individual_array)
        loss = loss_fn(predictions, y)
        train_loss += loss # Accumulate the loss over the batch
        optimiser.zero_grad() 
        loss.backward() 
        optimiser.step() 
    
    
    tl.append(train_loss/len(train_loader))

    activations_overall = np.array(activations_overall)
    activations_overall = activations_overall.tolist()

    mutual_info = ee.mi(y_train_for_mi, activations_overall)
    mutual_info_list.append(mutual_info)
 
final_loss = [tl[-1]]
final_loss_numpy = torch.stack(final_loss).cpu().detach().numpy()[0]
#print(f"Final training loss: {final_loss_numpy}")
train_loss_curve = torch.stack(tl)
plt.plot(train_loss_curve.detach().cpu().numpy(), label="Train loss")
plt.xlabel("Epoch")
plt.ylabel("Loss")
plt.legend()
plt.show()

plt.plot(mutual_info_list, label="Mutual information")
plt.xlabel("Epoch")
plt.ylabel("Mutual information")
plt.legend()
plt.show()

EDIT: Response to nbro's comment

The mutual information between two variables X and Y measures the information that they share. It measures how much knowing one of these variables reduces uncertainty about the other. It is a non-negative, symmetric quantity.

I wish to use treat the activations of a layer as a single multivariate variable, and the desired output as the other variable. The mutual information between these two, I believe, can be used to measure how learning progresses in a layer because as training progresses, the activations should tell you more about the desired outputs. So, I expect the mutual information to steadily increase (and then maybe plateau) with training duration/epochs. This idea is quite widely discussed in the literature, especially in the context of the information plane.

• https://www.semanticscholar.org/reader/267980e417f1d01a897e87fa409f64e2a76b96cd

• https://towardsdatascience.com/information-theory-of-neural-networks-ad4053f8e177#:~:text=Mutual%20Information%20along%20the%20Training%20Phase&text=With%20successive%20layers%20the%20mutual,output%20is%20low%20as%20well.&text=As%20we%20train%20the%20Neural,of%20information%20about%20the%20output.

When I say "desired outputs", I mean the targets given my input. My regression dataset consists of 1D inputs and 1D outputs (the "desired outputs").

Answer 1

NPEET uses a K Nearest Neighbors estimation for mutual information and they note in section 4.3 of their provided pdf doc that there is a chance for negative values in the estimated mutual information. They note their own example has this issue. They mention their shuffle.test function may help resolve the negatives.

https://github.com/gregversteeg/NPEET/blob/master/npeet_doc.pdf

In section 6.4 of the doc they detail further on shuffle.test for potentially resolving negative estimates.

The pdf doc notes that the mutual information follows the calculation from the following linked paper in equation 8

https://journals.aps.org/pre/abstract/10.1103/PhysRevE.69.066138

I see you are purposefully comparing the activations to the desired output, rather than the predicted output. As a sanity check it may be useful to see if the output mutual information behaves as expected.

The docs say these measurements default to base 2 thus are in bits and your mutual information estimates all < .15 bits as absolute values. I suspect this noise in your measurements is due to the estimation technique and there is less than .15 MI in the layer you compare the targets to.

I believe that'd be the weights not the activations unless you mean the output of the activation functions. I have not run your code and do not know exactly which layer or part of the network you assign as the activations in the code with line activations = model.layers[0:layer_index+2](X)

Your example has relatively small dimensions so I doubt it'd be a sample data amount too small vs measure being estimated, though it may be worth trying more samples just to check.