Approximating mutual information of high- dimensional variables using learned representations

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 of two kinds: the ambient dimensionality 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, we consider variables with $d=10$ to $d=5\cdot 10^{3}$ 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 $d$ bivariate Gaussians and concatenate the first components to construct samples of $X$, and concatenate the second components to construct $Y$. By choosing the covariance of each of the bivariate Gaussians, $I(X,Y)$ and $k$ can be tuned. To enforce a $k$-dimensional dependence structure, we can choose covariance matrices such that $\text{Cor}(X_{i},Y_{i})=\rho\textbf{Heav}(i-k)$ where Heav is the Heaviside step function, and $\text{Cor}(X_{i},Y_{j})=0\ \forall i\neq j$. The exact sampling procedure we use for these experiments is given in Appendix A.3.1, Algorithm 4.

Results of benchmarking MI estimators using these synthetic datasets are given in Fig. 2. For estimates from $N=2\cdot 10^{3}$ 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 $d=5\cdot 10^{3}$ ambient dimensions (Fig. 2d, 2i). The faithfulness of LMI approximation instead degrades with increasing $k$. Nonetheless, LMI approximation gives more absolutely and relatively accurate MI estimates than alternatives for 83% and 87% of tested settings respectively (Fig. 2e, 2j).

While the empirical results on multivariate Gaussians are 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 a technique for creating benchmark datasets by resampling real-world data. Briefly, we use correspondences between discrete labels and complex data (i.e. digit labels and digit images in MNIST) to transform simple discrete distributions into realistic high-dimensional distributions.

Specifically, we draw samples from a bivariate Bernoulli vector, $\mathbf{L}=[L_{x},L_{y}]\in\{0,1\}^{2}$ with prescribed pairwise correlation $\text{Cor}(L_{x},L_{y})=\rho$, where each value corresponds to a discrete label of a set of samples in a high-dimensional dataset (e.g. $0$ and $1$ correspond to images of 0s and 1s in MNIST). For each sample of $\mathbf{L}$, we replace each component with a random (without replacement) high-dimensional sample matching the label. For the example of MNIST, this transforms samples from $\mathbf{L}$ into pairs of images of 0s and 1s, represented as samples of random vectors $X,Y\in\mathbb{R}^{784}$.

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 $H(L_{x}|X)=H(L_{y}|Y)=0$ then $I(X;Y)=I(L_{x};L_{y})$ (shown in Appendix A.3.2, Theorem 4). And using our knowledge of $\rho$, $I(L_{x};L_{y})$ can be analytically computed.

We resample two different source datasets: (1) “binary” subset of MNIST, containing only images of 0s and 1s, with $5000$ samples and $784$ dimensions and (2) embeddings of a subset of protein sequences from E. coli and A. thaliana proteins, with $4402$ samples and $1024$ dimensions. For both source datasets, we validate the $I(L_{x};L_{y})\approx I(X;Y)$ approximation in Appendix A.3.3.

For each source dataset, we generate 200 benchmark datasets with true MI ranging from 0 to 1 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 [16]. On protein embedding datasets, variational estimators nearly always fail to estimate nonzero values – resulting in a rank correlation below 0 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. In contrast, LMI approximation yields estimates consistently close to the ground truth, with high relative and absolute accuracy.

Pretrained protein language models (pLMs) have recently seen widespread use, largely due to their convenient representations of protein sequences (vectors in $\mathbb{R}^{N}$, typically with $N\approx 10^{3}$) which can be used for transfer learning on downstream tasks [31, 32, 33]. While it is known that pLM sequence embeddings contain significant information about protein structure [34], it is not clear how well existing pLMs encode functional information. Recent work has shown that pLMs fail to capture some important functional properties, such as thermostability [35]. It remains unclear the extent to which pLM embeddings contain information about protein-protein interactions, which are essential to protein function. Here, we use an information-theoretic approach to quantify interaction information contained in 1024-dimensional sequence embeddings from the ProtTrans5 model [32].

We study two types of protein-protein interactions: kinase-target and ligand-receptor interactions. For both, the OmniPath database [36] contains lists of thousands of annotated pairs of interacting proteins. We consider each annotated pair to be a sample from a joint “interaction” distribution over pairs of sequence embeddings. For example, for kinase-target interactions, we consider kinase and target sequence embeddings as random variables $K\in\mathbb{R}^{1024},T\in\mathbb{R}^{1024}$, with joint distribution $P_{KT}(k,t)$. Then, using joint samples $\{(k_{1},t_{1}),\ldots,(k_{N},t_{N})\}$ we estimate mutual information between interaction partners, $I(K;T)$. We analogously estimate $I(L;R)$ for ligand-receptor interactions.

