Package: torch-mist
Mist - A PyTorch Mutual information Estimation toolkitMutual Information Estimation toolkit based on pytorch.
Please refer to thedocumentationfor additional details regardinginstallation,usage,
tutorials and pracical use-case example.InstallationThe package can be installed via pip as follows:$pipinstalltorch_mistUsageThetorch_mistpackage provides the basic functionalities for sample-based continuous mutual information estimation using modern
neural network architectures.Here we provide a simple example of how to use the package to estimate mutual information between pairs
of observations using the MINE estimator[2].
Consider the variables $x$ and $y$ as of shape[N, x_dim],[N, y_dim]respectively sampled from some joint distribution $p(x,y)$.
Mutual information can be estimated directly using theestimate_miutility function that takes care of fitting the estimator's parameters and evaluating mutual information.fromtorch_mistimportestimate_mifromsklearn.datasetsimportload_iris# Load the Iris Dataset as a pandas DataFrameiris_dataset=load_iris(as_frame=True)['data']# Estimate how much information the petal length and its width have in commonestimated_mi,estimator,train_log=estimate_mi(data=iris_dataset,# The dataset (as a pandas.DataFrame, many other formats are supported)x_key='petal length (cm)',# Consider the 'petal length' column as xy_key='petal width (cm)',# And the 'petal width` as yestimator='mine',# Use the MINE mutual information estimatormax_iterations=1000,# Number of maximum train iterationshidden_dims=[128],# Hidden layers for the critic architecture.)print(f"Mutual information estimated value:{estimated_mi}nats")Theestimate_mifunction supports additional data formats such as tuples(x, y)or dictionaries{'x':x, 'y':y}of tensors, but also datasets (seetorch.utils.data.Dataset) and torchDataLoaders.Additional flags can be used to customize the estimators, training and evaluation procedure, as detailed in thedocumentation.Command lineThetorch_mistpackage provides basic functionality to estimate mutual information directly from the command line.
Given a fileiris.csvcontaining the columns(sepal_1, sepal_2, petal_1, petal_2), one can estimate mutual information between
thesepalandpetal2-dimensional features with:mistdata=csvdata.filepath=iris.csvmi_estimator=jsx_key=sepaly_key=petalThe same flags and options provided by theestimate_mifunction are also available from command line.Additionally, internal properties of the estimator can also be easily specified:mistdata=csvdata.filepath=iris.csvmi_estimator=jsx_key=sepaly_key=petal\# Train on GPUdevice=cuda\# Use AdamW for the optimizationestimation.optimizer_class._target_=torch.optim.AdamW\# Use ELU as nonlinearities+mi_estimator.nonlinearity=torch.nn.ELU\# Change the batch size to 256params.batch_size=256# Log on weights and biaslogger=wandbTo visualize the full list of arguments use:mistdata=csv--helpThemistCLI is implemented usinghydraand the full configuration can be accessedhere.Advanced usageIt is possible to manually instantiate, train and evaluate the mutual information estimators.fromtorch_mist.estimatorsimportminefromtorch_mist.utils.trainimporttrain_mi_estimatorfromtorch_mist.utilsimportevaluate_mi# Instantiate the JS mutual information estimatorestimator=mine(x_dim=1,y_dim=1,hidden_dims=[64,32],)# Define x and y as the vectors of petal lengths and widths, respectivelyx=iris_dataset['petal length (cm)'].valuesy=iris_dataset['petal width (cm)'].values# Train it on the given samplestrain_log=train_mi_estimator(estimator=estimator,train_data=(x,y),batch_size=64,max_iterations=1000)# Evaluate the estimator on the entirety of the dataestimated_mi=evaluate_mi(estimator=estimator,data=(x,y),batch_size=64)print(f"Mutual information estimated value:{estimated_mi}nats")Note that the two code snippets above perform the same procedure.
Please refer to thedocumentationfor a detailed description of the package and its usage.EstimatorsEach estimator implemented in the library is an instance ofMutualInformationEstimatorand can be instantiated
through a simplified utility functions############################# Simplified instantiation #############################fromtorch_mist.estimatorsimportmineestimator=mine(x_dim=1,y_dim=1,neg_samples=16,hidden_dims=[64,32],critic_type='joint')or directly using the correspondingMutualInformationEstimatorclass########################### Advanced instantiation ###########################fromtorch_mist.estimatorsimportMINEfromtorch_mist.criticimportJointCriticfromtorchimportnn# First we define the critic architecturecritic=JointCritic(# Wrapper to concatenate the inputs x and yjoint_net=nn.Sequential(# The neural network architectures that maps [x,y] to a scalarnn.Linear(x.shape[-1]+y.shape[-1],32),nn.ReLU(True),nn.Linear(32,32),nn.ReLU(True),nn.Linear(32,1)))# Then we pass it to the MINE constructorestimator=MINE(critic=critic,neg_samples=16,)Note that the simplified and advanced instantiation reported in the example above result in the same model.The basic estimators implemented in this package are summarized in the following table:EstimatorTypeModelsHyperparametersNWJ[1]Discriminative$f_\phi(x,y)$MMINE[2]Discriminative$f_\phi(x,y)$M, $\gamma_{EMA}$InfoNCE[3]Discriminative$f_\phi(x,y)$MTUBA[4]Discriminative$f_\phi(x,y)$, $b_\xi(x)$MAlphaTUBA[4]Discriminative$f_\phi(x,y)$, $b_\xi(x)$M, $\alpha$JS[5]Discriminative$f_\phi(x,y)$MSMILE[6]Discriminative$f_\phi(x,y)$M, $\tau$FLO[7]Discriminative$f_\phi(x,y)$, $b_\xi(x,y)$MBA[8]Generative$q_\theta(y|x)$-DoE[9]Generative$q_\theta(y|x)$, $q_\psi(y)$-GM[6]Generative$q_\theta(x,y)$, $q_\psi(x)$, $q_\psi(y)$-L1OUT[4][10]Generative$q_\theta(y|x)$-CLUB[10]Generative$q_\theta(y|x)$-Binned[13]Transformed (Generative)$Q(x)$, $Q(y)$-PQ[11]Transformed (Generative)$Q(y)$, $q_\theta(Q(y)|x)$-in which the following models are used:$f_\phi(x,y)$ is acriticneural network with parameters $\phi, which maps pairs of observations to a scalar value.
Critics can be eitherjointorseparabledepending on whether they parametrize function of both $x$ and $y$ directly,
or through the product of separate projection heads ( $f_\phi(x,y)=h_\phi(x)^T h_\phi(y)$ ) respectively.$b_\xi(x)$ is abaselineneural network with parameters $\xi$, which maps observations (or paris of observations) to a scalar value.
When the baseline is a function of both $x$ and $y$ it is referred to as ajoint_baseline.$q_\theta(y|x)$ is a conditional variational distributionq_Y_given_Xused to approximate $p(y|x)$ with parameters $\theta$.
Conditional distributions may have learnable parameters $\theta$ that are usually parametrized by a (conditional) normalizing flow.$q_\psi(y)$ is a marginal variational distributionq_Yused to approximate $p(y)$ with parameters $\psi$.
Marginal distributions may have learnable parameters $\psi$ that are usually parametrized by a normalizing flow.$q_\theta(x,y)$ is a joint variational distributionq_XYused to approximate $p(x,y)$ with parameters $\theta$.
Joint distributions may have learnable parameters $\theta$ that are usually parametrized by a normalizing flow.$Q(x)$ and $Q(y)$ arequantizationfunctions that map observations to a finite set of discrete values.And the following hyperparameters:$M \in [1, N]$ is the number of samples used to estimate the log-normalization constant for each element in the batch.$\gamma_{EMA} \in (0,1]$ is the exponential moving average decay used to update the baseline in MINE.$\alpha \in [0,1]$ is the weight of the baseline in AlphaTUBA (0 corresponds to InfoNCE, 1 to TUBA).$\tau \in [0..]$ is used to define the interval $[-\tau,\tau]$ in which critic values are clipped in SMILE.Hybrid estimatorsThetorch_mistpackage allows to combine Generative and Discriminative estimators in a single hybrid estimators as proposed in[11][12].
Hybrid mutual information estimators combine the flexibility of discriminative mutual information estimators with the lower
variance of generative estimators.fromtorch_mist.estimators.hybridimportResampledHybridMIEstimatorfromtorch_mist.estimatorsimportjs,doe# Use the proposal r(y|x) to sample negatives instead of p(y)estimator=ResampledHybridMIEstimator(# Difference of Entropies generative estimatorgenerative_estimator=doe(x_dim=x.shape[-1],y_dim=y.shape[-1],hidden_dims=[32,32],),# NWJ discriminative estimatordiscriminative_estimator=js(x_dim=x.shape[-1],y_dim=y.shape[-1],hidden_dims=[32,32],neg_samples=16))Further details on the available hybrid mutual information estimators and additional details are reported in thetutorialavailable in thedocumentation.Training and EvaluationMost of the estimators included in this package are parametric and require a training procedure for accurate estimation.
Thetrain_mi_estimatorutility function supports either row dataxandyasnumpy.arrayortorch.Tensor.fromtorch_mist.utils.trainimporttrain_mi_estimator####################################### Training using tensors for x and y ######################################## By default 10% of the data is used for cross-validation and early stoppingtrain_log=train_mi_estimator(estimator=estimator,train_data=(x,y),batch_size=64,valid_percentage=0.1,)Alternatively, it is possible to use atorch.utils.DataLoaderthat returns eiter batches of pairs(batch_x, batch_y)or dictionaries of batches{'x': batch_x, 'y': batch_y}, withbatch_xof shape[batch_size, ..., x_dim]and[batch_size, ..., y_dim]respectively.############################## Training with DataLoaders ##############################fromtorch_mist.utils.dataimportSampleDatasetfromtorch.utils.dataimportDataLoader,random_split# We provide an utility to make the tensors into a torch.utils.data.Dataset object# This can be replaced with any other Dataset object that may load the data from diskdataset=SampleDataset(samples={'x':x,'y':y})# Split into train and validationtrain_size=int(len(dataset)*0.9)valid_size=len(dataset)-train_sizetrain_set,valid_set=random_split(dataset,[train_size,valid_size])# Instantiate the dataloaderstrain_loader=DataLoader(train_set,batch_size=64,shuffle=True,num_workers=8)valid_loader=DataLoader(valid_set,batch_size=64,num_workers=8)# Train using the specified dataloaders# Note that the validation set is optional but recommended to prevent overfitting.train_log=train_mi_estimator(estimator=estimator,data=train_loader,valid_data=valid_loader,)The two options result in the same training procedure, but we recommend usingDataLoaderfor larger datasets.BothDataLoaderandtorch.Tensor(ornp.array) can be used for theevaluate_mifunction.References[1]Nguyen, XuanLong, Martin J. Wainwright, and Michael I. Jordan. "Estimating divergence functionals and the likelihood ratio by convex risk minimization." IEEE Transactions on Information Theory 56.11 (2010): 5847-5861.[2]Belghazi, Mohamed Ishmael, et al. "Mutual information neural estimation." International conference on machine learning. PMLR, 2018.[3]Oord, Aaron van den, Yazhe Li, and Oriol Vinyals. "Representation learning with contrastive predictive coding." arXiv preprint arXiv:1807.03748 (2018).[4]Poole, Ben, et al. "On variational bounds of mutual information." International Conference on Machine Learning. PMLR, 2019.[5]Hjelm, R. Devon, et al. "Learning deep representations by mutual information estimation and maximization." arXiv preprint arXiv:1808.06670 (2018).[6]Song, Jiaming, and Stefano Ermon. "Understanding the limitations of variational mutual information estimators." arXiv preprint arXiv:1910.06222 (2019).[7]Guo, Qing, et al. "Tight mutual information estimation with contrastive fenchel-legendre optimization." Advances in Neural Information Processing Systems 35 (2022): 28319-28334.[8]Barber, David, and Felix Agakov. "The im algorithm: a variational approach to information maximization." Advances in neural information processing systems 16.320 (2004): 201.[9]McAllester, David, and Karl Stratos. "Formal limitations on the measurement of mutual information." International Conference on Artificial Intelligence and Statistics. PMLR, 2020.[10]Cheng, Pengyu, et al. "Club: A contrastive log-ratio upper bound of mutual information." International conference on machine learning. PMLR, 2020.[11]Federici, Marco, David Ruhe, and Patrick Forré. "On the Effectiveness of Hybrid Mutual Information Estimation." arXiv preprint arXiv:2306.00608 (2023).[12]Brekelmans, Rob, et al. "Improving mutual information estimation with annealed and energy-based bounds." arXiv preprint arXiv:2303.06992 (2023).[13]Kraskov, Alexander, Harald Stögbauer, and Peter Grassberger. "Estimating mutual information." Physical review E 69.6 (2004): 066138.ContributingInterested in contributing? Check out the contributing guidelines. Please note that this project is released with a Code of Conduct. By contributing to this project, you agree to abide by its terms.Licensetorch_mistwas created by Marco Federici. It is licensed under the terms of the MIT license.Creditstorch_mistwas created withcookiecutterand thepy-pkgs-cookiecuttertemplate.