/TC-estimation

Code for AISTATS 2023 paper - Estimating Total Correlation with Mutual Information Estimators

Primary LanguageJupyter Notebook

Total Correlation (TC) Estimation

This repo contains the implementation of the AISTATS 2023 paper: Estimating Total Correlation with Mutual Information Estimators.

The above figure shows the estimation performance of our TC-Tree and TC-Line estimators based on different mutual information (MI) estimators.

Introduction

We implement the TCLineEstimator and TCTreeEstimator in tc_estimators.py.

Both TC estimators are based on MI estimators (NWJ, MINE, InfoNCE, CLUB) in mi_estimators.py.

In tc_estimation.ipynb, we conduct a toy simulation to test the estimation ability of our TC estimators.

Method

Mutual information (MI) is a fundamental measurement of correlation between two variables:

Total correlation (TC) is an extension of MI for multi-variate scenarios:

We introduce two calculation paths to decomposite the total correlation into mutual information terms:

  • Line-like decomposition:

  • Tree-like decomposition:

The calculation paths are demonstrated in the following figure:

Citation

Welcome to cite our paper if the code is useful:

@InProceedings{pmlr-v206-bai23a,
  title = {Estimating Total Correlation with Mutual Information Estimators},
  author = {Bai, Ke and Cheng, Pengyu and Hao, Weituo and Henao, Ricardo and Carin, Larry},
  booktitle = {Proceedings of The 26th International Conference on Artificial Intelligence and Statistics},
  pages = {2147--2164},
  year = {2023},
  volume = {206},
  series = {Proceedings of Machine Learning Research},
  month = {25--27 Apr},
  publisher = {PMLR},
  url = {https://proceedings.mlr.press/v206/bai23a.html}
}