🌊 FMMI: Flow Matching Mutual Information Estimation

Official implementation of the FMDoE and FMMI estimators. See arXiv:2511.08552 for the full paper.

Description

This package implements the Flow Matching Difference of Entropies (FMDoE) estimator, as well as some special cases of it (e.g., the Flow Matching Mutual Information (FMMI) estimator) using PyTorch.

The core idea of the method is captured by the following equation:

$$\mathsf{h}(X_1) - \mathsf{h}(X_0) = \mathbb{E} , \mathrm{div} , v(X_T,T), \qquad T \sim \mathrm{U}[0;1],$$

where $v(x,t)$ transforms $X_0$ to $X_1$ via $\partial x_t / \partial t = v(x_t,t)$. Both the velocity field and the expectation can be estimated efficiently using the Flow Matching technique (in practie, we approximate $v$ with a neural network).

Features

• General Difference of Entropies estimator.
• Mutual information estomator.
• $\Omega$-information estimator.

Getting started

Examples

Examples of the application can be found in source/examples

Documentation

In development.

Cite

If you use our method, please consider citing our paper:

@misc{butakov2025fmmi,
    title={FMMI: Flow Matching Mutual Information Estimation}, 
    author={Ivan Butakov and Alexander Semenenko and Valeriya Kirova and Alexey Frolov and Ivan Oseledets},
    year={2025},
    eprint={2511.08552},
    archivePrefix={arXiv},
    primaryClass={cs.LG},
    url={https://arxiv.org/abs/2511.08552}, 
}