SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Pawel Czyż*1,2 Frederic Grabowski*3
Julia E. Vogt $^{4,5}$  Niko Beerenwinkel $^{\dagger 1,5}$  Alexander Marx $^{\dagger 2,4}$
$^{1}$ Department of Biosystems Science and Engineering, ETH Zurich  $^{2}$ ETH AI Center, ETH Zurich
$^{3}$ Institute of Fundamental Technological Research, Polish Academy of Sciences
$^{4}$ Department of Computer Science, ETH Zurich  $^{5}$ SIB Swiss Institute of Bioinformatics
Mutual information is a general statistical dependency measure which has found applications in representation learning, causality, domain generalization and computational biology. However, mutual information estimators are typically evaluated on simple families of probability distributions, namely multivariate normal distribution and selected distributions with one-dimensional random variables. In this paper, we show how to construct a diverse family of distributions with known ground-truth mutual information and propose a language-independent benchmarking platform for mutual information estimators. We discuss the general applicability and limitations of classical and neural estimators in settings involving high dimensions, sparse interactions, long-tailed distributions, and high mutual information. Finally, we provide guidelines for practitioners on how to select appropriate estimator adapted to the difficulty of problem considered and issues one needs to consider when applying an estimator to a new data set.
Estimating the strength of a non-linear dependence between two continuous random variables (r.v.) lies at the core of machine learning. Mutual information (MI) lends itself naturally to this task, due to its desirable properties, such as invariance to homeomorphisms and the data processing inequality. It finds applications in domain generalization [Li et al., 2022, Ragonesi et al., 2021], representation learning [Belghazi et al., 2018, Oord et al., 2018], causality [Solo, 2008, Kurutach et al., 2018], physics [Keys et al., 2015], systems biology [Selimkhanov et al., 2014, Grabowski et al., 2019, Uda, 2020], and epidemiology [Young et al., 2023].
Over the last decades, the estimation of MI has been extensively studied, and estimators have been developed, ranging from classical approaches based on histogram density [Pizer et al., 1987], kernel density estimation [Moon et al., 1995] to  $k$ -nearest neighbor [Kozachenko and Leonenko, 1987, Kraskov et al., 2004] and neural estimators [Belghazi et al., 2018, Oord et al., 2018, Song and Ermon, 2020]. However, despite the progress that has been achieved in this area, not much attention has been focused on systematically benchmarking these approaches.
Typically, new estimators are evaluated assuming multivariate normal distributions for which MI is analytically tractable [Darbellay and Vajda, 1999, Kraskov et al., 2004, Suzuki, 2016]. Sometimes, simple transformations are applied to the data [Khan et al., 2007, Gao et al., 2015], moderately high-dimensional settings are considered [Lord et al., 2018, Lu and Peltonen, 2020], or strong dependencies are evaluated [Gao et al., 2015]. Beyond that, Song and Ermon [2020] study the
*Equal contribution † Joint supervision
Figure 1: Visualisations of selected proposed distributions. Two correlated Gaussian variables  $X$  and  $Y$  (1) can be transformed via the Gaussian CDF into correlated uniform variables (2), into a long-tailed distribution via the "half-cube" mapping  $t \mapsto t\sqrt{|t|}$  (3), into a multi-modal distribution (4) or embedded in the three-dimensional space via the composition of the Swiss roll mapping and Gaussian CDF (5). Color on the plot corresponds to the original  $Y$  variable.
self-consistency (additivity under independence and the data processing inequality) of neural MI estimators in the context of representation learning from image data.
Our Contributions In this work we show a method of developing expressive distributions with known ground-truth mutual information (Sec. 2), propose forty benchmark tasks and systematically study the properties of commonly used estimators, including representatives based on kernel or histogram density estimation, kNN estimation, and neural network-based estimators (Sec. 3). We address selected difficulties one can encounter when estimating mutual information (Sec. 4), such as sparsity of interactions, long-tailed distributions, invariance, and high mutual information. Finally, we provide recommendations for practitioners on how to choose a suitable estimator for particular problems (Sec. 6). Our benchmark is designed so that it is simple to add new tasks and estimators. It also allows for cross-language comparisons — in our work we compare estimators implemented in Python, R and Julia. All of our experiments are fully reproducible by running Snakemake workflows [Mölder et al., 2021]. Accompanying code is available at http://github.com/cbg-ethz/bmi.
Overall, our key findings from the benchmark can be summarized as follows:
- Testing on multivariate normal distributions gives a biased and overly optimistic view of estimator performance. In this setting, canonical correlation analysis (CCA), a model-based approach, emerges as the most effective method — even if the model is slightly misspecified.
- Compared to classical estimators, neural estimators excel in high-dimensional settings, capturing sparse interactions across multiple dimensions.
- The popular KSG estimator [Kriskov et al., 2004] is accurate in low- and moderate-dimensional settings, but its performance suffers on problems involving high-dimensions or sparse interactions.
- Although MI is invariant to a wide range of transformations (Theorem 2.1), numerical estimates, even with large sample sizes, are not. Thus, MI may not be suitable when metric invariance to specific transformation families is required.
- Multivariate Student distributions pose an important and hard challenge to many mutual information estimators. This effect is partially, but not fully, attributed to their long tails.
2 Mutual Information: Estimation and Invariance
We start by recalling the definition of MI, then discuss the problem of estimating it, and introduce the estimators used in the benchmark.
Consider two r.v.  $X$  and  $Y$  with domains  $\mathcal{X}$  respectively  $\mathcal{V}$ , joint probability measure  $P_{XY}$  and marginal probability measures  $P_{X}$  and  $P_{Y}$ , respectively. If  $P_{XY}$  is absolutely continuous with respect to  $P_{X} \otimes P_{Y}$  (e.g.,  $X$  and  $Y$  are finite r.v.), MI is equal to the following Kullback-Leibler divergence<sup>2</sup>:
$$
\mathbf {I} (X; Y) = \mathbf {D} _ {\mathrm {K L}} \left(P _ {X Y} \| P _ {X} \otimes P _ {Y}\right) = \int \log f \mathrm {d} P _ {X Y},
$$
where  $f = \mathrm{d}P_{XY} / \mathrm{d}(P_X\otimes P_Y)$  is the Radon-Nikodym derivative. If the absolute continuity does not hold, then  $\mathbf{I}(X;Y) = +\infty$  [Pinsker and Feinstein, 1964, Theorem 2.1.2]. If  $\mathcal{X}$  and  $\mathcal{V}$  are Euclidean spaces and measures  $P_{XY}, P_X$  and  $P_Y$  have probability density functions (PDFs) with respect to the Lebesgue measure, then the Kullback-Leibler divergence can be written in terms of the PDFs.
2We use the natural logarithm, meaning that the mutual information is measured in nats.
However, we will later consider distributions which do not have PDFs with respect to the Lebesgue measure and the general Radon-Nikodym derivative must be used (see "Swiss roll" in Fig. 1).
In almost all applications,  $P_{XY}$  and  $P_{X} \otimes P_{Y}$  are not known and one needs to estimate MI from a finite sample from  $P_{XY}$ , i.e., a realization  $\left((x_1, y_1), \ldots, (x_N, y_N)\right) \in (\mathcal{X} \times \mathcal{Y})^N$  of  $N$  i.i.d. r.v.  $(X_1, Y_1), \ldots, (X_N, Y_N)$  distributed according to the joint distribution  $P_{XY}$ . As a MI estimator we will denote a family of measurable mappings, indexed by the sample size  $N$  and spaces  $\mathcal{X}$  and  $\mathcal{Y}$  (in most applications, Euclidean spaces of various dimensionalities):  $e_N^{\mathcal{X}\mathcal{Y}}: (\mathcal{X} \times \mathcal{Y})^N \to \mathbb{R}$ . Composing this mapping with the r.v. representing the whole data set we obtain a real-valued r.v.  $E_N^{XY}$ . For a given  $P_{XY}$  we can consider the distribution of  $E_N^{XY}$ . It is often summarized by its first two moments, resulting in the bias and the variance of the estimator. Understanding the bias of an estimator however requires the knowledge of ground-truth MI, so the estimators are typically tested on the families of jointly multivariate normal or uniform distributions, with varying  $N$  and ground-truth  $\mathbf{I}(X; Y)$  [Khan et al., 2007, Lord et al., 2018, Holmes and Nemenman, 2019, Song and Ermon, 2020].
Constructing Expressive Distributions and Benchmarking Besides considering different types of distributions such as (multivariate) normal, uniform and Student, our benchmark is based on the invariance of MI to the family of chosen mappings. Generally MI is not invariant to arbitrary transformations, as due to the data processing inequality  $\mathbf{I}(f(X);g(Y))\leq \mathbf{I}(X;Y)$ . However, under some circumstances MI is preserved, which enables us to create more expressive distributions by transforming the variables.
Theorem 2.1. Let  $\mathcal{X},\mathcal{X}^{\prime},\mathcal{Y}$  and  $\mathcal{Y}'$  be standard Borel spaces (e.g., smooth manifolds with their Borel  $\sigma$ -algebras) and  $f\colon \mathcal{X}\to \mathcal{X}'$  and  $g\colon \mathcal{Y}\to \mathcal{Y}'$  be continuous injective mappings. Then, for every  $\mathcal{X}$ -valued r.v.  $X$  and  $\mathcal{Y}$ -valued r.v.  $Y$  it holds that  $\mathbf{I}(X;Y) = \mathbf{I}(f(X);g(Y))$ .
This theorem has been studied before by Kraskov et al. [2004, Appendix], who consider diffeomorphisms and by Polyanskiy and Wu [2022, Th. 3.7], who assume measurable injective functions with measurable left inverses. For completeness we provide a proof which covers continuous injective mappings in Appendix A. In particular, each of  $f$  and  $g$  can be a homeomorphism, a diffeomorphism, or a topological embedding. Embeddings are allowed to increase the dimensionality of the space, so even if  $X$  and  $Y$  have probability density functions,  $f(X)$  and  $g(Y)$  do not need to have them. Neither of these deformations change MI. Thus, using Theorem 2.1, we can create more expressive distributions  $P_{f(X)g(Y)}$  by sampling from  $P_{XY}$  and transforming the samples to obtain a data set  $((f(x_1), g(y_1)), \ldots, (f(x_N), g(y_N)) \in (\mathcal{X}' \times \mathcal{Y}')^N$ .
While offering a tool to construct expressive distributions, another valuable perspective to consider with this problem is estimation invariance: although MI is invariant to proposed changes, the estimate may not be. More precisely, applying an estimator to the transformed data set results in the induced r.v.  $E_N^{f(X)g(Y)}$  and its distribution. If the estimator were truly invariant, one would expect that  $\mathbb{E}\left[E_N^{f(X)g(Y)}\right]$  should equal  $\mathbb{E}\left[E_N^{XY}\right]$ . This invariance is rarely questioned and often implicitly assumed as given (see e.g. Tschannen et al. [2020] or Murphy [2023, Sec. 32.2.2.3]), but as we will see in Sec. 4.3, finite-sample estimates are not generally invariant.
3 Proposed Benchmark
In this section, we outline forty tasks included in the benchmark that cover a wide family of distributions, including varying tail behaviour, sparsity of interactions, multiple modes in PDF, and transformations that break colinearity. As our base distributions, we selected multivariate normal and Student distributions, which were transformed with continuous injective mappings.
We decided to focus on the following phenomena:
1. Dimensionality. High-dimensional data are collected in machine learning and natural sciences [Buhlmann and van de Geer, 2011, Ch. 1]. We therefore change the dimensions of the  $\mathcal{X}$  and  $\mathcal{Y}$  spaces between 1 and 25.
2. Sparsity. While the data might be high-dimensional, the effective dimension may be much smaller, due to correlations between different dimensions or sparsity [Lucas et al., 2006]. We therefore include distributions in which some of the dimensions represent random noise,
which does not contribute to the mutual information, and perform an additional study in Section 4.1.
3. Varying MI. Estimating high MI is known to be difficult [McAllester and Stratos, 2020]. However, we can often bound it in advance — if there are 4000 image classes, MI between image class and representation is at most  $\log 4000 \approx 8.3$  nats. In this section, we focus on problems with MI varying up to 2 nats. We additionally consider distributions with higher MI in Section 4.4.
4. Long tails. As Taleb [2020] and Zhang et al. [2021] argue, many real-world distributions have long tails. To model different tails we consider multivariate Student distributions as well as transformations lengthening the tails. We conduct an additional study in Section 4.2.
5. Robustness to diffeomorphisms. As stated in Theorem 2.1, mutual information is invariant to reparametrizations of random variables by diffeomorphisms. We however argue that when only finite samples are available, invariance in mutual information estimates may not be achieved. To test this hypothesis we include distributions obtained by using a diverse set of diffeomorphisms and continuous injective mappings. We further study the robustness to reparametrizations in Section 4.3.
While we provide a concise description here, precise experimental details can be found in Appendix D and we visualise selected distributions in Appendix F. We first describe tasks employing one-dimensional variables.
Bivariate Normal For multivariate normal variables MI depends only on the correlation matrix. We will therefore consider a centered bivariate normal distribution  $P_{XY}$  with  $\operatorname{Cor}(X, Y) = \rho$  and  $\mathbf{I}(X;Y) = -0.5\log(1 - \rho^2)$ . We chose  $\rho = 0.75$ .
Uniform Margins As a first transformation, we apply the Gaussian CDF  $F$  to Gaussian r.v.  $X$  and  $Y$  to obtain r.v.  $X' = F(X)$  and  $Y' = F(Y)$ . It is a standard result in copula theory (see Lemma D.2 or e.g., Nelsen [2007] for an introduction to copulas) that  $F(X) \sim \mathrm{Uniform}(0,1)$  resp.  $F(Y) \sim \mathrm{Uniform}(0,1)$ . The joint distribution  $P_{X'Y'}$ , however, is not uniform. Mutual information is preserved,  $\mathbf{I}(X';Y') = \mathbf{I}(X;Y)$ . For an illustration, see Fig. 1. A distribution  $\mathbb{P}$  transformed with Gaussian CDF  $F$  will be denoted as Normal CDF @ P.
Half-Cube Map To lengthen the tails we applied the half-cube homeomorphism  $h(x) = |x|^{3/2}$  sign  $x$  to Gaussian variables  $X$  and  $Y$ . We visualize an example in Fig. 1 and denote a transformed distribution as Half-cube @ P.
Asinh Mapping To shorten the tails we applied inverse hyperbolic sine function  $\operatorname{asinh} x = \log \left(x + \sqrt{1 + x^2}\right)$ . A distribution transformed with this mapping will be denoted as  $\operatorname{Asinh} @ P$ .
Wiggly Mapping To model non-uniform lengthscales, we applied a mapping
$$
w (x) = x + \sum_ {i} a _ {i} \sin (\omega_ {i} x + \varphi_ {i}), \qquad \sum_ {i} | a _ {i} \omega_ {i} | <   1.
$$
Due to the inequality constraint this mapping is injective and preserves MI. The parameter values can be found in Appendix D. We denote the transformed distribution as Wiggly @ P.
Bimodal Variables Applying the inverse CDF of a two-component Gaussian mixture model to correlated variables  $X$  and  $Y$  with uniform margins we obtained a joint probability distribution  $P_{X'Y'}$  with four modes (presented in Fig. 1). We call this distribution Bimodal and provide technical details in Appendix D.
Additive Noise Consider independent r.v.  $X \sim \mathrm{Uniform}(0,1)$  and  $N \sim \mathrm{Uniform}(-\varepsilon, \varepsilon)$ , where  $\varepsilon$  is the noise level. For  $Y = X + N$ , it is possible to derive  $\mathbf{I}(X;Y)$  analytically (see Appendix D for the formula and the parameter values) and we will call this distribution Uniform (additive noise= $\varepsilon$ ).
Swiss Roll Embedding As the last example in this section, we consider a mixed case in which a one-dimensional r.v.  $X \sim \mathrm{Uniform}(0,1)$  is smoothly embedded into two dimensions via the Swiss roll mapping, a popular embedding used to test dimensionality reduction algorithms (see Fig. 1 for visualisation and Appendix D for the formula). Note that the obtained distribution does not have a PDF with respect to the Lebesgue measure on  $\mathbb{R}^3$ .
Next, we describe the tasks based on multivariate distributions.
Multivariate Normal We sampled  $(X,Y) = (X_{1},\ldots ,X_{m},Y_{1},\ldots ,Y_{n})$  from the multivariate normal distribution. We model two correlation structures: "dense" interactions with all off-diagonal correlations set to 0.5 and "2-pair" interactions where  $\operatorname{Cor}(X_1,Y_1) = \operatorname{Cor}(X_2,Y_2) = 0.8$  and there is no correlation between any other (distinct) variables (Multinormal (2-pair)). For a latent-variable interpretation of these covariance matrices we refer to Appendix D.2.
Multivariate Student To see how well the estimators can capture MI contained in the tails of the distribution, we decided to use multivariate Student distributions (see Appendix D) with dispersion matrix $^3$  set to the identity  $I_{m+n}$  and  $\nu$  degrees of freedom. Contrary to the multivariate normal case, in which the identity matrix used as the covariance would lead to zero MI, the variables  $X$  and  $Y$  can still interact through the tails.
Transformed Multivariate Distributions As the last set of benchmark tasks we decided to transform the multivariate normal and multivariate Student distributions. We apply mappings described above to each axis separately. For example, applying normal CDF to the multivariate normal distribution will map it into a distribution over the cube  $(0,1)^{m + n}$  with uniform marginals. To mix different axes we used a spiral diffeomorphism (denoted as Spiral @ P); we defer the exact construction to Appendix D but we visualise this transformation in Fig. 5.
Estimators In our benchmark, we include a diverse set of estimators. Following recent interest, we include four neural estimators, i.e., Donsker-Varadhan (D-V) and MINE [Belghazi et al., 2018], InfoNCE [Oord et al., 2018], and the NWJ estimator [Nguyen et al., 2007, Nowozin et al., 2016, Poole et al., 2019]. As a representative for model-based estimators, we implement an estimator based on canonical correlation analysis (CCA) [Murphy, 2023, Ch. 28], which assumes that the joint distribution is multivariate normal. For more classical estimators, we include the Kraskov, Stögbauer and Grassberger (KSG) estimator [Kriskov et al., 2004] as the most well-known  $k$ -nearest neighbour-based estimator, a recently proposed kernel-based estimator (LNN) [Gao et al., 2017b], as well as an estimator based on histograms and transfer entropy (both implemented in Julia's TransferEntropy library). A detailed overview of the different classes of estimators is provided in Appendix C. Further, we describe the hyperparameter selection in Appendix E.
Preprocessing As a data preprocessing step, we standardize each dimension using the mean and standard deviation calculated from the sample. Other preprocessing strategies are also possible, but we did not observe a consistent improvement of one preprocessing strategy over the others (see Appendix E).
3.1 Benchmark Results
We show the results for  $N = 10000$  data points in Fig. 2. Neural estimators and KSG perform better than alternative estimators on most problems, with KSG having often low sample requirements (see Appendix E). The simple model-based estimator CCA obtained excellent performance at low sample sizes (see Appendix E), even for slightly transformed multivariate normal distributions. Finally, LNN and the two Julia estimators (histogram and transfer entropy), work well in low dimension but are not viable in medium- and high-dimensional problems (dim  $X + \dim Y \geq 4$ ).
Overall, we observed that for most  $1 \times 1$ -dimensional and multivariate normal problems MI can be reliably estimated. Difficulties arise for sparse interactions, atypical (Student) distributions, and for significantly transformed distributions. This suggests that the typical evaluations of MI estimators are overly optimistic, focusing on relatively simple problems.
KSG is able to accurately estimate MI in multivariate normal problems, however performance drops severely on tasks with sparse interactions (2-pair tasks). For  $5 \times 5$  dimensions the estimate is about  $70\%$  of the true MI, and for  $25 \times 25$  dimensions it drops to  $10\%$ , which is consistent with previous observations [Marx and Fischer, 2022]. Meanwhile, performance of neural estimators is stable. We study the effect of sparsity in more detail in Sec. 4.
Student distributions have proven challenging to all estimators. This is partially due to the fact that long tails can limit estimation quality (see Sec. 4), and indeed, applying a tail-shortening asinh transformation improves performance. However, performance still remains far below multivariate normal distributions with similar dimensionality and information content, showing that long tails are
3The dispersion matrix of multivariate Student distribution is different from its covariance, which does not exist for  $\nu \leq 2$
Figure 2: Mean MI estimates of nine estimators over  $n = 10$  samples with  $N = 10000$  points each against the ground-truth value on all benchmark tasks grouped by category. Color indicates relative negative bias (blue) and positive bias (red). Blank entries indicate that an estimator experienced numerical instabilities.
only part of the difficulty. For neural estimators, the issue could be explained by the fact that the critic is not able to fully learn the pointwise mutual information (see Appendix E).
Interestingly, the performance of CCA, which assumes a linear Gaussian model, is often favorable compared to other approaches and has very low sample complexity (see Appendix E). This is surprising, since in the case of transformed distributions the model assumptions are not met. The approach fails for Student distributions (for which the correlation matrix either does not exist or is the identity matrix), the Swiss roll embedding and multivariate normal distributions transformed using the spiral diffeomorphism. Since the spiral diffeomorphism proved to be the most challenging transformation, we study it more carefully in Sec. 4.
4 Selected Challenges
In this section we investigate distributions which proved to be particularly challenging for the tested estimators in our benchmark: sparse interactions, long tailed distributions and invariance to diffeomorphisms. Additionally we investigate how the estimators perform for high ground truth MI.
4.1 Sparsity of Interactions
In the main benchmark (Fig. 2) we observed that estimation of MI with "2-pair" type of interactions is considerably harder for the KSG estimator than the "dense" interactions. On the other hand, neural estimators were able to pick the relevant interacting dimensions and provide better estimates of MI in these sparse cases.