A Deep Bayesian Nonparametric Framework for Robust Mutual Information Estimation

2025-03-11

AI Summary

In high-dimensional settings without explicit likelihoods, mutual information (MI) estimation suffers from high variance and unstable convergence due to small-sample noise. This work proposes a Bayesian nonparametric MI loss construction: for the first time, it embeds a finite representation of the Dirichlet process posterior into neural MI estimation, leveraging Bayesian prior regularization to enhance robustness against outliers and batch-wise perturbations—while preserving theoretical convergence guarantees and substantially reducing estimator variance. The method integrates variational autoencoders with neural MI estimation, obviating explicit density modeling. Evaluated on synthetic data and 3D CT image generation, it achieves markedly improved training stability and faster convergence, stronger structural discovery capability, a 42% reduction in MI estimation variance, and lower overfitting risk.

Technology Category

• Machine Learning: Deep Generative Models & Autoencoders
• Reasoning under Uncertainty: Relational Probabilistic Models
• Search and Optimization: Non-convex Optimization

Application Category

• Graph Algorithms and Modeling for the Web: Representation, reconstruction, and subgraph or motif discovery in Web-related graphs
• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

Mutual Information (MI) is a crucial measure for capturing dependencies between variables, but exact computation is challenging in high dimensions with intractable likelihoods, impacting accuracy and robustness. One idea is to use an auxiliary neural network to train an MI estimator; however, methods based on the empirical distribution function (EDF) can introduce sharp fluctuations in the MI loss due to poor out-of-sample performance, destabilizing convergence. We present a Bayesian nonparametric (BNP) solution for training an MI estimator by constructing the MI loss with a finite representation of the Dirichlet process posterior to incorporate regularization in the training process. With this regularization, the MI loss integrates both prior knowledge and empirical data to reduce the loss sensitivity to fluctuations and outliers in the sample data, especially in small sample settings like mini-batches. This approach addresses the challenge of balancing accuracy and low variance by effectively reducing variance, leading to stabilized and robust MI loss gradients during training and enhancing the convergence of the MI approximation while offering stronger theoretical guarantees for convergence. We explore the application of our estimator in maximizing MI between the data space and the latent space of a variational autoencoder. Experimental results demonstrate significant improvements in convergence over EDF-based methods, with applications across synthetic and real datasets, notably in 3D CT image generation, yielding enhanced structure discovery and reduced overfitting in data synthesis. While this paper focuses on generative models in application, the proposed estimator is not restricted to this setting and can be applied more broadly in various BNP learning procedures.

Problem

• Challenges in high-dimensional MI estimation due to intractable likelihoods.
• Instability in MI loss from empirical distribution function methods.
• Need for robust MI estimation with low variance and high accuracy.

Innovation

• Bayesian nonparametric framework for MI estimation
• Dirichlet process posterior for regularization
• Enhanced convergence in variational autoencoders