Approximating mutual information of high-dimensional variables using learned representations
Authors: Gokul Gowri
Xiaokang Lun
Allon M Klein
Peng Yin
Venue: NeurIPS 2024 spotlight (2024)
Keywords: mutual information
high-dimensional data
computational biology
scRNA-seq
protein language models
Primary area: machine_learning_for_other_sciences_and_fields

Abstract: Mutual information (MI) is a general measure of statistical dependence with widespread application across the sciences. However, estimating MI between multi-dimensional variables is challenging because the number of samples necessary to converge to an accurate estimate scales unfavorably with dimensionality. In practice, existing techniques can reliably estimate MI in up to tens of dimensions, but fail in higher dimensions, where sufficient sample sizes are infeasible. Here, we explore the idea that underlying low-dimensional structure in high-dimensional data can be exploited to faithfully approximate MI in high-dimensional settings with realistic sample sizes. We develop a method that we call latent MI (LMI) approximation, which applies a nonparametric MI estimator to low-dimensional representations learned by a simple, theoretically-motivated model architecture. Using several benchmarks, we show that unlike existing techniques, LMI can approximate MI well for variables with $> 10^3$ dimensions if their dependence structure is captured by low-dimensional representations. Finally, we showcase LMI on two open problems in biology. First, we approximate MI between protein language model (pLM) representations of interacting proteins, and find that pLMs encode non-trivial information about protein-protein interactions. Second, we quantify cell fate information contained in single-cell RNA-seq (scRNA-seq) measurements of hematopoietic stem cells, and find a sharp transition during neutrophil differentiation when fate information captured by scRNA-seq increases dramatically. An implementation of LMI is available at *latentmi.readthedocs.io.*

