This repository contains code for training simple fully connected and convolutional neural networks on the MNIST dataset. Additionally, it provides tools for calculating mutual information using different methods to analyze the layers' information retention during training.
-
Clone the repository:
git clone https://github.com/aidinattar/info-bottleneck.git cd your-repo-name
-
Install the required packages:
pip install -r requirements.txt
To train a simple fully connected network on the MNIST dataset:
-
Navigate to the directory:
cd info_bottleneck
-
Run the training script:
python main.py --model mlp
To train a simple convolutional neural network on the MNIST dataset:
-
Navigate to the directory:
cd info_bottleneck
-
Run the training script:
python main.py --model cnn
The mutual information is calculated using different methods (binning
, kde
, kraskov
). The mutual information calculation is integrated within the NetworkTrainer
class and can be specified during initialization.
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
# MNIST dataset and DataLoader
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
train_dataset = datasets.MNIST(root='./data', train=True, transform=transform, download=True)
train_loader = DataLoader(dataset=train_dataset, batch_size=64, shuffle=True)
val_dataset = datasets.MNIST(root='./data', train=False, transform=transform, download=True)
val_loader = DataLoader(dataset=val_dataset, batch_size=64, shuffle=False)
# Initialize the trainer and train the model
trainer = SimpleCNNTrainer(train_loader, val_loader, activation='relu', optimizer=optim.Adam, epochs=10, device='cuda')
trainer.train()
This code is based on the Information Bottleneck principle for neural networks. For more information, refer to the following resources:
- On the Information Bottleneck Theory of Deep Learning
- Information Bottleneck for Deep Learning
- Mutual Information Neural Estimation
- Information Bottleneck
The current implementation focuses on fully connected and convolutional neural networks. Future work could include:
- Extending the mutual information analysis to other network architectures such as recurrent neural networks (RNNs) or transformers.
- Applying the mutual information analysis to more complex datasets beyond MNIST.
- Investigating the impact of different mutual information estimation methods on the performance and interpretability of neural networks.
This project is licensed under the MIT License. See the LICENSE file for details.
Special thanks to the authors of the referenced papers and the open-source community for providing valuable resources and tools that made this project possible.