REMEDI: Corrective Transformations for Improved Neural Entropy Estimation

TL;DR

REMEDI addresses the challenge of estimating differential entropy in high-dimensional settings by marrying a flexible base density learned via cross-entropy with a Donsker-Varadhan-based corrective term. The method provides a tractable empirical loss that combines KNIFE-style cross-entropy and a neural DV term, with theoretical guarantees of consistency under mild, non-compact-support conditions. Empirically, REMEDI improves entropy estimation on synthetic datasets, enhances Information Bottleneck performance on MNIST/CIFAR-10/ImageNet, and enables generative sampling via rejection or Langevin dynamics. This approach yields practical benefits for information-theoretic objectives in deep learning and offers a bridge to generative modeling, with avenues for future exploration in higher dimensions and diffusion-model connections.

Abstract

Information theoretic quantities play a central role in machine learning. The recent surge in the complexity of data and models has increased the demand for accurate estimation of these quantities. However, as the dimension grows the estimation presents significant challenges, with existing methods struggling already in relatively low dimensions. To address this issue, in this work, we introduce $\texttt{REMEDI}$ for efficient and accurate estimation of differential entropy, a fundamental information theoretic quantity. The approach combines the minimization of the cross-entropy for simple, adaptive base models and the estimation of their deviation, in terms of the relative entropy, from the data density. Our approach demonstrates improvement across a broad spectrum of estimation tasks, encompassing entropy estimation on both synthetic and natural data. Further, we extend important theoretical consistency results to a more generalized setting required by our approach. We illustrate how the framework can be naturally extended to information theoretic supervised learning models, with a specific focus on the Information Bottleneck approach. It is demonstrated that the method delivers better accuracy compared to the existing methods in Information Bottleneck. In addition, we explore a natural connection between $\texttt{REMEDI}$ and generative modeling using rejection sampling and Langevin dynamics.

Figures (18)

• Figure 1: KNIFE training curves with error bars on $8$-dimensional triangle and uniform ball datasets. It is observed that increasing the number of components $M$ for KNIFE leads to overfitting in both datasets.
• Figure 2: Results on two moons dataset. In the middle we see what direction (positive or negative) REMEDI affects the base distribution. To the right is the unnormalized distribution implied by $q(x) e^{T(x)}$.
• Figure 3: Results on one-dimensional triangle dataset. On the bottom is the data distribution, compared to the density that KNIFE and REMEDI (up to a constant) has learned.
• Figure 4: Training curve from the 8-dimensional triangle dataset. The horizontal dashed line indicates where the KNIFE training phase ends and REMEDI takes over.
• Figure 5: Plot showing test error of the Information Bottleneck methods vs $\beta$ on benchmark image classification datasets (error bars represent standard deviations). For most $\beta$ values, consistently REMEDI performs better than other methods on MNIST and ImageNet. On CIFAR10, the classification errors are similar for all the methods. However, REMEDI exhibits the lowest classification error across the $\beta$ values.

Theorems & Definitions (27)

• Proposition 1.1: Donsker-Varadhan
• Definition 1.2
• Definition 1.3
• Proposition 1.4
• proof
• Definition 1.5
• Lemma 1.6: Lemma 3.1 in van de Geer
• proof
• Definition 1.7
• Lemma 1.8