InfoBridge: Mutual Information estimation via Bridge Matching

This paper introduces InfoBridge, a novel method that leverages diffusion bridge models to frame mutual information estimation as a domain transfer problem, thereby providing an unbiased estimator that outperforms conventional approaches on diverse benchmarks including protein language model embeddings.

The Gist

The Big Picture: Measuring the "Secret Handshake"

Imagine you have two friends, Alice and Bob. You want to know: How much do they talk to each other? How much does Alice's mood influence Bob's mood?

In the world of data science, this question is called Mutual Information (MI). It measures how much knowing one thing tells you about another.

• If Alice and Bob are strangers, their MI is zero.
• If they are best friends who finish each other's sentences, their MI is high.

The Problem:
Calculating this "friendship score" is incredibly hard when you have millions of variables (like pixels in an image or words in a book). Traditional math tools break down in these high-dimensional worlds. They either get confused by the noise or take forever to compute.

The Solution:
The authors of this paper built a new tool called InfoBridge. Instead of trying to count every conversation Alice and Bob ever had, they watch how they move together over time to figure out how connected they are.

---

The Core Idea: Building a Bridge

To understand InfoBridge, imagine a river separating two towns: Town A (where Alice lives) and Town B (where Bob lives).

1. The Old Way (The "Noise" Approach)

Previous methods tried to estimate the connection by building a random path from a foggy swamp (pure noise) to Town A, and another from the swamp to Town B. Then, they compared the two paths.

• The Flaw: It's like trying to guess how well two people know each other by watching them both walk out of a foggy forest. It's messy, prone to errors, and the "fog" (mathematical bias) makes the results unreliable.

2. The New Way (InfoBridge)

InfoBridge takes a different approach. It builds a direct bridge between Town A and Town B.

• The Metaphor: Imagine you have a specific pair of shoes. One shoe is in Town A, the other in Town B.
• The Process: You ask a magical walker to start at the shoe in Town A and walk to the shoe in Town B.
  • Scenario 1 (Connected): If Alice and Bob are best friends, the walker finds a smooth, straight, easy path between the shoes. They are "in sync."
  • Scenario 2 (Disconnected): If they are strangers, the walker has to wander aimlessly, taking a long, chaotic, winding route to get from one to the other.

The Magic Trick:
The paper proves that the length and chaos of the path the walker takes is a perfect mathematical measure of how connected Alice and Bob are.

• Straight path = High Mutual Information (They are close).
• Winding path = Low Mutual Information (They are far apart).

How the "Magic Walker" Works (The Tech)

The paper uses a concept called Diffusion Bridge Matching. Think of this as a "time-lapse video" of the walker.

1. The Setup: The walker starts at a specific point (Alice's data) and must end at a specific point (Bob's data).
2. The Learning: The computer trains a neural network (a smart AI) to act as the walker. It learns the "drift" or the "push" needed to get from A to B.
   • If the data is highly connected, the AI learns a very direct, predictable push.
   • If the data is random, the AI has to push in all directions, creating a chaotic path.
3. The Calculation: The AI calculates the difference between the "Connected Path" (Alice to Bob) and the "Random Path" (Alice to a Stranger).
   • The bigger the difference in how the paths behave, the higher the Mutual Information.

Why is this Better? (The "Why Should You Care?")

The authors tested InfoBridge on three types of challenges, and it won every time:

1. Simple Data (Low-dimensional): Like basic math problems. InfoBridge was as good as the best existing tools.
2. Complex Data (Images): Imagine trying to measure the connection between two blurry, 32x32 pixel images. Old tools got lost in the noise. InfoBridge saw the pattern clearly.
   • Analogy: It's like being able to hear a whisper in a hurricane while everyone else is deafened by the wind.
3. Real-World Data (Proteins): They tested it on protein structures (the building blocks of life). This is huge data. InfoBridge accurately measured how different parts of a protein "talk" to each other, while other methods failed completely or gave wild, wrong numbers.

The "Secret Sauce"

The paper's main breakthrough is framing the problem as "Domain Transfer" (moving from A to B) rather than "Generative Modeling" (creating A from nothing).

• Old Way: "Let's try to create a picture of a cat from scratch, then a picture of a dog from scratch, and compare them." (Hard, noisy, biased).
• InfoBridge Way: "Let's take a real cat and a real dog, and watch how we can morph one into the other." (Direct, efficient, unbiased).

Summary

InfoBridge is a new, super-accurate ruler for measuring how much two things depend on each other.

• It uses a bridge instead of a fog.
• It watches movement instead of just counting.
• It works on images, proteins, and complex data where other tools fail.

It's like upgrading from a rusty, broken compass to a GPS that works perfectly even in a storm. This allows scientists to better understand everything from how AI learns to how our DNA works.

Technical Summary

1. Problem Statement

Mutual Information (MI) is a fundamental measure of non-linear dependence between random variables, widely used in machine learning for representation learning, regularization, and neural network analysis. However, estimating MI is notoriously difficult due to the curse of dimensionality, long-tailed distributions, and high MI values.

• Limitations of Existing Methods:
  • Non-parametric estimators (e.g., KSG, kernel density) fail in high dimensions.
  • Discriminative estimators (e.g., MINE, InfoNCE) often suffer from high variance or require massive batch sizes.
  • Generative estimators (e.g., MINDE) based on standard diffusion models frame MI estimation as a generative task (noise-to-data). This introduces a non-zero bias term that only vanishes as the diffusion time horizon approaches infinity, and the trajectories (noise-to-data) are complex and difficult to learn, leading to high variance in estimates.

2. Methodology: InfoBridge

