SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We trained NWJ estimator with the default hyperparameters on  $N = 5000$  data points and visualised the critic predictions (with 1 subtracted), true PMI, and 500 samples from the distribution in Fig. 8. The data is drawn from a bivariate Student distribution with dispersion matrix
$$
\Omega = \left( \begin{array}{c c} 1 & 0. 5 \\ 0. 5 & 1 \end{array} \right)
$$
and varying degrees of freedom  $\nu \in \{1,2,3,5,8\}$ . Additionally, we repeated the experiment for the bivariate normal distribution with covariance matrix given by the dispersion, which corresponds to the limit  $\nu \to \infty$ .
Qualitatively we see that the critic does not learn the correct PMI for  $\nu \in \{1,2\}$ , with better similarities for  $\nu \in \{3,5,8\}$ . For the normal distribution, the critic qualitatively approximates well the PMI function in the regions where the data points are available.
We hypothesise that these problems can be partly attributed to the longer tails of Student distributions (as the critic may be unable to extrapolate to the regions where few points are available) and partly to the analytical form of the PMI function — for the considered bivariate normal distribution the PMI is given by a simple quadratic function, which is not the case for the bivariate Student distribution.
E.2 Sample Size Requirements
We investigated the minimal number of samples needed to estimate the mutual information changing the number of samples in  $N \in \{100, 500, 1000, 3000, 5000, 10000\}$ . We say that sample  $\ell$  suffices
to solve the task if for every  $\ell' \geq \ell$  the estimator provides the estimate between 2/3 and 3/2 of the ground-truth mutual information.
We present the minimal number of samples  $\ell$  as a function of benchmark task and considered estimator in Fig. 9. Note that CCA and KSG typically need very small amount of samples to solve the task. As neural estimators use batch size of 256 and split the sample into training and test data sets, it is not possible for them to have  $\ell < 1000$ , which is the smallest number of samples considered greater than  $2 \cdot 256 = 512$ .
E.3 Preprocessing
In Fig. 2 and Fig. 9 we transformed all samples by removing the mean and scaling by the standard deviation in each dimension calculated empirically from the sample. This strategy ensures that the scales on different axes are similar. Note that this strategy is not always suitable (e.g., if the distance metric used in the kNN approach should weight select dimensions differently) and makes the observed data points not i.i.d. (as the "shared" information, sample mean and variances, has been used to rescale the variables), although the effect of it seems to be negligible on most tasks in the benchmark.
Additionally, we considered two other preprocessing strategies, implemented using the scikit-learn package [Pedregosa et al., 2011]: making each marginal distribution uniform  $\mathrm{Uniform}(0,1)$  by calculating empirical CDF along each axis (see Section D for the technical lemmata discussing this operation) and making each marginal distribution a Gaussian variable  $\mathcal{N}(0,1)$ . These approaches may affect the effective sample size more strongly than the previous rescaling strategy, as estimating empirical CDF requires more data points than estimating mean and standard deviation.
We used  $N = 10000$  data points and rerun the benchmark (see Fig. 10). Overall, the impact of the "uniformization" strategy appeared to be insubstantial for most tasks, and in certain cases negative. The "Gaussianisation" strategy improved estimation for Student distributions and helped histogram based approaches in low-dimensional tasks. Interestingly, both approaches worsen the performance of CCA on Student distributions.
While the improvements gained from "Gaussianisation" were minor, the approach is simplistic, treating each axis separately. Ideally, appropriate preprocessing would be able to transform  $P_{XY}$  into a better-behaved distribution  $P_{f(X)g(Y)}$ , and thus simplify the estimation. In certain cases such preprocessing could (partially) undo transformations considered in our benchmark, and in effect make the estimators invariant (or more robust) to these transformations (cf. Section 4). We consider the design of informed preprocessing strategies an important open problem.
E.4 Hyperparameter Selection
CCA We used the implementation provided in the scikit-learn [Pedregosa et al., 2011] Python package (version 1.2.2). We set the latent space dimension to the smaller of the dimensions of the considered random vectors  $X$  and  $Y$ .
Histograms Adaptive estimators such as JIC [Suzuki, 2016] or MIIC [Cabeli et al., 2020] were not applicable to the high-dimensional data we considered, hence we resort to using an equal-width estimator implemented in TransferEntropy library (version 1.10.0). We fixed the number of bins to 10, although the optimal number of bins is likely to depend on the dimensionality of the problem.
KSG As a default, we used the KSG-1 variant from the rmi package with 10 neighbors. However, as shown in our ablation in Fig. 11, changing the number of neighbors or using the KSG-2 variant did not considerably change the estimates. We compared the used implementation with a custom Python implementation and one from TransferEntropy library. The version implemented in the TransferEntropy suffers from strong positive bias.
LNN We used the implementation from the rmi package (version 0.1.1) with 5 neighbors and truncation parameter set to 30. Ablations can be seen in Fig. 11.
Neural Estimators We split the data set into two equal-sized parts (training and test) and optimized the statistics network on the training data set using the Adam optimiser with initial learning rate set to 0.1. We used batch size of 256 and ran each algorithm for up to 10000 training steps with early stopping (checked every 250 training steps). We returned the highest estimate on the test data set.
Figure 10: Estimates for preprocessed data with  $N = 10000$  samples. Top: margins transformed into uniform distributions. Bottom: margins transformed into normal distributions.
We chose the architecture using a mini-benchmark with four ReLU neural networks considered on two neural estimators (see Fig. 11) and chose a neural network with two hidden layers (network M in Table 2).
As neural estimators rely on stochastic gradient descent training, we implemented heuristic diagnostics used to diagnose potential problems with overfitting and non-convergence. We present the number of runs susceptible to these issues in Fig. 12. We did not include these runs when computing statistics in the main benchmark figures (Fig. 2 and Fig. 9).
Figure 11: Left: ablation of hyperparameters of the KSG and LNN estimators (with  $n$  neighbors and threshold  $t$ ). Right: ablation of neural network architectures.
Architecture	Hidden layers
D	8, 8, 8
L	24, 12
M	16, 8
S	10, 5
Table 2: Network architectures and their corresponding hidden layer specifications.
F Visualisation of Selected Tasks
In the following, we visualise a selection of benchmark tasks. Each figure visualises a sample of  $N = 10000$  points from  $P_{XY}$ . The upper diagonal represents the scatter plots  $P_{X_iX_j}$ ,  $P_{X_iY_j}$  and  $P_{Y_iY_j}$  for all possible combinations of  $i \neq j$ . The lower diagonal represents the empirical density fitted using a kernel density estimator with 5 levels. The diagonal represents the histograms of the marginal distributions  $P_{X_i}$  and  $P_{Y_i}$  for all possible  $i$ .
We present selected  $1 \times 1$ -dimensional tasks in Fig. 13 and Fig. 14; the  $2 \times 1$ -dimensional Swiss-roll distribution in Fig. 15 and selected  $2 \times 2$ -dimensional tasks in Fig. 16. Finally, in Fig. 17 we present two  $3 \times 3$ -dimensional problems, visualising the effect of the normal CDF and the spiral diffeomorphism.
Figure 12: Number of runs (out of 10) when our convergence and overfitting diagnostics diagnosed potential problems in neural estimator training.
Figure 13: Top row: additive noise model. Bottom row: bivariate normal distribution and transformation making each variable a bimodal mixture of Gaussians.
G Author Contributions
Contributions according to the Contributor Roles Taxonomy:
1. Conceptualization: PC, FG, AM, NB.
2. Methodology: FG, PC.
3. Software: PC, FG.
4. Validation: FG, PC.
5. Formal analysis: PC, FG.
6. Investigation: FG, PC, AM.
Figure 14: Top row: bivariate normal distribution transformed so that the margins are uniform and asinh-transformed bivariate Student distribution with 1 degree of freedom. Bottom row: bivariate normal distribution transformed with the half-cube and the wiggly mappings.
7. Resources: FG, NB.
8. Data Curation: FG, PC.
9. Visualization: FG, PC.
10. Supervision: AM, NB.
11. Writing - original draft: PC, AM, FG.
12. Writing - review and editing: PC, AM, FG, NB, JEV.
Figure 15: Visualisation of the Swiss-roll distribution.
Figure 16: Visualisations of selected  $2 \times 2$ -dimensional problems. Top row: multivariate normal distributions with dense and sparse interactions. Bottom row: multivariate Student distribution with 2 degrees of freedom.
Figure 17: Visualisations of selected  $3 \times 3$ -dimensional problems. Top: multivariate normal distribution transformed axis-wise with normal CDF, to uniformise margins. Bottom: multivariate normal distribution transformed with the spiral diffeomorphism.