Approximating mutual information of high-
dimensional variables using learned representations
Gokul Gowri 2,,, Kino-Kang Lun Allon M. Klein s. and Peng Yin 2.
Wyss Institute for Biologically Inspired Engineering Department of Systems Biology, Harvard University
corresponding authors: harvard.edu, citio__lein@hmm.cm klein@hhssharrarrd.ddd harvard.edu
senior authors: A.M.K. and P.Y. co-supervised this work.
Abstract
Mutual information (MI) is general measure of statistical dependence with widespread application across the sciences. However, estimating MI between multi- dimensional variables is challenging because the number of samples necessary converge to accurate estimate scales unfavorably with dimensionality. In practice, existing techniques can reliably estimate MI in up to tens dimensions, but fail in higher dimensions, where sufficient sample sizes infeasible. Here, we explore the idea that underlying low-dimensional structure in high-dimensional data be exploited to faithfully approximate MI in -ishdimensiooall settings with realistic sample sizes. We develop method that we call latent MI (LMI) approximation, which applies nonparametric MI estimator low-dimensional representations learned by simple, theoretically-motivaee model architecture. Using several benchmarks, we show that unlike existing techniques, LMI can approximate MI well for variables with > 10xh dimensions if their dependence structure captured by low-dimensional representations. Finally, showcase LMI on two open problems in biology. First, we approximate MI between protein language model (pLM) representations of interacting proteins, and find that pLMs encode non-trivial information about protein-proteie interactions. Second, we quantify cell fate information contained in single-cell RNA-seq (scRNA-seq) mea- surements of hematopoietic stem cells, and find sharp transition during neutrophil differentiation when fate information captured by seRNA-seq increases dramati- cally. An implementation of LMI is available at latentmi. readthedocs. io.
Introduction
Mutual information is a universal dependence measure which has been used to descrihe relationships between variables in wide variety of complex systems: developing embryos [1] artificial neural networks (2). flocks of birds (3) and more. Its widespread use can be attributed to at least two of its appealing properties: equitability and interpretability.
Many dependence measures are inequitable, meaning that they are biased toward relationships of specific forms [4. For example, Pearson correlations quantify the strength of linear relationships, and Spearman correlations quantify the strength of monotonic relationships. Inequitability can be particularly problematic for complex systems, where relationships can be nonlinear, non-monotonic, or involve higher-order interactions between multidimensional variables [5]. Mutual information (MI) stands out equitable measure that capture relationships of any form, and generalizes across continuous, discrete, and multidimensional variables [6]. And when scaled consistently. MI provides universal currency in interpretable units which be understood the number of 'bits'
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

of information shared between variables 77. MI can also be interpreted through decomposition into pointwise mutual information (pMI) 8 9 which attributes dependence to specific pairs of values.
MI can be defined as the Kullback-Leibler divergence, DKL, of a joint distribution from the product of its marginals. For absolutely continuous random variables X, Y defined over X, y, with joint distribution Pxy and marginal distributions Px, Py
I(X;Y) (xx(Pxx||Px & Py) J (xr(z,)) log Pxy v) (x
dydr
(1)
In practice, Pxy is often unknown, and Y( must be estimated from observations { B:)
that sparsely sample Pxy. While nonparametrie MI estimators have been remarkably successful for variables with single dimension [TO MOLE 13] MI estimation for high-dimensional variables
remains a significant challenge. Nonparametric MI estimation suffers from the curse of dimensionality accurate estimation requires number of samples that scales exponentially with the dimensionality of the variables
14
An exciting recent approach to scaling MI estimates to high dimension is the use of variational bounds on KI. divergence to reduce the MI estimation problem to a gradient descent optimization problem 15 16. MI estimators based on variational bounds indeed empirically perform well for data with ones to tens of dimensions [12]. but still suffer from the curse dimensionality [[4], and can exhibit high variance [7] To our knowledge, techniques have been shown to reliably estimate MI in practice for variables with hundreds or thousands dimensions - a regime relevant to many fields,
including genomics, neuroscience, ecology, and machine learning [88 19 5 20
More generally, has been shown that no technique can accurately estimate MI from finite samples without making strong assumptions about the distribution from which samples are drawn (1]. resulting in a fundamental tension between the theoretical appeal of MI and the practical difficulty of its estimation. One way to resolve this tension is to develop alternative measures of statistical dependence [14 (2), which retain desirable properties of MI, but are tractable to estimate. Sliced random low-dimensional linear projections ("slices") of nigh-ddmensiooal data, is a notable example of such an approach [14] 23]. While sliced MI is useful for many problems [24] 25 it does not retain the interpretability (in bits) of classical MI, and is inequitable, as it quantifies information that can be extracted through linear projections [4].
MI, which is the average of MI estimates
Here, we take complementary approach to sliced measures. Rather than considering alternatives to classical MI, we ask if it is possible to make strong, yet reasonable, assumptions about data which enable tractable MI estimation. In this work, we consider the empirically supported assumption that complex systems have underlying low-dimensional structure [5]. The usefulness of this assumption relies on our ability to identify low-dimensional structure in data. We will explore methods for learning now-dimensional representations which are useful for MI estimation and highlight examples where the methods can still fail.
Specifically, we propose latent mutual information (LMI) approximation, which applies a non- parametric MI estimator to mutually informative compressed representations of nigh-dimensional variables. To learn such representations, we design a simple neural network architecture motivated by information -theoreticct principles. We demonstrate, using synthetic multivariate Gaussian data, that LMI approximation can be stable for variables reaching thousands of dimensions, provided their dependence has low-dimensional structure. We then introduce an approach for resampling real data to generate benchmark datasets of two high-dimensional variables where ground truth mutual information is known. Using this approach, we evaluate the ability of LMI to capture statistical dependence in two types of real data: images and protein sequence embeddings. Finally, we apply LMI to two open problems in biology: quantifying interaction information in protein language model embeddings, and quantifying cell fate information in the gene expression stem cells.
2
Approach
Our goal is to approximate MI from digh-diiensional data using low-dimensional representations which capture dependence structure. Our specific approach is to use neural networks to map variables X, Y to low-dimensional representations Za, Zyy Then, we use the well-established nonparametric
MI estimator introduced in [10] to estimate [Zzi Zy).

---

X
Y
y
2
Z
27
Z
X
Y
LMI(X; ÍKsG (Zzi zy)
Figure 1: Workflow of latent MI approximation a) Embed high-dimensional data in low-
dimensional space such that mutually informative structure is preserved. b) The KSG estimator [[0] is used to estimate MI by averaging over pointwise MI (pMI) contributions.
The central challenge here learn Ze, Zy such that f(Z:: 2,) = (() One sensible approach
would be to use autoencoders [26] or other popular nonlinear dimensionality reduction techniques [27] 28 to compress each variable separately. While such an approach could yield a good approximation if compression is perfectly lossless, it can result in poor approximation if compression is lossy. An illustrative example is two variables each with thousands of independent dimensions but single pair of strongly dependent dimensions two separate lossy compressions unlikely to preserve the rare dependent components.
We can make this intuition precise using properties of entropy and MI 77.. For simplicity, let us consider an approximation using only one compressed variable, Zy = f(Y). By rewriting MI in terms of differential entropy, denoted h. we see that for absolutely continuous X, Y with finite differential entropy,
- M(X) M(XIV) A(X) M(XIZ)) M(XIZ,) M(XI))
(2)
In the case of lossless compression, (x) = 1(X(Y), so
- - = (( - h(X()) =0
(3)
However, if compression does not perfectly preserve information, it is possible that A(X|Z,) h(X[Y). Since h(XlY) is intrinsic to the data and independent of learned representations, minimiz- ing Zy.) is equivalent minimizing 1(X; L(X; 2yy. This points to approach to learn representations suitable for approximating (X;Y):: regularizing pair of autoencoders to learn compressed representations Z, f(X) and Zy g(Y)) while minimizing h (((lZy) and h(Yl
Because directly minimizing conditional entropies intractable, we instead minimize a convenient proxy, which is the mean-squared error (MSE) loss of networks that predict one variable from another. The connection between conditional entropy and reconstruction loss long been appreciated as way to interpret autoencoders through an thforratioo-theorrtiic lens (29 (30. Here, we observe that this connection can be applied to learn representations which lend themselves to MI estimation. We explicitly show that minimizing cross- prediction loss from Zμ to Y is equivalent to minimizing upper bound on the conditional entropy (Y (2,) in Appendix 1. Theorem 1.
Applying cross- prossppredi regularization to pair of autoencoders results in a network architecture (Fig. 1) with one encoder for each variable and four decoders which reconstruct each variable from each latent code. We train the networks by minimizing the sum of the MSE reconstruction loss for each decoder. More precisely, for variables X, with dimensionality dx,dy. we optimize encoders Ex, Ey, and decoders D Dxx, DXY, DYY, Dyx to minimize LAEC LAF + Lc. where
1 LAB
=ElIIX (Ex dy
EilyY Dyy
(4)
dx
1 Lc
1 - Dxy
E| X (yx(EE
(5)
dx
dy
3

---

This is not the only way one could regularize autoencoders to preserve mutually informative structure. We design and empirically characterize some alternatives in Appendix A.2. We find that multiple regularization approaches can be effective, but cross- pressicodiitionn comes with unique benefits. For example eross-predictive networks can be dissected to attribute high-dimensional MI estimates to specific dimensions, as demonstrated in Appendix A.2.3.
While the specific architecture of encoders and decoders could be carefully chosen for each estima- tion problem (e.g. convolutional layers for image data), here we use multilayer perceptrons with a priori determined hidden layer sizes for all problems. This is intentional: a useful MI estimator should not need extensive parameter selection. Every LMI estimate shown in this paper (exclud- ing Appendix and Section 3.3) uses the default parameters of our library. equivalent to running 1mi. estimate (X_ [a_samlle,, Y samples).
To ensure that optimizing based cross-reconstruccion does not introduce spurious dependence due to overfitting, we learn representations and estimate MI on different subsets of the data. That is, for N joint samples, we train the network using subset of N/2 samples, then estimate MI by applying the estimator of [TO] to latent representations of the remaining N/2 samples. A high-level overview of MI estimate using LMI approximation is given in Algorithm 1.
We also state and prove some basic properties of LMI approximation in Appendix 1.3, Theorems and 3, namely that (Zz: Zy) s (X:Y), and that I (Zzi Zy) 0f (X;Y) = 0.
Algorithm Estimating MI using LMI Approximation
Require: N joint samples {z 31) N of random variables X,
Require: Encoders Ex, Ey, decoders Dxx..... Dyy parameterized by .......,6 randomly split into two subsets of N/2 samples, Duain Dest
optimize 01,...,6, 0g to minimize LAEC on Drain
encode Dest using E1, Ez, 01, 02, yielding {(Zg, 2 N/2 return
N/2
)
3 Empirical evaluation
Next, we empirically study the effectiveness of LMI approximation with 16 dimensional latent space (8 dimensions per variable), in comparison with three popular estimators: the nonparametric estimator from [10), referred to as KSG, and the variational bound estimators from [5] 31 referred to as MINE and InfoNCE, respectively (implementation details in Appendix A4).
3.1 Evaluating mutual information estimators on synthetic data
We first consider the problem of MI estimation between multivariate Gaussian distributions, because ground truth MI can be analytically computed, and dimensionality can be easily tuned. We consider the scalability of MI estimators with increasing dimensionality two kinds: the ambient dimension- ality of the data, denoted d, and the intrinsic dimensionality of the dependence structure, denoted k. We benchmark the performance of estimators in the regime of high ambient dimensionality and low intrinsic dimensionality. Specifically, consider variables with d = 10 to d 5 10 ambient dimensions and k = 1 to k = 9 dimensional dependence structure.
To generate samples from two d-dimensional random variables X, Y with k-dimensional dependence structure. we sample multivariate Gaussians with a covariance matrix of a prescribed from. In particular, we enforce pairwise correlation p between & components of each variable, with redundant duplicate components as well as independent components constituting the remain dl - & dimensions. [[X; Y) can then be easily computed analytically from the covariance matrix. The exact sampling procedure we use for these experiments is given in Appendix 1.3.1, Algorithm 4.
Results of benchmarking MI estimators using these synthetic datasets are given in Fig. 2. For estimates from N = 10 samples, we find that, as expected, the performance of existing estimators degrades with d, with near complete failure for d >> 100 (Fig. 2a-c, 2f-h.). In contrast, applying LMI approximation results in stable estimates up to = 103 ambient dimensions (Fig. 2d, 21). The faithfulness of LMI approximation instead degrades with increasing k. Nonetheless, LMI

---

INFONCE
b)
LHI
IIFONCE
9)
UNI
30
30
MSE
lowes: MIE
50
50
KT
hggest Kendal
400
800
LMI
1200
1200
190
5000
5000
1.23
30
5000
5500
50
Intrissic
Intrisiic
1.00
200
K3G
Intrinsic
200
150
360
4)
KSG
d)
MINE
3 500
100
h)
KSG
1)
HINE
-.4
560
0.50
10-
0.23
2000
50
50
52
2000
440--
400-
4395
5300
1000
400
4020
5000
5000
1.5.7% 7
5000
3
6
5200
3.5
5000
Intrraii: dimmtssors
3 3
005
Figure 2: MI estimator performance scaling with increasing dimensionality. a) d) Absolute accuracy measured by mean-squared error over 10 estimates per setting, with ground truth MI between 0 and 2 bits, and 2 10 samples per estimate (1:1 training split for neural network- based estimators). e) Estimator with highest absolute accuracy in each setting. Ties broken randomly. f) i) Relative accuracy measured by Kendall J rank correlation of estimates with ground truth. j) Estimator with highest relative accuracy in each setting. Ties broken randomly.
approximation gives more absolutely and relatively accurate MI estimates than alternatives for 83% and 87% of tested settings respectively (Fig. 2e, 2j). We also find that these benchmarking results remain qualitatively similar for multivariate Gaussian data after various Ml-preserving transformations (Appendix A.3.4), similar to the benchmarking approach of [12].
3.1.1 Empirically quantifying convergence rates of MI estimators synthetic data
The principle enabling the scalability of LMI approximation that the number of samples requires to converge is limited by & rather than d when k < d. We empirically demonstrate this by quantifying the convergence rates of MI estimators on the synthetic Gaussian datasets described above. We generate datasets with sample numbers in N E [102, 10*), and ambient dimensionalities in d E [0), each with single correlated dimension between variables (k- 1), and bit MI. For each estimator and each ambient dimensionality d, we empirically determine the number of samples required to achieve I(X,Y) i(x, Y) < bits, with linear interpolation between tested sample numbers.
As expected, methods that do not explicitly learn low-dimensional representations (InfoNCE, MINE, KSG) require increasing numbers of samples to estimate MI with error below = 0.1 (Fig. 3a). KSG fails to estimate MI for 2 13 for any N, while MINE and InfoNCE scale slightly better, failing for d 2 25 and d 2 37 respectively. In contrast, the sample requirements of LMI remain qualitatively stable no more than 4- 103 samples are necessary for an accurate estimate.
While the convergence behavior of LMI is mostly unaffected by varying d, it is sensitive to varying k. When the same experiment performed with increasing numbers of correlated dimensions the limit where k: = d, the convergence behavior of LMI is no longer favored over other estimators (Fig. 3c). The performance of all estimators dramatically decreases with k: such that larger error tolerance must be chosen for informative convergence estimates. The dependence of variational bound estimator convergence on k is. to our knowledge, not explained by existing theory (16] 32]. In the intermediate case of k [0.1 d] (Fig. 3b), we find that LMI convergence is fast with low k, but becomes slow k grows, nonetheless remaining favorable compared to other estimators.
3.2 Evaluating mutual information estimators on resampled real-world data
While the empirical results on multivariate Gaussians reassuring, they are not representative of performance on real data, where low intrinsic dimensionality is not known a priori, and distributions can be non-Gaussian. To better understand the behavior of LMI in more realistic settings, we introduce technique for creating benchmark datasets by resampling real-world data. Briefly, we correspondences between discrete labels and complex data (i.e. digit labels and digit images in MNIST) to transform simple discrete distributions into realistic high-dimensiooal distributions. A similar approach is explored in a concurrent benchmarking study (33.
Specifically, we draw samples from a bivariate Bernoulli vector, [Tm, Tyy 132 with prescribed pairwise correlation Cor( Lax Ly) = p. where each value corresponds to discrete label
5

---

6-0.1
Samp recdod
10000
3000
[=01
10000
8000
c+0
Sampres rnpded
10000
8000
<-00
Edi inlinooor
USS
0000
6000
5000
4000
4200
4000
InfONCE
2000
200n
2000
10
30 Ammilent dimennsinns
20
40
50
10
20
30
40
50
0
10
20 30
50
Ambient dimersions
Arrikent dimernions
Figure 3: Number of samples required to achieve Y) ((x)-((x))))) ((), < a) Data with low-rank dependence structure, with e 0.1. b) Moderate-rank dependence structure, with < 0.2. c) Full-rank dependence structure, with 0.4. "+" marker indicates that N > 10a samples are required for accurate estimates for all larger d. Ground truth MI bit in all cases.
of a set of samples in high- sigh--imensional dataset (e.g. 0 and correspond to images of Os and Is in MNIST). For each sample of L, we replace each component with random (without replacement) high-dimensional sample matching the label. For the example of MNIST, this transforms samples from into pairs of images of Os and 1s, represented as samples of random vectors X, Y. E TR7744
Under the assumption that discrete labels can be uniquely identified by high-dimensional vectors, high- dimensional MI is identical to the discrete label MI. That is, assuming HL,, X) (T.,, 0 then 1(X; Y = I (Lz; Lyy) (shown in Appendix A.3.2, Theorem 4). And using our knowledge of [[2x; Lyy can be analytically computed.
We resample two different source datasets: (1) "binary" subset of MNIST, containing only images of Os and 1s, with 5000 samples and 784 dimensions and (2) embeddings of subset of protein sequences from E. coli and A. thaliana proteins, with 4402 samples and 1024 dimensions. For both source datasets, validate the I(L.i L.) = YY) approximation in Appendix A.3.3.
For each source dataset, we generate 200 benchmark datasets with true MI ranging from 0 to bits. We estimate MI on each dataset with each estimator (Fig. 4a, 4b), and quantify absolute accuracy (MSE), relative accuracy (rank correlation with ground truth), and runtime for each estimator (Fig. 4c). For both types of source data, we find that variational bound estimators have high variance, in line with previous observations [6]. On protein embedding datasets, variational estimators nearly always fail to estimate nonzero values resulting in rank correlation below O for InfoNCE. The KSG estimator, while achieving high relative accuracy, systematically underestimates MI. resulting in low absolute accuracy. Furthermore, the amount by which it underestimates true MI is different between the two datasets indicating inequitability. contrast, LMI approximation yields estimates consistently close the ground truth, with high relative and absolute accuracy.
3.3 Constructing and studying problems where LMI fails
While LMI performs well on diverse set of benchmarks, it is not infallible. LMI can fail when (1) its representations do not capture dependence structure, or (2) when KSG fails to accurately estimate MI in latent space. Because the limitations of KSG are well-studied [34,, here we focus failure mode (1) and identify problems where LMI learns representations which result in poor MI estimates.
LMI trivially fails to learn useful representations when the dimensionality of the latent space is too small to capture the dependence structure of the variables (e.g. in Fig. 3c). This limitation can be partially overcome by evaluating LMI with a latent space large enough to capture dependence structure, however a priori knowledge of the appropriate embedding dimension is rarely possible. One heuristic approach (Appendix A.5.2) is to make estimates with several latent space sizes, and choose the size which maximizes the estimate.
LMI accuracy can also degrade when MSE cross- pross-prediction loss fails to constrain latent represen- tations, reducing the LMI model to pair of independent autoencoders. This can happen when certain symmetries are present in the data, such as variables X, Y for which E[X Y = E[X] and MSE-minimizing
T(X; Y) 2 0, two examples of which are shown in Fig. 5a. For such variables, predictor of X becomes independent of Y, and cross-prediction loss no longer has useful regu- larization effect. In these cases, if independent autoencoders are insufficient to capture dependence structure, LMI will fail to accurately estimate MI.

---

a)
MINITT benchmarking
b)
ProtTransS benchmarking
c)
1.0
Estimetor
Estimator
LMI
KSG
MINE
InfoNCE
LNI
0.8
LMI
(MNISST)
KSGG
KSGG
0.8
Rank (MNIST)
0.95
0.92
0.61
0.59
0.5
InfONCE
0.4
InfoNCE
Rundime (MNIST)
7.9
39.7
12.18
12.01
0.4
MSE
0.01
0.32
0.33
0.2
Rank con IP15)
0.94
0.92
0.26
-0.06
0.0
0.0
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
Labal MI
Label KI
Figure 4: Performance of MI estimators on resampled real datasets. a) Estimates on resampled pairs of MNIST digits, with 5 10 samples and 784 dimensions. b) Estimates on resampled pairs of ProtTrans5 sequence embeddings, with 4.4 10 samples and 1024 dimensions. e) Statistics of estimator accuracy and runtime (in seconds), for each dataset type.
Nolsy circle
Gaussian mixture
0.8
Estimatar
0.6
Escimese
0.4
KSG
0.4
0.2
0.2
InfONCE
0.0
0.0
X
Exclusive dimersions
Exclusive dimensions
Figure 5: Constructing problems where LMI fails. a) Examples symmetric distributions where
[[X[Y]= = [[[ b) MI estimates from 10 samples of symmetric variables exclusive (independent normal) dimensions are added. Ideal estimators are invariant. LMI latent space is 1D per variable.
We construct benchmark to illustrate this limitation and understand if other estimators suffer from the same limitation in practice. In brief, we generate samples from variables with single pair of symmetric dimensions and increasing numbers of independently normally distributed dimensions. The exact sampling procedure is given Appendix A.3.5. An ideal estimator should not vary with the number of independent dimensions as true MI is invariant. In the case with no independent dimensions, we expect LMI be accurate up the limitations of KSG because independent autoencoders are sufficient to learn mutually informative representations. As the number of independent dimensions increases, LMI estimates should quickly degrade.
As anticipated, LMI estimates implemented with an MSE cross-prediction regularization decay quickly (Fig. 5b). However, this behavior is not unique to LMI: all studied estimators decay near with fewer than 10 dimensions. In this dimensionality, MINE and InfoNCE typically perform well [1], so their failure not due merely the dimensionality of the problem but also due to the nature of the distributions. The estimators do not agree even in the 1D case without independent dimensions, suggesting that symmetric distributions may be generally problematic for MI estimation, similar to long-tailed distributions [12]
Because this failure of LMI specifically arises from artifact of the MSE loss, one might conclude avoiding MSE cross-prediction regularization may be sufficient to improve LMI estimation. We explore this possibility in Appendix A.2.4.
4
Applications
4.1
Quantifying interaction information in protein language model embeddings
Pretrained protein language models (pLMs) have recently seen widespread use, largely due to their convenient representations of protein sequences (vectors R- typically with N Rs 103) which can be used for transfer learning on downstream tasks 35
While it is known that pLM sequence embeddings contain significant information about protein structure [38], not clear how well existing pLMs encode functional information. Recent work has shown that pLMs fail to capture
36
37
7

---

Kinase Kinase-target dependence
d)
InfoNCE
Ligand-receptor receptor classification
Ligand
Kinase target classification
Extimatorr
Shuffle
MINE
Unshuffied
LMI KSG
0.6
0.5
Estimate
Ligand receptor dependence
Estimator
InfoNCS
Shuffied
0.4
0.4
Unshuffled
0.2
0.2
LMI
0.74
TPR=FFRR
TPR=F9R
0.5 0.8 1.0
Figure 6: Quantifying dependence between participants of protein interactions. a) b) MI estimates between interaction partners, compared to randomly permuted data. c) - d) ROC curves of density ratio classifier distinguishing annotated interacting pairs from unannotated "negative" samples, for all pairs of 170 held-out proteins. Averages over 20 random hold-out splits.
some important functional properties, such thermostability [39]. It remains unclear the extent to which pL. embeddings contain information about protein-proted protein interactions, which are essential to protein function. Here, contained in (02--iimensi sequence embeddings from the ProtTrans5 model [36].
information-theoreticc approach quantify interaction information
We study two types of protein-prooe protein interactions: kinase-targes target and ligand-receptor interactions. For both, the OmniPath database [40] contains lists of thousands annotated pairs of interacting proteins. We consider each annotated pair to be sample from joint "interaction" distribution over pairs of sequence embeddings. For example, for kinase-targer interactions, we consider kinase and target sequence embeddings random variables K E 12004 T 1024 with joint distribution PKt (k, t). Then, using joint samples {(k), t.)..... ... (KN (xN,t)), (N)] tion partners, TT) We analogously estimate 1 (L; R) for ligand-receptor interactions.
} we estimate mutual information between interac-
If pLM embeddings capture interaction information, MI between interaction partners should be significantly above 0 bits. Applying LMI approximation, we estimate [(L; R) = 2.8 bits and TT: 0.8 bits. To test the significance of these values, we estimated MI from shuffled data, and found T K;T) ss [L; Rt) x 0 (both mean and standard deviation < 0.05), across 20 random shuffles. These results indicate that PL. embeddings contain information about both types of interactions, and contain more information about ligand-recepton interactions than kinase- target interactions. In contrast, existing estimators yield far lower estimates, with MINE estimates indicating independence for both types of interactions (Fig. 6a, 6b). To validate the LMI estimates of dependence, we next operationally verify the presence interaction information.
If contein-prroein interactions can be predicted for set of held-out proteins based on sequence embeddings, then sequence embeddings must contain interaction information. To see if this the case, we extend LMI to predict protein interactions from sequence embeddings. For prediction, our goal is to predict whether a held-out pair of sequence embeddings (1, is an annotated P.alle))
ligand-receptor pair. One way to do this is estimating the log density ratio log ((())-Px())' and setting a threshold above which sequence pairs are predicted to be annotated interactions.
We make simple modification the KSG estimator yield the desired density ratio estimates (given in Algorithm 2). and use these estimates (with latent approximation) to predict interaction annotations. For 20 different random splits of 170 held out proteins, we use density ratio estimates to classify all 2.89 10 pairs of held out proteins as interacting or non-interacting. The receiver operating characteristic (ROC) curves for predictions of both interaction types are shown in Fig. 6c, 6d, with mean AUC-ROC scores of 0.87 and 0.74 for ligand-receptor and kinase-target interactions respectively. These results demonstrate that protein interactions can be predicted better than random chance using ProtTransS embeddings, suggesting that pL embeddings capture information about kinase-target and ligand-receptor interactions. And in line with the LMI estimates, ligand-receptor interactions are better predicted than kinase-target interactions.
8

