InfoBridge: Mutual Information estimation via Bridge Matching

Diffusion bridge models have recently become a powerful tool in the field of generative modeling. In this work, we leverage them to address another important problem in machine learning and information theory, the estimation of the mutual information (MI) between two random variables. Neatly framing MI estimation as a domain transfer problem, we construct an unbiased estimator for data posing difficulties for conventional MI estimators. We showcase the performance of our estimator on three standard MI estimation benchmarks, i.e., low-dimensional, image-based and high MI, and on real-world data, i.e., protein language model embeddings. The code for our estimator can be found at: https://github.com/SKholkin/infobridge