Distance-Preserving Representations for Genomic Spatial Reconstruction

TL;DR

dp-VAE tackles the challenge of recovering spatial tissue context from gene expression data by introducing a distance-preserving regularizer that enforces latent representations to reflect spatial geometry learned from reference spatial-transcriptomics datasets.During inference, spatial coordinates are recovered or imputed by solving a distance-geometry optimization over the latent embeddings, avoiding the need for spatial data in the input and enabling broad applicability.The work provides a theoretical connection between the distance-preserving loss and distortion/bi-Lipschitz properties, and demonstrates robust performance, out-of-sample generalization, and transfer-learning potential across 27 public datasets, with observed limitations in zero-shot cross-domain scenarios that can be mitigated by fine-tuning.Overall, dp-VAE offers a scalable framework to integrate spatial context into genomics analyses, broadening access to spatial insights for diverse single-cell studies without relying on spatial measurements at inference time.

Abstract

The spatial context of single-cell gene expression data is crucial for many downstream analyses, yet often remains inaccessible due to practical and technical limitations, restricting the utility of such datasets. In this paper, we propose a generic representation learning and transfer learning framework dp-VAE, capable of reconstructing the spatial coordinates associated with the provided gene expression data. Central to our approach is a distance-preserving regularizer integrated into the loss function during training, ensuring the model effectively captures and utilizes spatial context signals from reference datasets. During the inference stage, the produced latent representation of the model can be used to reconstruct or impute the spatial context of the provided gene expression by solving a constrained optimization problem. We also explore the theoretical connections between distance-preserving loss, distortion, and the bi-Lipschitz condition within generative models. Finally, we demonstrate the effectiveness of dp-VAE in different tasks involving training robustness, out-of-sample evaluation, and transfer learning inference applications by testing it over 27 publicly available datasets. This underscores its applicability to a wide range of genomics studies that were previously hindered by the absence of spatial data.

Figures (14)

• Figure 1: The model architecture of dp-VAE with vanilla VAE as its backbone. During the training stage, the distance-preserving loss ($\mathcal{L}_{\textrm{DP}}$) enforces distance preservation between the spatial domain ($S$) and the latent space ($Z$), enforcing the encoder network (e.g., parameterized as a fully connected neural network) to capture meaningful signals of the tissue spatial characteristics of the provided gene expression ($X$). During the inference/test stage, the extracted latent representations are fed to an optimizer, such as multidimensional scaling (MDS), if it is a spatial reconstruction task instead of a spatial context imputation task.
• Figure 2: Pearson correlation between the pairwise distance of gene expression ($X$) and the pairwise distance ($S$) of spatial location, computed on 27 spatial transcriptomics datasets. The red line indicates zero correlation, and the error bars are the $95\%$ confidence intervals of the correlation coefficient calculated using Fisher's Z transformation.
• Figure 3: An illustrative example of the masking procedure is shown. Suppose we have three spatial locations, so the mask matrix $\mathbf{G}$ is $3 \times 3$. One user-specified connection in $S$ is masked during training by setting the corresponding entries in $\mathbf{G}$. Consequently, the distance-preserving regularization in $\mathcal{L}_{\rm DP}$ is not enforced between those two locations.
• Figure 4: Illustration of forward and inverse mappings. The mapping that recovers the coordinates $S$ from pairwise distances $D_S$ is not unique or non-deterministic when $\left| \mathcal{I} \right| < 3$.
• Figure 5: Training stability for different choices of regularization coefficients for the distance-preserving loss over $1 \times 10^3$ epochs of training, averaged over 27 datasets.

Theorems & Definitions (7)

• Theorem 1
• Lemma 2
• proof
• Theorem 3
• proof
• Lemma 4: KL divergence of multivariate normals
• proof : Proof of \ref{['lem:kl-normal']}