---

Algorithm 2 k-nearest neighbor log density ratio estimator
Require:
joint samples {(3r4 91))
Require: query point (4a Qy)
let (ra,Ty) < k-th nearest neighbor sample of (9x, qy) in joint space (default 3)
W compute Chebyshev distance
let + 11(9. qy) Ty)loo
let de < 0, R.y < 0
count neighbors within d in marginal spaces
for each (I;, Vi) do
11 - xi 80 < then
n1. += 1
end if
if 11qu y/ 10o
< d then
By += end if
end for
W return estimate of log pl)
pla.)p(a))
return (k) + (NN) - (nx)) (my)). where vb is Digamma function
4.2 Identifying cell fate information in hematopoietic stem cells
Single cell RNA sequencing (scRNA-seq) measures the expression of 10 genes in single cells, which can be thought of samples of a gene expression state variable X E R These samples can be used to infer probability distribution over gene expression states, Px. To study the dynamics of gene expression, one approach is to make measurements at multiple timepoints .,,,, which can be thought of as samples of random variables X; E IR9, Lineage tracing is a technique where clonally related cells can be labelled with barcodes, allowing sampled cells from different timepoints to be matched with their "twins" from another. When combined with scRNA-seq. lineage tracing can be thought to provide samples from the joint distribution Px.....Xy
One fundamental question about cellular dynamics is whether the time evolution of gene expression state is dependent entirely on the current gene expression state. That is, if the "fate" of cell can be formally modeled as a Markov chain X; + X In some cases, cell behavior may be a function of hidden variables resulting in man-vaako dynamics. Using the data processing inequality (DPI) 77.. we know that if gene expression dynamics are Markovian, I(X;; X++1) 2 I(X;; Xi+2), and if DPI does not hold, then gene expression states are non-Markovian. This can, in principle, be tested using samples from the joint distribution Px.....Xx. Due to the difficulty of high- high-diimnsiooal MI estimation, previous work has used heuristic alternatives to indirectly test the Markovian assumption for lineage-traced scRNA-seq (LT-seq) data 42). Here, we will use LMI explicitly estimate high-dimensional MI from LT-seq data, and test the Markovian assumption for gene expression states.
We study previously published LT-seq data set of in vitro differentiating mouse hematopoietic stem cells 42). The dataset includes sister cells which are separated at day 2 of the experiment and allowed to differentiate in separate wells until day 6, with cell states sampled on both days 2 and 6. Under the Markovian assumption, this can be modeled as X2 + X6. X₂ - Xee where X the state of the cells on day 6 in the second well. By the DPI we should have T(X2; Xs) 2 7(Xa; Xe:). Using LMI, we estimate that 1(X2; Xτ) s 0.31 + 0.02 bits and [(6;; X Rs 0.98 + 0.01 bits (mean + SEM), over 20 random pairings of clonally related cells. This indicates that gene expression states are non-Markovian, in line with prior findings 42
Cell fate information manifests in transcriptomes sometime between days 2 and 6. By decomposing
LMI estimates into pointwise contributions (Appendix A.4.3, Algorithm 5), we can determine precisely when this hidden information emerges. Generally, we see that pMI increases along differentiation trajectories (Fig. 7b). Along the neutrophil trajectory, we quantatively compare cell fate information with neutrophil pseudotime computed by graph smoothing [42,, and find that PMI begins to rapidly increase around pseudotime value of 35 103, which is roughly aligned with the transition from granulocyte-myelee progenitor promyclocyte, as defined in 42).

---

pkj
boe
-14
Undiffrentsted
Monteyle
-12
0.8
Neuropphl
-10
Beso
0.6
Mog
Host
Lymphoid
0.2
Eos
Cer7_pC
2.0
0.0
10000
50000 60000
Keutrophil preudtimm
Figure 7: Quantifying cell fate information in single cell transcriptomes. a) 2D SPRING embed-
ding 44 of lineage-traced single cell RNA-sequencing data from 42 b) Pointwise decomposition of MI between sister cells across timepoints, as estimated using LMI. Hue applied to early timepoint cells, late timepoint and unbarcoded cells in grey. Star indicates neutrophil pseudotime value 35 10 <) Smoothed PMI (rolling average over 100 cells) across neutrophil differentiation trajectory. Vertical line denotes neutrophil pseudotime value 35 103.
Discussion
In this paper, we tested the hypothesis that low-dimensional structure can enable scalable estimation of MI from high-dimensional data. We introduced LMI approximation, which applies a nonparametric MI estimator to low-dimensional representations learned by neural networks. We quantified the effectiveness of LMI approximation, using multiple approaches spanning over 3000 benchmark datasets. Our results suggest that, unlike existing techniques, LMI approximation is generally effective for high-ding loeb-dimensional data with www-iimnsional structure, even if the number of available samples remains relatively low a regime where many real datasets reside [5]
We used LMI to study two open problems in biology. We show one example where LMI enables the use of information-theoree study [44]. the authors had indeed wished to estimate MI but were unable to do so and resorted to heuristic approaches. LMI may similarly help identify dependence in cellular dynamics in other systems 43 45 46]. In a different subfield of biology, we showed an example of using LMI to quantify functional information learned by pLMs. Our results suggest that nontrivial protein -preein interaction information is learned by ProtTrans5, motivating the development of interaction prediction tools based on pLMs. As the number of large pLMs grows [35 37
ideas to study the dynamics of gene expression. In the original
36
information-thcoret theoretic
approaches using LMI could help benchmark models.
Limitations The most prominent limitation of LMI follows directly from its motivating assumption: will fail to quantiffy dependence not captured by the learned representations. As result, it is easy to design synthetic datasets on which LMI will fail (see Fig. 3c). However, it is likely that many real datasets (beyond those explored in this paper) will be amenable to LMI approximation, as there is strong evidence that complex systems generally have low intrinsic dimensionality [5]. Our implementation of LMI approximation also inherits some limitations of the KSG estimator, notably that it fails for strongly dependent (near deterministically related) variables [32 34 To overcome
this, it may possible to apply previously developed corrections (34]. Finally, despite reassuring empirical results, few theoretical properties of LMI have been derived. This is an important line of future work, which could help precisely identify settings where LMI approximation is effective.
Broader impacts MI estimators have been used quantify moral and legal fairness 25 LMI approximation is not universally faithful, and more generally no MI estimator can be universally accurate (32. MI estimates must be interpreted with great care when applied human lives. The experiments (and pilot iterations) in this paper were performed on single NVIDIA RTX 3090, and resulted in estimated 45.36 kg CO2eq. Estimates made using [50].
Code availability The code necessary to reproduce all results from this paper are available at https: //i thub. com/ hwtp:///ggggcaaaaaaaetaatmmaaaaa intommoottiiiiiioooooooooiooooo The 1mi Python package can be found at https:: //tpts::///ttbbbb/t hwtpstttu..ommmttttttttt wwwp////////aa/anmam..io.mmo and documentation is hosted hhttps //latentmi readthedocs.
10

---

Acknowledgments and Disclosure of Funding
We thank Caroline Holmes, Ninning Liu, Sean McGeary, and Pippa Richter for thoughtful discussions. This work is supported by funding from NIH Pioneer Award DPIGMI33052, R01HG012926 to P.Y.. and Molecular Robotics Initiative the Wyss Institute.
Author contributions G.G. conceived of, designed the study, developed the software, analyzed the data, and wrote the paper. X.L. contributed study design. A.M.K. contributed to the design of the study, analysis of the data, wrote the paper, and supervised the study. P.Y. supervised the study. All authors reviewed, edited, and approve the paper.
References
[1] Julien O Dubuis, Gasper Tkacik, Eric F Wieschaus, Thomas Gregor, and William Bialek. Positional information, in bits. Proc. Natl. Acad. Sci. U. S. A., 110(4)):16301-1630 October 2013.
[2] Ravid Shwartz-Ziv and Naftali Tishby. Opening the black box of deep neural networks via information. arXiv [cs.LG], March 2017.
[3] Matthijs Meijers, Sosuke Ito, and Pieter Rein Ten Wolde. Behavior of information flow near criticality. Phys Rev E, 103(1): January 2021.
[4] David Reshef, Yakir A Reshef, Hilary K Finucane, Sharon R Grossman, Gilean Vean, Peter J Turnbaugh, Eric S Lander, Michael Mitzenmacher, and Pardis C Sabeti. Detecting novel associations in large data sets. Science, 334(6662):1118-1222 December 2011.
[5] Vincent Thibeault, Antoine Allard, and Patrick Desrosiers. The low-rank hypothesis of complex systems. Nat. Phys., pages 1-9, January 2024.
[6] Justin B Kinney and Gurinder S Atwal. Equitability, mutual information, and the maximal information coefficient. Proc. Natl. Acad. Sci. U. S. A., 111(9):335-3355) March 2014.
[7] Thomas M Cover and Joy A Thomas. Elements of Information Theory. Wiley & Sons, Incorporated, John, 2006.
[8] Yao-Hung Hubert Tsai, Han Zhao, Makoto Yamada, Louis-Philippe Morency, and Ruslan Salakhutdinov. Neural methods for point-wise wise dependency estimation. arXiv [cs.LG], June 2020.
[9] Xianghao Kong, Ollie Liu, Han Li, Dani Yogatama, and Greg Ver Steeg. Interpretable diffusion via information decomposition. arXiv [cs.LG], October 2023.
[10] Alexander Kraskov, Harald Stöghauer, and Peter Grassberger. Estimating mutual information. Phys. Rev. Stat. Nonlin. Soft Matter Phys., 69(6 Pt 2):066138, June 2004.
[11] Caroline Holmes and Ilya Nemenman. Estimation mutual information for real- valued data with error bars and controlled bias. Phys Rev E, 100(2-1):02244)) August 2019.
[12] Pawet Czyž, Frederic Grabowski, Julia E Vogt, Niko Beerenwinkel, and Alexander Marx. Beyond normal: On the evaluation mutual information estimators. arXiv [stat.ML], June 2023.
[13] Thalia Chan, Michael Stumpf, and Ann C Babtie. Gene regulatory network inference from single-cell data using multivariate information measures. Cell Syst, 5(3):251-267.c) September 2017.
[14] Ziv Goldfeld and Kristjan Greenewald. Sliced mutual information: A scalable measure of statistical dependence. arXiv [cs./T], October 2021.
[15] Mohamed Ishmael Belghazi, Aristide Baratin, Sai Rajeswar, Sherjil Ozair, Yoshua Bengio, Aaron Courville, and Devon Hjelm. MINE: Mutual information neural estimation. arXiv [Gs, January 2018

---

[16] Ben Poole, Sherjil Ozair, Aaron van den Oord, Alexander A Alemi, and George Tucker. On variational bounds of mutual information. arXiv [cs.LG], May 2019.
[17] Jiaming Song and Stefano Ermon. Understanding the limitations of variational mutual informa- tion estimators. arXiv [cs.LG], October 2019.
[18] Allon M Klein, Linas Mazutis, Ilke Akartuna, Naren Tallapragada, Adrian Veres, Victor Li, Leonid Peshkin, David A Weitz, and Mare W Kirschner. Droplet barcoding for single-cell transcriptomics applied to embryonic stem cells. Cell, 161(5):1187-1201, May 2015.
[19] Surya Ganguli and Haim Sompolinsky. Compressed sensing, sparsity, and dimensionality in neuronal information processing and data analysis. Annu. Rev. Neurosci, 35:485-508, April 2012.
[20] Ziv Goldfeld, Ewout van den Berg. Kristjan Greenewald, Igor Melnyk, Nam Nguyen, Brian Kingsbury, and Yury Polyanskiy. Estimating information flow in deep neural networks. arXiv [cs.LG], October 2018.
[21] David McAllester and Karl Stratos. Formal limitations on the measurement of mutual infor- mation. In International Conference Artificial Intelligence and Statistics, pages 875 884. PMLR, June 2020.
[22] Danqi Liao, Chen Liu, Benjamin W Christensen, Alexander Tong, Guillaume Huguet, Guy Wolf, Maximilian Nickel, Ian Adelstein, and Smita Krishnaswamy. Assessing neural network representations during training using noise-resilient diffusion spectral entropy, In 2024 58th Annual Conference on Information Sciences and Systems (CISS), pages 1-6. IEEE, March 2024.
[23] Ziv Goldfeld, Kristjan Greenewald, Theshani Nuradha, and Galen Reeves. k-sliced mutual information: A quantitative study of scalability with dimension. arXiv [cs.IT], June 2022.
[24] Dor Tsur, Ziv Goldfeld, and Kristjan Greenewald. Max-sliced mutual information. arXiv [cs.LG], September 2023.
[25] Yanzhi Chen, Wei-Der Sun, Yingzhen Li, and Adrian Weller. Scalable infomin learning. Adv Neural Inf. Process. Syst., abs/2302. 10701, February 2023.
[26] Hinton and Salakhutdinov. Reducing the dimensionality of data with neural networks. Science, 313657866:504-507,, July 2006.
[27] Etienne Becht, Leland McInnes, John Healy, Charles-Antoine Dutertre, Immanuel W H Kwok, Lai Guan Ng, Florent Ginhoux, and Evan W Newell. Dimensionality reduction for visualizing single-cell data using UMAP. Nat. Biotechnol., December 2018.
[28) Kevin R Moon, David van Dijk, Zheng Wang. Scott Gigante, Daniel Burkhardt, William Chen, Kristina Yim, Antonia van den Elzen, Matthew J Hirn, Ronald R Coifman, Natalia B Ivanova, Guy Wolf, and Smita Krishnaswamy. Visualizing structure and transitions in high- dimensional biological data. Nat. Biotechnol., 148-:1944221142, December 2019.
[29] Paseal Vincent, Hugo Larochelle, Yoshua Bengio, and Pierre-Antoine Manzagol. Extracting and composing robust features with denoising autoencoders. In Proceedings of the 25th international conference on Machine learning, ICML '08, pages 1096-1103, New York, NY, USA, July 2008.
Association for Computing Machinery.
[30] R Devon Hjelm, Alex Fedorov, Samuel Lavoie-Marchildon, Karan Grewal, Phil Bachman, Adam Trischler, and Yoshua Bengio. Learning deep representations by mutual information estimation and maximization. arXiv [stat.MI.), August 2018.
[31] Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding, arXiv [cs. July 2018.
[32] David McAllester and Karl Stratos. Formal limitations on the measurement of mutual informa- tion. arXiv [C), November 2018
[33] Kyungeun Lee and Wonjong Rhee, benchmark suite for evaluating neural mutual information estimators on unstructured datasets. arXiv [stat.ML), October 2024.
12

---

[34] Shuyang Gao, Greg Ver Steeg, and Aram Galstyan. Efficient estimation of mutual information for strongly dependent variables. arXiv [cs.IT), November 2014.
[35] Alexander Rives, Joshua Meier, Tom Sereu, Siddharth Goyal, Zeming Lin, Jason Liu, Demi Guo, Myle Ott, Lawrence Zitnick, Jerry Ma, and Rob Fergus. Biological structure and function emerge from scaling unsupervised learning to 250 million protein sequences. Proc. Natl, Acad. Sci. U. S. A., 118(15), April 2021.
[36] Ahmed Elnaggar, Michael Heinzinger, Christian Dallago, Ghalia Rihawi, Yu Wang, Llion Jones, Tom Gibbs. Tamas Feher, Christoph Angerer, Martin Steinegger, Debsindhu Bhowmik, and Burkhard Rost. ProtTrans: Towards cracking the language of life's i code through self-supervised deep learning and high performance computing. arXiv [cs.LGG] July 2020.
[37] Ethan C Alley, Grigory Khimulya, Surojit Biswas. Mohammed AlQuraishi, and George M Church. Unified rational protein engincering with sequence-based deep representation learning. Nat. Methods, 16(12):1315-1222 December 2019.
[38] Zhidian Zhang. Hannah K Wayment-Steele, Garyk Brixi, Haobo Wang, Matteo Dal Peraro, Dorothee Kern, and Sergey Ovehinnikov. Protein language models learn evolutionary statistics of interacting sequence motifs. bioRriv, page 2024.01.30.577999 January 2024.
[39] Francesca-Zhouufnn Li, Ava P Amini, Yisong Yue, Kevin K Yang, and Alex X Lu. Feature reuse and scaling: Understanding transfer learning with protein language models. bioRxiv, page 1024.02.00.578959 February 2024.
[40] Dénes Türei, Alberto Valdeolivas, Lejla Gul, Nicolàs Palacio-Escat, Michal Klein, Olga Ivanova, Márton Ölbei, Attila Gábor, Fabian Theis, Dezsō Módos, Tamás Koresmáros, and Julio Saez- Rodriguez. Integrated intra- and intercellular signaling knowledge for multicellular omics analysis. Mol. Syst. Biol., 1733):9923, March 2021.
[41] Shou-Wen Wang, Michael J Herriges, Kilian Hurley, Darrell N Kotton, and Allon M Klein. CoSpar identifies early cell fate biases from single-cell transcriptomic and lineage information. Nat. Biotechnol., February 2022.
[42] Caleb Weinreb, Alejo Caariguez-Fraaicelle Fernando D Camargo, and Allon M Klein. Lineage tracing transcriptional landscapes links state to fate during differentiation. Science, 367(6479), February 2020.
[43] Kunal Jindal, Mohd Tayyab Adil, Naoto Yamaguchi, Xue Yang. Helen C Wang, Kenji Kami- moto, Guillermo C Rivera- Preera-Gnzzzlez, and Samantha A Morris. Single-cell lineage capture across genomic modalities with CellTag-mull reveals fate-specific gene regulatory changes. Nat. Biotechnol., September 2023.
[44] Caleb Weinreb, Samuel Wolock, and Allon M Klein. SPRING: kinetic interface for visualizing high dimensional single-cell expression data. Bioinformatics, (4(7):1246-122) April 2018.
[45] Duncan M Chadly, Kirsten L Frieda, Chen Gui, Leslie Klock, Martin Tran, Margaret Y Sui, Yo- dai Takei, Remco Bouckaert, Carlos Lois, Long Cai, and Michael B Elowitz. Reconstructing cell histories in space with image-readable base editor recording. bioRxiv, page 2024.01.03.573434 January 2024.
[46] Daniel E Wagner and Allon M Klein. Lineage tracing meets single-cell omics: opportunities and challenges. Nat. Rev. Genet., 21(7):410-42) July 2020.
[47] Zeming Lin, Halil Akin, Roshan Rao, Brian Hie, Zhongkai Zhu, Wenting Lu, Nikita Smetanin, Robert Verkuil, Ori Kabeli, Yaniv Shmueli, Allan Dos Santos Costa, Maryam Fazel-Zarandi, Tom Sercu, Salvatore Candido, and Alexander Rives. Evolutionary-scale prediction of atomic- level protein structure with a language model. Science, 379 March 2023.
[48] Jaewoong Cho, Gyeongjo Hwang, and Changho Suh. A fair classifier using mutual information. In 2020 IEEE International Symposium on Information Theory (ISIT), pages 2521-2526. IEEE, June 2020.
13

