Code Package for: Mutual Information Can Be Reliably Estimated in High-Dimensional Data that Admit Low-Dimensional Representations

This folder contains all code used to generate the results and figures in the paper. The structure below summarizes the purpose of each file and its relation to the figures.

-------------------------------------------------------
Estimator and Model Implementation
-------------------------------------------------------

1. utils.py  
   - Core utility functions used across experiments.  
   - Includes data generation utilities, teacher model definitions, MLP architectures, and miscellaneous helpers.  
   - Based on prior work from [1,2].

2. models.py  
   - Contains implementations of the DSIB (deterministic) and DVSIB (variational/probabilistic) estimator models.  
   - Implements the core encoder-decoder logic used for the VSIB family.
   - Works for different classes of critics.
   - Based on [3].

3. estimators.py  
   - Implements all estimator variants evaluated in the paper: InfoNCE, SMILE, etc.  
   - These estimators are compatible with the training and evaluation pipelines.  
   - Based on [1,2].

-------------------------------------------------------
Data Preparation
-------------------------------------------------------

4. data_generation.ipynb  
   - Notebook to generate the Noisy MNIST and teacher-transformed synthetic datasets.  
   - Loads MNIST, applies augmentations, and generates matched (X, Y) pairs efficiently.

-------------------------------------------------------
Plotting and Analysis
-------------------------------------------------------

5. analysis_n_plotting.ipynb  
   - Notebook to analyze raw results and generate all plots in the paper.  
   - Includes logic for max-test smoothing, extrapolation to γ→0, weighted least squares, etc.  
   - All results from scripts below are loaded and processed here.

-------------------------------------------------------
Experimental Scripts
-------------------------------------------------------

6. infinite_low_dim_all_methods.py  
   - Runs all estimators in the low-dimensional regime (raw, cubed, and teacher-transformed variants).  
   - Used to produce the data for Figure 1 and SI Figure 8.  
   - Operates in the infinite data regime (fresh batch per iteration).

7. infinite_high_dim.py  
   - Similar setup to above, but with high-dimensional teacher-transformed data.  
   - Produces the data for the right panel of Figure 2.

8. infinite_high_dim_cca.py  
   - Runs CCA-based estimator on replicated and teacher-transformed high-dimensional data.  
   - Produces the data for the left and center panels of Figure 2.

9. finite_summarizable.py  
   - Tests estimator behavior on teacher models with varying true latent dimensionalities.  
   - Includes early stopping logic.  
   - Used to generate the data for Figure 5 (and can be used to generate Figure 4 as well).

10. mnist_subsampling.py  
    - Subsampling experiments on the Noisy MNIST dataset.  
    - Computes MI across embedding dimensionalities and dataset sizes.  
    - Used to generate the data for Figure 7 and can be used (with smaller networks) to generate the data for Figure 6.

11. ksg_subsampling.py  
    - Runs the KSG estimator in a subsampling setup.  
    - Used to generate the data for SI Figure 10.

-------------------------------------------------------
Environment
-------------------------------------------------------

12. environment.yaml  
   - Dependencies file

-------------------------------------------------------
References
-------------------------------------------------------
[1] Poole et al., "On Variational Bounds of Mutual Information" (2019)  
[2] Song & Ermon, "Understanding the Limitations of Variational Mutual Information Estimators" (2019)  
[3] Abdelaleem et al., "Deep Variational Symmetric Information Bottleneck" (2023)