If pLM embeddings capture interaction information, MI between interaction partners should be significantly above 0 bits. Applying LMI approximation, we estimate $I(L;R)\approx 2.8\text{ bits}$ and $I(K;T)\approx 0.8\text{ bits}$. To test the significance of these values, we estimated MI from shuffled data, and found $I_{\text{shuff}}(K;T)\approx I_{\text{shuff}}(L;R)\approx 0$ (both mean and standard deviation $<0.05$), across 20 random shuffles. These results indicate that pLM embeddings contain information about both types of interactions, and contain more information about ligand-receptor interactions than kinase-target interactions. In contrast, existing estimators yield far lower estimates, with MINE estimates indicating independence for both types of interactions (Fig. 5a, 5b). To validate the LMI estimates of dependence, we next operationally verify the presence of interaction information.

If protein-protein interactions can be predicted for a set of held-out proteins based on sequence embeddings, then sequence embeddings must contain interaction information. To see if this is the case, we extend LMI to predict protein interactions from sequence embeddings. For ligand-receptor prediction, our goal is to predict whether a held-out pair of sequence embeddings $(l,r)$ is an annotated ligand-receptor pair. One way to do this is estimating the log density ratio $\log\frac{P_{LR}(l,r)}{P_{L}(l)\cdot P_{R}(r)}$, and setting a threshold above which sequence pairs are predicted to be annotated interactions.

We make a simple modification to the KSG estimator to 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\cdot 10^{4}$ 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. 5c, 5d, with mean AUC-ROC scores of 0.87 and 0.74 for kinase-target and ligand-receptor interactions respectively. These results demonstrate that protein interactions can be predicted better than random chance using ProtTrans5 embeddings, suggesting that pLM 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.

Single cell RNA sequencing (scRNA-seq) measures the expression of $g\approx 10^{4}$ genes in single cells, which can be thought of samples of a gene expression state variable $X\in\mathbb{R}^{g}$. These samples can be used to infer a probability distribution over gene expression states, $P_{X}$. To study the dynamics of gene expression, one approach is to make measurements at multiple timepoints $t_{1},...t_{N}$, which can be thought of as samples of random variables $X_{i}\in\mathbb{R}^{g}$. 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 $P_{X_{1},\ldots,X_{N}}$ [37].

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 a cell can be formally modeled as a Markov chain $X_{i}\to X_{i+1}$. In some cases, cell behavior may be a function of hidden variables resulting in non-Markovian dynamics. Using the data processing inequality (DPI) [7], we know that if gene expression dynamics are Markovian, $I(X_{i};X_{i+1})\geq I(X_{i};X_{i+2})$, and if the DPI is does not hold, then gene expression states are non-Markovian. This can, in principle, be tested using samples from the joint distribution $P_{X_{1},..,X_{N}}$. Due to the difficulty of high-dimensional MI estimation, previous work has used heuristic alternatives to indirectly test the Markovian assumption for lineage-traced scRNA-seq (LT-seq) data [38]. Here, we will use LMI to explicitly estimate high-dimensional MI from LT-seq data, and test the Markovian assumption for gene expression states.

We study a previously published LT-seq data set of in vitro differentiating mouse hematopoietic stem cells [38]. 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 $X_{2}\to X_{6}$, $X_{2}\to X_{6^{\prime}}$ and by the DPI we should have $I(X_{2};X_{6})\geq I(X_{6};X_{6^{\prime}})$. Using LMI, we estimate that $I(X_{2};X_{6})\approx 0.31\pm 0.02\text{ bits}$ and $I(X_{6};X_{6^{\prime}})\approx 0.98\pm 0.01\text{ bits}$ (mean $\pm$ SEM), over 20 random pairings of clonally related cells. This indicates that gene expression states are non-Markovian, in line with prior findings [38, 39].

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. 6b). Along the neutrophil trajectory, we quantatively compare cell fate information with neutrophil pseudotime computed by graph smoothing [38], and find that pMI begins to rapidly increase around pseudotime value of $35\cdot 10^{3}$, which is roughly aligned with the transition from granulocyte-myeloid progenitor to promyelocyte, as defined in [38].