---

[49] Jian Kang, Tiankai Xie, Xintao Wu, Ross Maciejewski. and Hanghang Tong. InfoFair:
Information- -hthorettttt intersectional fairness. In 2022 IEEE International Conference on Big Data (Big Data), pages 1455-1464. IEEE, December 2022.
[50] Alexandre Lacoste, Alexandra Luccioni, Victor Schmidt, and Thomas Dandres. Quantifying the carbon emissions machine learning. arXiv [cs.CY], October 2019.
[51] Xavier Glorot and Yoshua Bengio. Understanding the difficulty training deep feedforward neural networks. In Yee Whye Teh and Mike Titterington, editors, Proceedings of the Thirteenth International Conference on Artificial Intelligence and Statistics, volume 9 of Proceedings of Machine Learning Research, pages 249-256, Chia Laguna Resort, Sardinia, Italy, 2010. PMLR.
[52] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Yang. Zach De Vito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang. Junjie Bai, and Soumith Chintala. PyTorch: An imperative style, high-performance deep learning library. arXiv [cs.LG), December 2019.
[53] Lukas Heumos, Anna Schaar, Christopher Lance, Anastasia Litinetskaya, Felix Drost, Luke Zappia, Malte D Lücken, Daniel C Strobl, Juan Henao, Fabiola Curion, Single-cell Best Practices Consortium, Herbert Schiller, and Fabian Theis. Best practices for single-cell analysis across modalities. Nat. Rev. Genet., pages 1-23, March 2023.
14

---