The authors propose InfoBridge, a novel MI estimator that reframes the problem as a domain transfer task using Diffusion Bridge Matching and Reciprocal Processes.

Theoretical Foundation

The core theoretical contribution is a decomposition of Mutual Information into the difference of drifts between two specific stochastic processes:

1. Reciprocal Processes: The authors define a joint reciprocal process $Q_\pi$ induced by the joint distribution $\pi(x_0, x_1)$ and an independent reciprocal process $Q^{ind}_\pi$ induced by the product of marginals $\pi(x_0)\pi(x_1)$.
2. KL Divergence via Girsanov Theorem: Using the Girsanov theorem, the paper proves that the KL divergence between these two diffusion processes (which equals the Mutual Information $I(X_0; X_1)$) can be expressed as the time-integrated expected squared difference between their drift functions:
   $$I(X_0; X_1) = \frac{1}{2\epsilon} \int_0^1 \mathbb{E}_{q_\pi(x_t, x_0)} \left[ \| v_{joint}(x_t, t, x_0) - v_{ind}(x_t, t, x_0) \|^2 \right] dt$$
   Where:
   • $v_{joint}$ is the drift of the process conditioned on the joint distribution.
   • $v_{ind}$ is the drift of the process conditioned on the independent distribution.
   • $\epsilon$ is the volatility coefficient.

Practical Algorithm

The method avoids the bias of infinite-time diffusion by using Bridge Matching, which learns the drift directly from data-to-data trajectories (Brownian Bridges).

• Drift Learning: The drifts $v_{joint}$ and $v_{ind}$ are approximated by a single neural network $v_\theta$ with a binary conditioning input $s \in \{0, 1\}$.
  • $s=1$ approximates the joint drift.
  • $s=0$ approximates the independent drift.
• Training Procedure:
  1. Sample pairs $(x_0, x_1)$ from the joint distribution.
  2. Sample time $t \sim U[0, 1]$.
  3. Sample intermediate points $x_t$ from the Brownian Bridge connecting $x_0$ and $x_1$.
  4. For the independent drift, permute $x_1$ to create independent pairs $(x_0, \hat{x}_1)$ and sample corresponding $x_t$.
  5. Minimize the regression loss between the network output and the target drift term $\frac{x_1 - x_t}{1-t}$.
• Estimation: Once trained, the MI is estimated by Monte Carlo integration of the squared difference between the two drift outputs over sampled trajectories.

3. Key Contributions

1. Unbiased Theoretical Framework: Unlike MINDE, which relies on finite-time approximations of infinite-time diffusion (introducing bias), InfoBridge provides an unbiased estimator (under ideal learning conditions) by leveraging the exact KL decomposition of reciprocal processes over a finite time horizon $[0, 1]$.
2. Domain Transfer Perspective: By framing MI estimation as a data-to-data transfer problem (learning the bridge between $x_0$ and $x_1$) rather than noise-to-data generation, the method learns "straighter" and easier trajectories (Schrödinger Föllmer processes), resulting in significantly lower variance.
3. Generalizability: The framework naturally extends to estimating general KL divergence, differential entropy, interaction information (for $>2$ variables), and MI between variables of different dimensionalities.
4. State-of-the-Art Performance: The method demonstrates superior accuracy and stability across diverse benchmarks compared to existing discriminative, flow-based, and diffusion-based estimators.

4. Experimental Results

The authors evaluated InfoBridge on four distinct benchmarks:

• Low-Dimensional Benchmarks: On standard synthetic datasets (e.g., Gaussian, Student-t, Swiss roll), InfoBridge achieved performance comparable to MINDE and superior to classical methods (MINE, KSG) and flow-based methods. It successfully handled heavy-tailed distributions (Cauchy) using an Asinh transform.
• Image Data Benchmarks: Using synthetic images (Gaussians and Rectangles) mapped to manifolds, InfoBridge achieved the lowest Mean Absolute Error (MAE) and lowest variance among all methods. It outperformed MINDE-C/J and MIENF, particularly in stability.
• Protein Embeddings (Real-World Data): Tested on 1024-dimensional protein language model embeddings (ProtTrans). InfoBridge was the only method to provide accurate estimates (MAE $\approx$ 0.04). Competitors like MINDE drastically overestimated MI (MAE > 9), and discriminative methods failed to capture the true values.
• High Mutual Information Regimes: In high-dimensional tasks ($d=160$, $MI=80$), InfoBridge remained accurate. In contrast, discriminative methods (MINE, InfoNCE) collapsed to near-zero estimates, and MINDE consistently underperformed or exhibited extreme variance.

Key Finding: InfoBridge consistently demonstrated 1–2 orders of magnitude lower variance in estimates compared to MINDE, attributed to the stability of learning data-to-data bridges versus noise-to-data generation.

5. Significance

• Solving the High-Dimensional MI Problem: InfoBridge addresses the critical bottleneck of estimating MI in high-dimensional, complex real-world data (like protein embeddings or images) where traditional methods fail.
• Theoretical Rigor: It establishes a rigorous link between reciprocal processes, Schrödinger bridges, and information theory, providing a bias-free alternative to existing diffusion-based estimators.
• Practical Impact: The method enables more reliable analysis of deep neural networks (e.g., Information Bottleneck analysis), better self-supervised learning objectives, and improved alignment in generative models (e.g., text-to-image).
• Efficiency: While computationally more intensive than discriminative methods, it is comparable to MINDE and offers a much better trade-off between accuracy and variance, making it the preferred choice for applications requiring precise MI quantification.

The code is available at: https://github.com/SKholkin/infobridge.