A Appendix supplemental material
this appendix, we first elaborate on the theory and implementation details of the LMI approximation. Then, we will discuss some alternative approaches one could use to implement the LMI approximation. Then, in Appendix A.3., we describe the details, assumptions, and motivations of our empirical evaluation benchmarks. In Appendix A.4., provide details relevant to reproducibility specifically, our implementations of existing estimators and data preprocessing methods (all of which are also included in our code supplement). Finally, discuss the problem of choosing appropriate latent space size for LMI models.
A.1
Theory and implementation of LMI
A.1.1 Theoretically motivating the cross-predictive representation learning architecture
The core theoretical underpinning of the network architecture used in the LMI approximation is that cross-predictive mean-squared loss is a proxy of conditional entropy. Here, explicitly show this.
Theorem 1. Let X = i......, Xa) and Z = =Zz..... Zul be absolutely continuous random vectors in IR and & respectively with finite differential eniropy. Let fa RR& - IR be function (a neural network parameterized by 0) to estimate = fa(2). For any 0.
2(X S a + log MSE(X.EE... X, X)
(6)
2
where a is positive constant and MSE X, X)= = E. B( - X1
Proof. From the chain rule for differential entropy, bound
h(XIZ)
(7)
Because the maximum entropy distribution with fixed variance is Gaussian we can bound
MX11Z) Σ log(2ree Var(X.[Z))
(8)
Σ log((2r (ire - B(X,1Z))
(9)
Because the expectation of random variable is its best estimator,
Σ 1 E -((x/)))))
Σ
log(2ne (2re E[(x, x,)
(10)
So with positive constant a,
1 MX(IZ) S aa 2 log MSE(X,
(11)
In LMI, Z corresponds to the latent representation of one input variable (Y), and corresponds to
the decoder which aims to reconstruct other variable X from the latent code Z. This result is very
similar to some used in information-theoretie interpretations of autoencoders 29 30 and can be
thought of continuous analog of Fano's inequality 77. The Corollary to Theorem 8.6.6 in 77)
case of this bound given
15

---

A.1.2 Implementation of the representation learning architecture
There many ways could implement the high-level network architecture suggested by Theorem 1, with different encoder and decoder architectures, and choices of hyperparameters. Here, we will describe our design choices. Our motivating philosophy was that the implementation details should not need to be tuned for specific estimation problems.
All LMI estimates presented in the main text use programatically determined parameters, requiring no user input beyond joint samples. By default, all latent representations have 16 dimensions (8 for each variable). Each encoder and decoder is multilayer perceptron (MLP) with three hidden layers, whose sizes are determined as a function of the dimensionality of the input variable. For a variable with dimensionality d, the encoder has hidden layer sizes L, L/2, L/4 with max(2 [ona(d)), 1024). Decoders have the same structure inverted, with hidden layer sizes L/4, L/2, L.
All MLP activations used are Leaky ReLUs, with negative slope 0.2, except the last layers of decoders, which have no activation. Cross-decoders are trained with 50% dropout after each activation layer. All weights are initialized using Xavier uniform initialization [51]. and optimized using Adam, with hyperparameters listed in Table 1. They are trained with batch size of 512, with 1 : train-validation splits, and a maximum of 300 epochs using early stopping procedure provided in Algorithm 3.
All models are implemented in Pytorch 52 NVIDIA RTX 3090.
All experiments in this paper were done using a single
Parameter
Value 10
Learning rate (a)
31
0.9
0.999
Epsilon (e)
10 10-7
Table 1: Adam optimizer parameters used in LMI.
Algorithm 3 Early Stopping
procedure EARLYSTOPPING(Mood) validation__ossee patience 30)
best_loss f 20
setiencce_cuntee <- 0
best_modee + None while patience cociseer_ountrrreeee patience do
current pourent_loss_t + validation Losp_oossmmmnn (oosee))
if current_ loss < best_ loss then
best_loss +- current_loss
patience_counten
best_model <- model
else patience_counte + patience_counter +
end if end while
return best loss, best mest_ end procedure
A.1.3 Theoretical properties of LMI approximation
The error of MI estimates using LMI approximation can be broadly attributed to two sources: error due to the representation approximation (i.e. [T(X;)) Y) - -T(Zz: Zy))),, and classical MI estimation error (i.e. | (Z;; Z,) - [[Z,; Z.)D. Here, we will explore the first source by deriving some basic properties of the representation approximation I(Zz; Zy).
Theorem Let X,Y Y be random vectors Rdd and Z5 fo(X), Zyy (V) where f.g are neural networks parameterized by 0, c. For any 0, 6.
16

---

[(Zi ((X:)
(12)
Proof. Because 2y = =Y) we have that X + Y - Zy form Markov chain. From the data
processing inequality,
I(XY)) 1xx
(13)
Similarly, we have Zy - X - Zz, and
((x;))2((3) 2 1(Zxi Zny
(14)
Theorem 3. Let X, Y be independent random vectors in R. such that YY; 0. Let Zu
fX(X), Zy = (Y) where f, g are neural networks parameterized by 0, S. For any 0, Φ.
Zg) =0
(15)
Proof. Due to the nonnegativity of mutual information, we know
I(Zzi Zg) 20
(16)
From Theorem 2, we have
I(Zz: 5 SIX;Y) =0
(17)
So.
0 0 I(zi: (,) 0 0
(18)
T(Zz; Z,y) =0
(19)
17

---

1W
30-
180
200
100
590
500
100
2000
3000
600
156
00
30
300
390
300
109
3039
3030
1003
INcSCC
VINE:
06
30
200
001
530
500
300
2230
500
500
5030
2009
(a) Absolute emror, measured by MSE.
(b) Relative error, measured by Kendall T.
Figure 8: Experiment from Figure 2, with alternate regularization approaches.
A.2 Alternate approaches to latent MI approximation
The broad goal of LMI, estimate I(X;Y) using Zz) where Za, Zy low-dimensional representations, is quite general and could be approached in many ways beyond cross-predictive regu- larization. Next, will discuss some alternative approaches, empirically explore their performance, and finally, show one unique ad advantage of the cross- possspreedittve representation learning architectures.
A.2.1 Alternate methods regularizing autoencoders for MI estimation
Another approach to learn Zu, Zy such that i(Z,; Z,,) I(X;Y) to regularize autoencoders to maximize I (Zn; Zy). This approach is sensible because the data processing inequality ensures that I(Zz; Zy) < ((XY)) So maximizing [Zzz; Zyy is equivalent to minimizing the approximation
emor T(X; Y) I(2z; Zy).
While maximizing directly I(Zz: Zg) intractable, we can build on the variational bounds explored in [15] BI.. We can add loss term LMIEE Lz; Zy) or LInfoNCE( Zz; Z,) to our autoencoder loss functions regularize latent codes preserve mutually informative structure.
We implement both of these approaches, and benchmark them using the approach described in Figure 2 and Section 3.1. We find that the estimates from these regularization approaches perform similarly, but slightly more poorly, than the cross-predictive regularization.
A.2.2 Comparing latent nonparametric and latent variational MI estimation
After learning low-din www-dimensionll representations, there are multiple estimators that could be used for latent MI approximation. This appendix evaluates two methods: the KSG nonparametric estimator or the InfoNCE variational estimator. Nonparametric nearest-neighhor estimators have several advantages low-dimensional settings: they generally require far fewer samples for accurate estimation [12, and yield pointwise mutual information decompositions (see Algorithm 5).
18

---

0.8
Estimate (bit3)
02
Latent InfoNiE
Latent KSG
50
100
150
200
250
300
Epoch
Figure 9: Convergence of multiple latent estimation approaches during training. Bold lines indicate averages over 100 trials. Ground truth is I bit MI.
MSE
0.4
0.2
Figure 10: Pixel-wise reconstruction error of cross-decoders in paired binary MNIST dataset where
L
Ly.
For completeness, we empirically compare both options here on one particular estimate. For Gaussian dataset generated by Algorithm 4 with d = 100, k 1, N = 5 10 and bit MI, we train autoencoders regularized by InfoNCE loss to maximize T(Z,; Z,y). After each training epoch, we measure LMI estimated using latent KSG, and using latent InfoNCE. Both are plotted for 100 trials Fig. 9. The latent KSG estimation converges quickly to true value, while the latent InfoNCE estimate converges somewhat slowly to value below the ground truth.
A.2.3 Interpreting decoders with element-wise reconstruction error
Beyond performance differences, benefit of using cross-predictive networks regularize latent representations is that the cross-decoders themselves are useful. For example, by inspecting the dimension-wise reconstruction error of decoders, we can attribute MI estimate to the predictability of certain dimensions. For a network trained on the "binary" MNIST dataset with Le Ly, visualizing the dimension-we reconstruction error of cross-predictive decoder reveals the pixels that contain information about digit identity (Fig. 10).
Pixels with low reconstruction error are likely to be "well-explained" by the other high- dimennioon variable, while pixels with high reconstruction error are poorly explained. However, this reasoning is universally applicable. Dimensions with variation do not contribute to MI, but have very low reconstruction error. In Fig. 10, the outermost pixels are examples of this.
A.2.4 Effect of regularization on failure for symmetric and exclusive variables
Because the failure to preserve MI specifically arises from distributions with symmetry under MSE loss, one might conclude avoiding cross-prediction loss may be sufficient to improve LMI estimation. Indeed, we demonstrate this for two models, LMI-MINE and LMI-InfoNCE, described in Appendix 1. which are regularized by maximizing a variational lower bound on T(Z,; Zyy) rather than cross-prediction. In our Gaussian benchmark, we found these regularizations to be less effective than
19

---

circle
Gaussian mixture
0.8
Estimetorr
0.6
Etimate
Estimate
LMI-MINE
0.4
LMI-InfONCE
0.2
InfoNCE
0.0
-25
0.0
2.5
0.0
2.5
5.0
7.5
Exclusive dimensions
Exclusive dimensions
Figure 11: a) Examples of symmetric distributions where [[ = [[X]. b) MI estimates from 10 samples of symmetric variables exclusive dimensions added (increasing non-mutual information). Ideal estimators are invariant. All LMI latent spaces are 1 dimension per variable. LMI-MINE and LMI-InfONCE denote alternately regularized LMI models.
prossppredition: performance decayed more quickly with increasing intrinsic dimensions (Fig. 11). Here, however, these alternate models perform better than LMI, although ultimately they still fail with exclusive dimensions. Many practical data sets may not suffer from the symmetries considered here, but these failure modes are nonetheless instructive clarifying how choices made in learning latent representations may drive loss of MI. In practice, if one suspects that their data may have structure not captured by MSE cross-prediction loss (e-g. symmetries), the benchmarking results suggest that it may be advisable to use one of the alternate regularization methods. In our software library, this is simple to do, e.g. 1mi (x, Y, regularizer= mmodels........oooeooooo. AEMINE").
20

---

A.3 Details of experimental evaluation benchmarks
In this section, we will describe the details of our experimental evaluation benchmarks from Section 3. First, we will describe how we generate multivariate Gaussian datasets. Then, we will provide theoretical and empirical validation for the cluster-based benchmarking approach.
A.3.1 Generating multivariate Gaussian datasets with low-dimensional dependence structure
The algorithm used to sample multivariate Gaussians in Figure 2 is given in Algorithm 4. Briefly, generate samples of l correlated dimensions by sampling bivariate Gaussians. We then complete the remaining d' & ambient dimensions using half "redundant" dimensions (copies of the k correlated dimensions) and half "nuisance" dimensions (independent univariate Gaussians).
Algorithm 4 Generating multivariate Gaussian datasets with low-dimensional dependence structure
Require: ambient dimensionality d
Require: dependence structure dimensionality kt
Require: number of nuisance dimensions n (default (d - k))
Require:
number of samples N
Require: ground truth MI
W Compute cov matrix yield dependence b
let + v6 3.5 21-2-22/77 2-2b/ 25/E
let Σ 4 [[6, pl. [p, 3/5]]
Sample bivariate Gaussians for & dependent dimensions
for i in 1..k do
let X,, N samples from (0, Σ) end for Duplicate random dimensions for redundant dimensions for i in I..(d - (k + 1)) do
let 5- kni let X14k. Y +- X,,Y, Y
end for \I Sample univariate Gaussians for nuisance dimensions
for i On.(n - 1) do let Xa-s f samples from (0, 1) let Ya-i + N samples from (N end for return (al. [YI...YYa
A.3.2 Theoretical justification for label MI approximation of nigh-dimensional MI
For the benchmarking setup described in Section 3.2, we have Ly - Lg + and Lg + Ly + Y. We will show that Y(X;Y) I(Lx; Ly) under the condition that H (L, HL,, Y== =
Theorem 4. Let Tegy Tyy be Bernoulli random variables. Let X, Y be absolutely continuous random vectors in T&N such that I(Lg; ([,,) T(L.; X I.,) = 0 and HI, (X) HL,,IY) 0. Then,
I(X;Y) - [Las Ly)
(20)
Proof. Due conditional independence,
I(L,; = + (La;i\\L)
(21)
Using the chain rule for mutual information 7
I(Lx; L.y ((Lii Ly) + [(Lzi (Lz; I(L;;)) I(bnj Lyll)
(22)
Applying the chain rule again,
[LL: Ly) I(Lo; Y) + + L, -(x)) - ((L;i)
(23)
21

---

Digit
0
1
UMAP2
UMAP2
Species
A. thaliana
E. coli
UMAP1
UMAP1
(a) UMAP of binary MNIST subset.
(b) UMAP of proteome subsets.
Figure 12: Validating assumptions of cluster-based benchmarking, by visualizing separation of clusters.
Due to the non-negativity of mutual information we have
T(Izi Lg) 2 I(X:Y)
(24)
Because Tur, Tyy are discrete, we can bound
Ly) I(X;Y)+++++) H(L (X) H(L,lY)
(25)
So H( ((L,)Y) = O, we have
S I(X:)
(26)
Combining this with (24), have
I(XYY) I(bni Ly)
(27)
A.3.3 Validating the assumptions of cluster-based benchmarking setups
The effectiveness of the benchmarking setup in Section 3.2 relies the assumption that H(2, (X) H[L,)Y) R 0. We provide evidence that this is the case for MNIST Os and 1s, and sequence embeddings of E. Coli and A. Thaliana proteins. First, provide qualitative evidence by showing that label clusters (digits and species respectively) are well separated on UMAP visualizations of each dataset (Fig. 12). This indicates that the label of sample can be reliably determined from its nigh-dimensional representation (such that H (L,,) (X) 2o O). We make this notion more quantitatively precise by showing that logistic regression can predict the labels of held-out samples with high accuracy. Over 100 random 1 1 train-test splits, logistic regression achieves mean validation accuracy 0.98 with standard error of mean < 10 10-3 for both datasets. Logistic regression classifiers are trained with La penalty and X = 1.
A.3.4 Synthetic multivariate Gaussian benchmarking with nonlinear transformations
known that benchmarking MI estimators on Gaussian synthetic data has many shorteomings 12 In addition to addressing this with the resampling- benchmarking approch in Section 3.2, we also perform a subset of the multivariate Gaussian benchmarks after data been nonlinearly transformed.
22

---

MINE
Tranaformation
Transformat
Asinh
Half cube
None
2.0
Uniform margins
Intrinsic dimensions
Intrinsic dimensions
1.5
Tra Transformation
KSG
isformation
InfoNCE
1.0
Asinh
0.5
Half cube
None
0.0
Uniform margins
4
5
6
7 8
1
2
3
8
9
Intrinsic dimensions
Intrinsic dimensions
Figure 13: Performance (measured as MSE) on subset of multivariate Gaussian benchmark after non- linear transformations (defined in (12). All estimation problems are with 1000 ambient dimensions.
Rather than using the exact tasks from 12]. we have developed complementary high-dimensional benchmarks by applying some of the transforms proposed in [12 our multivariate Gaussians from Fig. 2. This allows us to create versions of the existing Fig. 2 with "half-cube", "asinh", and "uniform marginal" distributions (as defined in (12) We explore only the most challenging settings from Fig. 2: those with 1000 ambient dimensions and 1-9 intrinsic dimensions. We find that LMI performance qualitatively similar, though identical, for untransformed data and all three transformations (Fig. 13).
A.3.5 Sampling procedure for symmetric and exclusive variables
To sample from symmetric and exclusive variables (section 3.3, Fig. 5), we concatenate independent normally distributed dimensions to symmetric variables. Below, we give the procedure for sampling from symmetric variables.
Noisy circle For each sample, angle 0 is drawn uniformly from the interval [0, 2m). For each angle, the corresponding x- and y-coordinates of each sample are x cos(f) + and y = sin(f) + where and cy are independently sampled from a uniform distribution in the range [-0.25, 0.25]. This results in points approximately lying on the cireumference of a unit circle, with random noise.
Gaussian mixture
We sample from a two component Gaussian mixture, with each component
1
0.9
1
-0.9 and both with mean
having equal weight and covariance matrices
and
0.9
1
0.9
0.
23

---

Dims. 10
Dims. 100
Dims. 1000
1.0
Estimate
Estimator
0.5
MINE
InfONCE
0.0
500
1000
1000
500
1000
Max layer size
Max layer size
Max layer Size
Figure 14: Variational bound estimators with increasing critic complexity, evaluated multivariate Gaussians with I-dimensional dependence structure. Each estimation problem has 10 samples and ground truth MI of bit.
A.4 Experimental reproducibility details
Here, will briefly summarize key details necessary for the reproduction of experiments in this paper. All of this information, and more details, can be found (albeit in less easily readable form) in our code supplement.
A.4.1 Preprocessing ProtTrans5 embeddings
We downloaded ProtTrans5 embeddings of all H. sapiens, A. thaliana, and E. coli proteins directly from the UniProt database. Embeddings are from prottrans. saattrrsss___ae _____________ [36]. All proteins longer than 12 10 residues excluded. These embeddings are then unit variance normalized, and values are clipped at 10 and -10.
A.4.2 Preprocessing hematopoiesis lineage tracing scRNA-seq data
We first downloaded all data from Experiment 1 of the data repository from Then, we prepro- cessed using the Scanpy best practices [53], normalizing total reads per cell to 10 log transforming, filtering for the 10 most highly variable genes, and finally unit variance normalizing and clipping values at 10 and -10. We used the inferred diffusion pseudotime and SPRING embeddings computed in 42 For pseudotime analysis, we omit cells with pseudotime value below 10 because many are Lymphoid-fater rather than Neutrophil-fand conditions, we identified all clonal barcodes that appeared in both conditions, and randomly sampled single cell with each barcode from each condition. Because there are often several cells with the same clonal barcode the same sample, there are many possible random pairings clonally related cells.
To generate joint samples of clones between two
A.4.3 MINE and InfoNCE implementation details
We use the implementations of MINE and InfoNCE from 112 with their default parameter choices. To summarize, architectures have two hidden layers with sizes (16, 8) and are optimized using Adam.
Choosing critic architectures Because the benchmark tasks of 12 are dramatically different from those considered in this work (tens of dimensions opposed to thousands), consider the possibility that the parameters of MINE and InfoNCE used in 2 are not suitable for our use case.
To determine if other critic architectures could be more suitable, we first tested two layer architectures with layer sizes (L, L/2) for various L from to 1024 on multivariate Gaussian data sets. We considered variables with dimensions 10, 100, and 1000, with 1 bit MI, 1-dimensional dependence structure, and 1 samples. We find that increasing L does not improve estimation quality, and the architecture used by 212] indeed is optimal for all three tested settings (Fig. 14). We suspect that this because increasing model complexity does not help when sample size remains small.
On the MNIST benchmarking setup, we verify that the effectiveness of LMI estimation over MINE and InfoNCE are not merely from model complexity, by using crities with the same complexity as the LMI encoders. In line with the hypothesis that large critic architectures are not suitable for the small sample size regime, we find that with high complexity crities, MINE and InfoNCE, fail entirely on the MNIST benchmark (Fig. 15).

---

1.0
Estimator
InfoNCE
0.8
MINE
Eslimale
0.6
.abel MI
Figure 15: MNIST benchmarking for neural estimators with critic complexity equivalent to LMI encoders, over 20 datasets with true MI between and 1.
A.4.4 KSG implementation details for pointwise decompositions
We implement KSG with A nearest neighbors. For protein interaction data, due to large sample numbers resulting in high computational cost, we obtain KSG estimates by averaging over estimates on batches of data containing 10 samples.
We slightly adjust the KSG estimator yield pointwise mutual information estimates. While the original KSG estimator [10] takes sample expectation over computed pointwise mutual information values, we simply return an array of pointwise estimates rather than the average. For completeness, the algorithm is given in Algorithm 5.
Algorithm 5 KSG estimator for pointwise estimates
Require: joint samples {(z, 2) Require: parameter ls (default k = 3) let pmis <
for each (I;, do
find k-th nearest neighbor in joint space (æts, (k) compute Chebyshev distance d = [1( yk) - (I, 1100
let nı 1-0,0,11 0, fly +-
for each (Ir Mj do then
if 1| Iz
n- +=
end if
if ||ı; Vi < d then
Ny +=
end if
end for end for return pmis
( (kk) + (NN) (((x)) - (bny))
25

---

A.5 Choosing appropriate latent space size
All estimates the main text of this paper is made with 8 latent dimensions per variable. Without optimizing this parameter, we find that LMI performs reasonably well compared to other estimators on diverse benchmarks and real-world problems. However, the choice of latent dimensions per variable is somewhat arbitrary, and certainly not optimal for most estimation problems. Here, we outline some principles for how one might choose more optimal number of latent dimensions.
A.5.1 General principles for choosing latent space size
Intuitively, there is tradeoff which arises when changing latent space size. As the latent space gets larger, the capacity of the compressed representation increases, and we might expect that representation quality increases (with the caveat that representation quality can be limited by sample sparsity). However, as the latent space gets larger, MI estimation in latent space becomes more difficult, due to the curse of dimensionality for nonparametric MI estimators (14). As such, the ideal choice smallest possible latent space size which captures the dependence structure variables. In practice, this size can be difficult to determine in rigorous way. Instead, suggest heuristic approach.
A.5.2 Heuristic approach to choosing latent space size
From Theorem 2, simple extension of the data processing inequality, know that I(Z,; S I(X;Y) With the caveat that the inequality is not guaranteed hold for the estimated i(z Zy). we can reason that the parameter choices that maximize f(Zz; Zyy) are likely ideal. As such, one sensible way to choose latent space size is to try several, and use that which yields the largest estimate.
As an example, let consider multivariate Gaussian data with 4-dimensional dependence structure, 1000D ambient space, with ground truth MI of bit, and 5 10 samples (generated analogously to Figure 2). The most accurate estimate comes from LMI with optimal choice of 4 latent dimensions, but all tested choices (including the 2D latent space which cannot fully capture dependence) improve over MINE and InfoNCE. If we the heuristic approach, would choose the optimal latent space size. If we had arbitrarily chosen 8, we would be within 5% of the "optimal" latent dimension estimate.
LMI-2
LMI-4
LMI-6
LMI-8
MINE
InfoNCE -0.001417
Estimate
0.295732
0.719762
0.670632
0.686974
-0.000001
Table 2: LMI estimates with varying latent space size, with & dimensions per variable denoted as LMI-k. Data is multivariate Gaussian, generated as in Figure 2, with d 103, k = 4, and N 5 103
26

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
The results stated in the abstract and introduction are summaries of the experimental results shown in Sections 3 and (figures 2-6).
Justification:
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as is clear that these goals attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: There is subsection of the Discussion dedicated to limitations of the LMI approach. These limitations are considered throughout the paper. For instance, in abstract, highlight the necessity of low dimensional intrinsic dependence structure.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "I "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full of assumptions and complete (and correct) proof?
27

---

Answer: [Yes]
Justification: Assumptions about distributions (such as absolute continuity, finite entropies) are specified in theoretical analysis (Theorems 1-4).
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: Key details necessary to reproduce the experimental results are given in the appendix. All experiments can be reproduced using the supplementary code. It is carefully annotated match each result in the paper to a specific Jupyter notebook.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
28

---

5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification: As stated before, all experiments be reproduced using the supplementary code. It is carefully annotated to match each result in the paper to a specific Jupyter notebook. The code reproducibility "best practices" were followed.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips. cc/
for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
mtiiss.cc/public//ee
for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: All training and test details are provided in the Appendix sections imple- mentation, and can be found in the supplementary code.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We report SEMs and standard deviations in the text when relevant and feasible. No plots contain error bars.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emoor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
29

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative eтror rates).
If error bars reported in tables plots. The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We report the hardware used for all experiments, and the projected overall environmental impact of all experiments (failed and pilot) reported in the paper. We do not provide further granularity because the experiments have rather modest compute requirements the entirety of the experiments this paper can be reproduced using commercial NVIDIA GPU (RTX 3090) in about one day (by running every Jupyter notebook in the code supplement).
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: bttttp::///euriiie /tt./enurip.../ www.riis.ccpppulule cc/publlic/thiiccsulllllll...
Answer: [Yes]
Justification: We have reviewed and adhered to the guidelines. We mention potential social and environmental impacts in the Discussion section.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
30

---

Justification:
We discuss the potential social impact of MI estimators.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negatíve societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of MI.).
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (c.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification:
We do not release data or models with high risk for misuse.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification:
We cite the library used to implement the LMT approximation (PyTorch).
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
31

---

The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperssi thoode com comddatasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
The code developed for latent MI approximation (the 1mi library provided in the supplement) well documented.
Justification:
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
This work does not involve human subjects or crowdsourcing.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This work does not involve human subjects crowdsourcing.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
32

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
33


Reviews:
[{"review_id": "yMA1fVsWiR", "reviewer": "Reviewer_NtLr", "summary": "This paper introduces an approach to approximate mutual information (MI), by applying a nonparametric MI estimator to the learned representations. The representations are learned by minimizing a weighted sum of the reconstruction loss and the prediction loss. The authors conducted experiments on both synthetic data and biological data.\n\nNo higher ratings: The representation learning part has some technical flaws.\n\nNo lower ratings: The evaluation methods and discussions provide new understandings and suggest some new applications.", "questions": "None.", "limitations": "The authors have adequately acknowledged the limitations of the proposed method.", "rating": 5, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The evaluations are comprehensive with detailed discussions. The demonstrations suggest broader applications of MI in scientific domains.\n2. The paper is well organized; presentations are clear and easy to follow.", "weaknesses": "Despite the comprehensive evaluations and discussions, the proposed approach has a main technical issue: the learned representations might not capture the dependence between $X$ and $Y$, making the LMI fail.\n\n1. Ideally, the reconstruction loss (autoencoder) alone can preserve the information of $X$, $Y$. However, it can be very inefficient when there are many redundancies in $X, Y$. A counter-example is when $X = (U, W), Y = (V, W)$, and $U$, $V$ have much more information than $W$;\n2. The prediction loss highly depends on the $X$, $Y$; in the worst case, the prediction loss does not reveal any information about the dependence. A simple example is when the densities of $X$, $Y$ are symmetric. Two particular cases: (1) when $X$ and $Y$ are uniformly distributed on the unit circle; (2) when $X$ and $Y$ are generated from the mixture Gaussian: $ (\\mathcal{N}(0, \\begin{bmatrix}1, 1/2; 1/2, 1 \\end{bmatrix} + \\mathcal{N}(0, \\begin{bmatrix}1, -1/2; -1/2, 1 \\end{bmatrix})/2$. In both cases, the best predictor is simply zero due to the symmetry.\n3. By combining the above examples, one can construct examples where the proposed loss does not provide informative information.\n4. However, the above counterexamples, especially for prediction loss, can never appear for Gaussian data. This makes the evaluation of Gaussian synthetic data less convincing.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for this thoughtful review. You raise one major point of concern:\n\n“***\\[...\\] the learned representations might not capture the dependence between $X$ and $Y$, making the LMI fail.***”\n\nWe agree that there are situations where LMI will fail (as is broadly true for all MI estimators \\[2\\]), and aim to clearly identify them in our paper. You are correct that we had failed to consider the setting you describe, and we have now studied it more carefully. Our work has been strengthened as a result. The key new findings are:\n\n1. As the reviewer correctly identified, certain distributions are problematic for LMI. We refer to the general class raised by the reviewer as distributions with “symmetry and exclusivity”.  \n2. However, all other methods we evaluated also fail in practice, and thus this is a problem more general than for LMI.  \n3. For LMI, the failure is specific to a choice of regularization, and our paper proposed more than one regularization method (Appendix A.2.1). We now show that the alternate regularization methods recover some functionality for symmetric and exclusive variables.\n\nIn our revised paper, we will include discussion of symmetry and exclusivity in the Limitations section, and our new empirical results in the Appendix.\n\nWe now address the specific points raised in more detail:\n\n***“1. \\[...\\] A counter-example is when $X=(U, W),  Y=(U,V)$ and $U,V$ have much more information than $W$;”***\n\nThe reviewer is correct that when variable exclusive information outweighs shared information, self-reconstruction loss is insufficient – however the regularization term can still induce $Z\\_x, Z\\_y$ that share mutual information. Thus, this scenario, representative of “exclusivity”, is alone not problematic for LMI, as seen in our Figure 2 benchmark.\n\n***“2. \\[...\\] in the worst case, the prediction loss does not reveal any information. A simple example is when $X, Y$ are symmetric. \\[...\\]”***\n\nWe deeply appreciate this key insight. The two examples (“O” and “X”) are instances of distributions for which $\\\\mathbb{E}\\[X|Y\\] \\= \\\\mathbb{E}\\[X\\]$, meaning the MSE minimizing predictor of $X$ is independent of $Y$. As such, the cross-prediction loss will not have a meaningful impact on $Z\\_y$. This “symmetry” alone is not problematic for LMI, because the self-reconstruction can still induce meaningful $Z\\_x, Z\\_y$, as seen in Rebuttal Figure (RF) 2, introduced later in the response.\n\n***“3. By combining the above examples, one can construct examples where the proposed loss does not provide informative information.”***\n\nWe agree. While LMI can handle each above example in isolation, a combination can cause LMI to fail. Symmetry is problematic for cross-prediction, and exclusivity is problematic for self-reconstruction. Given this limitation, we now explore (1) whether other methods suffer from the same limitation, and (2) whether it can be overcome by changing the cross-prediction loss.\n\nTo explore these questions, we develop a benchmark to empirically study MI estimation in cases of exclusive symmetric distributions. Our benchmark involves estimating MI from samples of symmetric distributions with increasing levels of variable exclusive information. We choose two symmetric distributions, shown in RF1: (1) the Gaussian mixture proposed by the reviewer and (2) a noisy circle, modified from the reviewer’s suggestion to be well-behaved. To add exclusive information, we inflate each variable with independently normally distributed dimensions. The ideal estimator should not vary with the number of exclusive dimensions, as true MI is invariant. Following the reviewer’s argument, with no exclusive information, LMI should yield a good estimate (due to self-reconstruction loss), but the estimate should quickly decay to 0 as exclusive information is increased.\n\nIn the benchmark, we first confirm that LMI behaves as predicted by the reviewer. The decay in LMI occurs quickly, already with \\<10 exclusive dimensions (RF2).  \n\n**Result 1: other MI estimation methods suffer from the same limitation**\n\nWe found that the problem is not unique to LMI – all studied estimators fail in the symmetry and exclusivity setting, with \\<10 exclusive dimensions (RF2). Note that in this dimensionality, MINE and InfoNCE typically perform well \\[1\\], so their failure is not due merely to the dimensionality of the problem but also due to the nature of the distributions. The estimators do not agree even in the 1D case without exclusivity, suggesting that symmetric distributions may be generally problematic for MI estimation, similar to long-tailed distributions \\[1\\].\n\n**Result 2: handling symmetry by avoiding L2 cross-prediction loss**\n\nFor LMI, failure under symmetry is an artifact of the L2 cross-prediction loss. An alternate regularization approach may be able to handle distributions with symmetry and exclusivity. We proposed two such models in Appendix A.2.1, which maximize a lower bound on $I(Z\\_x, Z\\_y)$ via MINE and InfoNCE in latent space. In our Gaussian benchmark, we found these regularizations to be less effective than cross-prediction: performance decayed more quickly with increasing intrinsic dimensions (Figure 7). Now, in the symmetry and exclusivity benchmark, we find that these alternate models, LMI-MINE and LMI-InfoNCE, perform better than LMI, but still fail with 8 exclusive dimensions. In our eyes, moderate improvement in the case of symmetry and exclusivity is outweighed by the worse scaling with intrinsic dimension, so we prefer the cross-prediction loss for general use. If a practitioner suspects that their data may be symmetric and exclusive, it may be advisable to use one of the alternate regularization methods. In our software library, this is simple to do, e.g. `lmi(X, Y, regularizer=\"models.AEMINE\")`. \n\n***“4. However, the above counterexamples \\[...\\] never appear for Gaussian data.”***\n\nThis point is now fully addressed by the additional benchmarks already discussed above.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "CSJnTmMVs2", "reviewer": "Reviewer_13JG", "summary": "This paper proposes latent MI (LMI), a method for estimating the mutual information (MI) between two high-dimensional multivariate random variables. For that, the technique uses the non-parametric MI estimator from [KSG04] on lower-dimensional latent representations that are learned by neural networks such that their MI is close to the one between the original variables. The paper provides some theoretical motivation for the proposed method. Finally, there is an experimental evaluation and comparison with other state-of-the-art methods for MI estimation, together with applications to problems in biology.\n\n\n[KSG04] Kraskov, A., Stögbauer, H., Grassberger, P. (2004). Estimating mutual information. Physical Review E—Statistical, Nonlinear, and Soft Matter Physics, 69(6), 066138.", "questions": "1) While the paper only apply the LMI estimator to problems in biology, have the authors considered other applications? For instance, the MI plays a central role in supervised learning in which the feature vector $X$ is high-dimensional (e.g. a MNIST image) while the label $Y$ consists in a few dimensions. Typically $X$ would be compressed but not $Y$, resulting in quite different latent spaces (different dimension and nature) unlike in the examples introduced in the paper.\n\n2) Could the authors explain the choice to apply the KSG estimator to the latent representations in their LMI method? Have the authors considered combining their method to more modern methods such as MINE [BBR18]? \n\n3) The last paragraph of Section 3.2 refers to \"Fig. 3a, 3b\" and \"Fig 3c\" while it should be to Fig. 4. Please make sure the figures are referenced properly throughout the paper.\n\n4) Related to the previous point, please be consistent when referring to figures. It should be \"Fig. \\#\" or \"Figure \\#\" (preferably the former), but both are used.\n\n5) Some relevant references (in my opinion) on MI estimation could be added in the introduction e.g., [NZH19], [GVG15], [MAK20]. \n\n[BBR18] Belghazi, M. I., Baratin, A., Rajeswar, S., Ozair, S., Bengio, Y., Courville, A., & Hjelm, R. D. (2018). Mine: mutual information neural estimation. arXiv preprint arXiv:1801.04062.\n\n[NZH19] Noshad, M., Zeng, Y., Hero, A. O. (2019, May). Scalable mutual information estimation using dependence graphs. In ICASSP 2019-2019 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP) (pp. 2962-2966). IEEE.\n\n[GVG15] Gao, S., Ver Steeg, G., Galstyan, A. (2015, February). Efficient estimation of mutual information for strongly dependent variables. In Artificial intelligence and statistics (pp. 277-286). PMLR.\n\n[MAK20] Mukherjee, S., Asnani, H., Kannan, S. (2020, August). CCMI: Classifier based conditional mutual information estimation. In Uncertainty in artificial intelligence (pp. 1083-1093). PMLR.", "limitations": "The authors addressed most of the limitations of their work in the \"Limitations\" section of their paper. Please however see the 'Weaknesses' section.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1) Pragmatic presentation of the problem setup and the proposed solution.  \n\n2) Interesting approach to MI estimation in (very) high dimensions, an open research subject, exploiting the informative low-dimensional structure of variables, which is a trendy approach (representation learning). \n\n3) Some numerical illustrations of the proposed method, focused on the interpretability of the estimator, which is of paramount utility when it comes to applying mutual information to real-world problems. \n\n4) Most limitations of this work are acknowledged by the authors.", "weaknesses": "1) Theoretical justification for the method is rather simplistic, as acknowledged by the authors, and relies on potentially loose approximations such as the data processing inequality.\n\n2) Overall, the proposed method consists in applying an existing estimator to pre-processed input variables, in the form of latent representations. It seems to lack a joint design, which results in two additive and independent sources of error (one from the representation, the other from the estimation itself). \n\n4) Some minor concerns are raised in the 'Questions' field.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for this insightful review. \n\n***“Overall, the proposed method consists in applying an existing estimator to pre-processed input variables, in the form of latent representations. It seems to lack a joint design, which results in two additive and independent sources of error (one from the representation, the other from the estimation itself).”***\n\nThe reviewer is correct in noting that LMI can have two sources of error: representation error, and estimation error. This is a conscious choice, made because KSG has many useful properties which variational bound neural estimators lack. So we will address it in conjunction with the later question:\n\n***“Could the authors explain the choice to apply the KSG estimator to the latent representations in their LMI method? Have the authors considered combining their method to more modern methods such as MINE?”***\n\nWe considered the possibility of using MINE or InfoNCE as latent estimators, and showed one empirical result in Appendix A.2.2 where latent KSG outperforms latent InfoNCE. However, our choice to use KSG is deeper than this one result. The KSG estimator has a number of useful properties, most notably:\n\n1. KSG easily decomposes into pointwise mutual information estimates. Decomposing MI estimates into their pMI values allows us to “explain” which samples contribute to an MI value. In real-world applications, this is extremely useful. We show this concretely in Sections 4.1 and 4.2, where pMI decomposition enables prediction of protein interactions (4.1) and identifying transition points during neutrophil differentiation (4.2).  \n2. KSG is far more sample efficient in low-dimensional settings than MINE and InfoNCE. This is well-established: see Figure 9 of \\[1\\].  \n3. In general, MI estimators struggle in high MI settings \\[2\\]. This phenomenon is particularly well studied for KSG, and corrections have been developed for situations with high MI \\[3\\]. We are optimistic that these corrections can be applied to LMI as well.\n\n***“While the paper only apply the LMI estimator to problems in biology, have the authors considered other applications? For instance, the MI plays a central role in supervised learning in which the feature vector is high-dimensional (e.g. a MNIST image) while the label consists in a few dimensions. Typically $X$ would be compressed but not $Y$, resulting in quite different latent spaces (different dimension and nature) unlike in the examples introduced in the paper.”***\n\nThank you for this interesting question. In principle, LMI easily adapts to this situation. Reusing notation from Section 2, if $X$ is our compressed variable, we must simply train networks $E\\_X, D\\_{XY}, D\\_{XX}$ with the loss $\\\\mathcal{L} \\= \\\\text{MSE}(X, D\\_{XX}(E\\_X(X)) \\+ \\\\text{MSE}(Y, D\\_{XY}(E\\_X(X))$. Then we can estimate $\\\\hat{I}\\_{KSG} (Y; E\\_X(X))$.\n\nThough our existing software implementation does not generalize to this kind of model, we can still input a one-dimensional variable which gets “encoded” and “decoded” into latent space, such that the network must learn the identity function. While this approach includes some unnecessary parameters, we empirically evaluate its effectiveness in the problem of measuring the mutual information between MNIST digits and their labels. That is, $I(X; L)$, where $X$ is the 784-dimensional distribution over images, and $L$ is the 1-dimensional distribution over digit identities. The ground truth here is roughly $I(X; L) \\\\approx \\\\log\\_2 (10) \\\\approx 3.3$ bits.\n\nDespite the facts that (1) we are processing images without inductive bias from convolutional layers, (2) we are “wasting” parameters by using our default architecture choices, and (3) we are not carefully treating the discrete-continuous mixture, we find that LMI works reasonably well, significantly outperforming MINE and InfoNCE. We show the results in the table below. We will include this analysis in the Appendix of our revised paper.\n\n|  | LMI | MINE | InfoNCE | Ground truth |\n| ----- | ----- | ----- | ----- | ----- |\n| I(X;L) | 2.37 | 0.02 | 0.47 | \\~3.3 |\n\n***“The last paragraph of Section 3.2 refers to \"Fig. 3a, 3b\" and \"Fig 3c\" while it should be to Fig. 4\\. Please make sure the figures are referenced properly throughout the paper. Related to the previous point, please be consistent when referring to figures. It should be \"Fig. \\#\" or \"Figure \\#\" (preferably the former), but both are used.”***\n\nWe apologize for figure reference errors and inconsistency. Our revised paper will correct these.\n\n***“Some relevant references (in my opinion) on MI estimation could be added in the introduction e.g., \\[NZH19\\], \\[GVG15\\], \\[MAK20\\].”***\n\nThank you for pointing out these relevant references which will strengthen our revised introduction.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ByMOHlhaPN", "reviewer": "Reviewer_8utg", "summary": "The focus of the paper is on approximating mutual information (MI) between multidimensional variables. This problem is challenging as the approximation of the MI suffers from the curse of dimensionality. The authors propose a method that approximates the MI via an embedding in a lower dimensional space. They test their methods on embedding from proteins language model, scRNA-seq data, and toy datasets.", "questions": "- In Fig.2 , when you estimate 10 MIs, are the samples required to evaluate these MIs seen during training ? \n- In Fig.4 c), could you also include training time when training is required. \n- The results on toy datasets really highlight the benefits of the method, but not its limitations. For instance, it would be interesting to show results by varying the latent dimensions of the autoencoder. Especially, what happens if the latent dimension is smaller than $k$ ?", "limitations": "The authors address the limitation of their work at the end of the manuscript.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "The approach is theoretically founded, and I appreciate that the authors also focus on making it practical. For instance, their model is intentionally simple to avoid doing large parameter sweep. The claim of the paper are well validated on toy datasets for which we can vary the dependence and ambient dimension, while having a close form solution of the MI.", "weaknesses": "- I would advise reporting the standard deviation in all tables (error bar in figures).\n- The presentation of the results could be improved, for example a few figures have labels that are not readable (e.g. Fig.7).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for this review. Below we address specific questions and concerns.\n\n***“I would advise reporting the standard deviation in all tables (error bar in figures).”***\n\nThank you for raising this. We apologize for omissions and will include s.d. where applicable in our revision. We have included a version of Figure 4 with s.d. as Rebuttal Figure (RF) 4\\. \n\n***“The presentation of the results could be improved, for example a few figures have labels that are not readable (e.g. Fig.7).”***\n\nWe apologize for the lack of clarity. We will increase font size in our revision.\n\n***“In Fig. 2, when you estimate 10 MIs, are the samples required to evaluate these MIs seen during training?”***\n\nThe samples used to estimate MI are *not seen* during training (for LMI, InfoNCE, MINE). This is the case for all estimates in the paper. We will clarify this in our Appendix section on implementation details. Anecdotally, we find all neural estimators (including LMI) generally perform worse when estimates are made on training data.\n\n***“In Fig.4 c), could you also include training time when training is required.”***\n\nThe “runtime” column includes training time for all methods. We will clarify this in the caption.\n\n***\\[common with x45D\\] “\\[...\\] For instance, it would be interesting to show results by varying the latent dimensions of the autoencoder. Especially, what happens if the latent dimension is smaller than $k$?”***\n\nThank you for this insightful question. We did not adequately address this in our initial submission, and will include discussion in our revised manuscript. Below, we outline some general principles for choosing latent space size, give a heuristic approach to choosing latent space size, and show that even suboptimal latent space size often yields state-of-the-art performance in high-dimensional settings.\n\n***General principles for choosing latent space size***\n\nThere is a clear tradeoff which arises when changing latent space size. As the latent space gets larger, the capacity of the compressed representation increases, and we might expect that representation quality increases (with the caveat that representation quality can be limited by sample sparsity). However, as the latent space gets larger, the MI estimate in latent space becomes more difficult. As such, the ideal choice is the smallest possible latent space size which captures the dependence structure of the variables. In practice, this size is hard to determine rigorously.\n\n***Heuristic approach to choosing latent space size***\n\nFrom Theorem 2 of our paper, a simple extension of DPI, we know that $I(Z\\_x; Z\\_y) \\\\leq I(X; Y)$. With the caveat that the inequality is not guaranteed to hold for the estimated $\\\\hat{I}(Z\\_x; Z\\_y)$, we can reason that the parameter choices that maximize $\\\\hat{I}(Z\\_x; Z\\_y)$ are likely ideal. As such, one sensible way to choose a latent space size is to try several, and use that which yields the largest estimate. This is computationally reasonable to do, and we will include this suggestion for practitioners in the Discussion of our revised paper.\n\n***Empirically exploring sensitivity to latent space size***\n\nIn our paper, we find that a suboptimal choice of latent space size can still be effective. Every single estimate in the main text of the paper uses 8 latent dimensions per variable, despite datasets with dependence structure of varying intrinsic dimensionality. For example, in Figure 2, LMI with 8 latent dimensions performs better than existing techniques across datasets where the ideal choice ranges from 1 to 9 dimensions. Notably, even in the case where the LMI latent space is smaller than the intrinsic dimensionality of the dependence structure (the 9 intrinsic dimension column), it outperforms alternate approaches in high ambient dimensions.\n\nHere, we will give a more direct example for multivariate Gaussian data with 4 dimensional dependence structure, in 1000D ambient space, with ground truth MI of 1 bit. We estimate MI from 5e3 samples using various latent space sizes, and compare this to MINE and InfoNCE. \n\n|  | LMI-2 | LMI-4 | LMI-6 | LMI-8 | MINE | InfoNCE |\n| :---- | :---- | :---- | :---- | :---- | :---- | :---- |\n| Estimate | 0.295732 | 0.719762 | 0.670632 | 0.686974 | \\-0.000001 | \\-0.001417 |\n\nThe most accurate estimate comes from LMI with optimal choice of 4 latent dimensions, but all tested choices (including the 2D latent space which cannot fully capture dependence) improve over MINE and InfoNCE. If we use the heuristic approach, we would choose the optimal latent space size. If we had arbitrarily chosen 8, we would be within 5% of the 4 latent dimension estimate. \n\nAs another example, in real world data, we show the sensitivity of the $I(K;T)$ estimate from the protein sequence embedding dataset described in Section 4.1 (RF3). We also visualize how model loss scales with latent space size, and show that both the estimate and model loss stabilize around 8 latent dimensions.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "sX4tRi1ssU", "reviewer": "Reviewer_x45D", "summary": "This paper explores the idea that underlying low-dimensional structure in high-dimensional data can be exploited to approximate mutual information (MI) efficiently and with a reasonable number of samples.  The approach learns a low-dimensional embedding of high-dimensional data using a neural network architecture similar to an autoencoder.  The MI of the resulting low-dimensional embedding is estimated using a nearest-neighbor approximation.  The paper includes extensive experimental evaluation including synthetic Gaussian data, resampled non-Gaussian data, and two examples in the computational biology domain.", "questions": "See \"Weaknesses\" section", "limitations": "An obvious limitation of this work is not explicitly discussed.  The latent dimension is introduced as a design variable that must be known (estimated) to compute the MI measure.  Sensitivity to the choice of latent dimension is not directly addressed.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "This is a very nicely-written paper.  The problem addressed is important and the ideas and approach are interesting.  While the idea is relatively straightforward, the authors provide extensive experimentation to convince the reader that their approach is useful.  This reviewer particularly appreciates the inclusion of domain-specific open problems as evidence of the efficacy of the proposed approach.", "weaknesses": "The synthetic multivariate Gaussian data evaluation is lacking in some respects.  The authors consider a only very specific form of dependence, namely linear dependence that obeys the stochastic process discussed in Sec. 3.1 (second paragraph) and that has low intrinsic dimensionality.  Because Gaussians model only linear dependence it is likely that MI estimation based on linear projections (such as Sliced MI) would likely perform well in this setting.  Yet the experiments fail to provide comparison to Sliced MI or related methods.  A more thorough set of synthetic data experiments, including those that exploit low-dimensional structure, is proposed in (Paweł et al. 2024).\n\nThe motivation and justification for the proposed approach is a little misleading.  In particular, the authors misstate the findings in McAllester and Stratos (2018); this reviewer suspects the intended reference is McAllester and Stratos (2020).  In that paper the authors do not demonstrate that statistical efficiency of estimators (MINE and InfoNCE) depend on dimension of the random variables, but rather they show a strong dependence on the value of MI.  In particular they show that sample complexity scales exponentially with the value of MI.  This finding is not consistent with the claim (L52) \"_More generally, it has been shown that no MI estimator can be accurate without making strong assumptions about the distribution..._\". In fact the estimators referred to are consistent in the infinite sample limit, but they can require a prohibitive number of samples for accurate estimates.\n\nA secondary motivation that is questionable is that the proposed method (L62) makes \"_strong, yet reasonable, assumptions about data which enable tractable MI estimation._\" In fact the MI of the latent embedding is not tractable.  Indeed, the authors use a nonparametric estimator (KSG) to approximate the latent state MI.\n\n**Detailed Comments**\n* Sec. 3.1 : The stochastic process described seems overly complicated.  Isn't it equivalent to restrict the covariance matrix for a multivariate Gaussian to a known structure?\n* Fig. 3c : This figure is not very informative due to the failure of the majority of methods ; perhaps consider a larger epsilon value?\n* L163 : Change \"MIME\" to \"MINE\"\n* Sec. 3.2 (last paragraph) : Change figure references from Fig. 3 to Fig. 4\n* Sec. 4.2 : It isn't totally clear what $X_{6'}$ is referring to as it is not explicitly defined\n\n\n**References**\nCzyż, Paweł, et al. \"Beyond normal: On the evaluation of mutual information estimators.\" _Advances in Neural Information Processing Systems_ 36 (2024).\n\nMcAllester, David, and Karl Stratos. \"Formal limitations on the measurement of mutual information.\" _International Conference on Artificial Intelligence and Statistics_. PMLR, 2020.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for this thoughtful review. Below we address specific points:\n\n***“\\[...\\] the experiments fail to provide comparison to Sliced MI or related methods.”***\n\nWe appreciate the reviewer’s insight that sliced MI is likely to capture dependence well in high dimensions, particularly for Gaussian data (where CCA is sufficient \\[1\\]). We have now applied sliced MI to a subset of our multivariate Gaussian benchmark, shown in Rebuttal Figure (RF) 7\\. While the relative accuracy of sliced MI is near perfect, its absolute accuracy is low. This is because sliced MI is not a method to estimate MI — it is a distinct measure of dependence with different interpretation from classical MI (as discussed in \\[4\\]). We fully agree with the broader point that the more interesting use-case of LMI is for non-Gaussian data, and we explore this next.\n\n***“A more thorough set of synthetic data experiments, including those that exploit low-dimensional structure, is proposed in (Paweł et al. 2024).”***\n\nWe have now included an additional benchmark which considers high-dimensional distributions diversified using the transformations from \\[1\\] (RF6). The results are qualitatively similar to our Fig. 2 (LMI dramatically outperforming other methods in the high ambient dimensions, low intrinsic dimensions regime).\n\nIn our initial submission, we had opted to develop our own “realistic” non-Gaussian benchmarking approach (Section 3.2, Fig. 4\\) rather than use the tasks of \\[1\\], because the maximum dimensionality studied in \\[1\\] is 25, the vast majority of tasks were $\\\\leq 5$ dimensions, and the distributions, while diverse, are still analytically defined. \n\nNow, rather than using the specific tasks from \\[1\\], we have developed complementary high-dimensional benchmarks by applying the transforms proposed in \\[1\\] to our multivariate Gaussians from Fig. 2\\. This allows us to create versions of the existing Fig. 2 with “half-cube”, “asinh”, and “uniform marginal” distributions (as defined in \\[1\\]). \n\nDue to time constraints in the rebuttal period, we explore only the most challenging settings from Fig. 2, those with 1000 ambient dimensions and 1-9 intrinsic dimensions. We find that LMI performance is similar for untransformed data and all three transformations (RF6). We will include these results, as well as results in varying ambient dimensions, in our revised paper.\n\n***“In particular, the authors misstate the findings in McAllester and Stratos (2018); this reviewer suspects the intended reference is McAllester and Stratos (2020). In that paper the authors do not demonstrate that statistical efficiency of estimators (MINE and InfoNCE) depend on dimension of the random variables, but rather they show a strong dependence on the value of MI.”***\n\nWe are grateful that the reviewer caught our mis-citation. We will correct the reference to McAllester and Stratos (2020) and remove its citation in L48 (about curse of dimensionality, leaving only a reference to \\[4\\]).\n\n***This finding is not consistent with the claim (L52) \\[...\\] In fact the estimators referred to are consistent in the infinite sample limit, but they can require a prohibitive number of samples for accurate estimates.***\n\nThank you for pointing this out. We will correct the statement to “*More generally, it has been shown that no technique can accurately estimate MI **from finite samples** without making strong assumptions...*\" with the assumptions being, at a minimum, that $I(X, Y) \\< O(\\\\log(N))$ for $N$ samples.\n\n***“In fact the MI of the latent embedding is not tractable. Indeed, the authors use a nonparametric estimator (KSG) to approximate the latent state MI.”***\n\nWe apologize for the wording. We had meant tractable in the more colloquial sense, we will replace it with “feasible” for clarity.\n\n***“\\[...\\] Isn't it equivalent to restrict the covariance matrix for a multivariate Gaussian to a known structure?”***\n\nWe agree that the process described in Section 3.1 is unsatisfyingly complicated. We’ll update and simplify exposition in the revised manuscript. It is indeed equivalent to restrict the covariance matrix, for example to the following structure\n\n  * $\\\\text{Cov}(X\\_i, Y\\_i) \\= \\\\rho$ for $i \\\\in \\\\{1, …, k\\\\}$,   \n  * $\\\\text{Cov}(X\\_i, X\\_i) \\= \\\\text{Cov}(Y\\_i, Y\\_i) \\= 1$  for $i \\\\in \\\\{1, …, d\\\\}$,   \n  * $\\\\text{Cov}(X\\_i, X\\_j) \\= \\\\text{Cov}(Y\\_i, Y\\_j)  \\= \\\\text{Cov}(X\\_i, Y\\_j) \\= 0$ for all $i \\\\neq j$  \n  * $\\\\text{Cov}(X\\_i, Y\\_i) \\= 0$ for $i \\\\in \\\\{k+1, …d\\\\}$\n\n***“Fig. 3c : This figure is not very informative due to the failure of the majority of methods; perhaps consider a larger epsilon value?”***\n\nThank you for this suggestion. We have included a version with $\\\\epsilon \\=0.8$ in RF5. We caution that this corresponds to an 80% error in the estimate.\n\n***“L163 : Change \"MIME\" to \"MINE\"; Sec. 3.2 (last paragraph) : Change figure references from Fig. 3 to Fig. 4”***\n\nThank you for catching these typos. We will fix them in the revision.\n\n***“Sec. 4.2 : It isn't totally clear $X\\_{6’}$ what is referring to as it is not explicitly defined”***\n\nWe apologize for the lack of clarity. In our revision we will define it as gene expression state in the separated well of cells on day 6 after harvest.\n\n***“An obvious limitation of this work is not explicitly discussed. The latent dimension is introduced as a design variable that must be known (estimated) to compute the MI measure. Sensitivity to the choice of latent dimension is not directly addressed.”***\n\nWe fully agree that this is an important aspect of our work which warranted more careful treatment. In our response to Reviewer 8utg, we empirically study and discuss sensitivity to latent dimension. Due to space constraints, we kindly point the reviewer to the response to 8utg, specific section denoted \\[common with x45D\\].", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We deeply appreciate the thoughtful feedback shared by the reviewers. In our responses to reviewers, we share some additional analyses, discussions, and clarifications (companion figures to responses are included in the attached .pdf). Overall, we think we have managed to address all of the concerns and questions raised. Below we summarize the most significant new content:

1. We study the failure of LMI for distributions with “symmetry and exclusivity”, show that such distributions are also troublesome for existing estimators, then propose and empirically evaluate some potential solutions. \[in response to NtLr\]  
2. We explicitly discuss tradeoffs in the choice of latent space size, give a heuristic approach to choosing latent dimension, and discuss the consequences of choosing suboptimally. \[Response to 8utg, x45D\]  
3. We show that LMI outperforms alternatives in an MI estimation problem when only one variable requires compression (MI of MNIST image and label). \[Response to 13JG\]  
4. We show that the performance of LMI on multivariate Gaussian data remains qualitatively similar across a more diverse set of distributions, by constructing high dimensional analogs of the MI estimator benchmark tasks in \[1\]. \[Response to x45D\]

**Below are the references used across all responses, centralized due to space constraints.**

\[1\] Czyż, P., Grabowski, F., Vogt, J. E., Beerenwinkel, N. & Marx, A. Beyond normal: On the evaluation of mutual information estimators. *arXiv \[stat.ML\]* (2023).   

\[2\] McAllester, D. & Stratos, K. Formal Limitations on the Measurement of Mutual Information. in *International Conference on Artificial Intelligence and Statistics* 875–884 (PMLR, 2020).  

\[3\] Gao, S., Steeg, G. V. & Galstyan, A. Efficient estimation of mutual information for strongly dependent variables. *arXiv \[cs.IT\]* (2014).   

\[4\] Goldfeld, Z. & Greenewald, K. Sliced mutual information: A scalable measure of statistical dependence. *arXiv \[cs.IT\]* (2021).

Meta-review: This paper focuses on learning representations of X and Y so that the MI can be more easily estimated. While existing techniques such as MINE have long done this, this paper provided novel criteria using reconstruction losses to ensure that the learned representations preserve the necessary information, and this procedure outperforms MI estimation baselines. I support the consensus view of the reviewers in favor of acceptance. Authors - please remember to add the promised additions, in particular the section on limitations.
Decision: Accept (spotlight)