Interpretable models for scRNA-seq data embedding with multi-scale structure preservation

This paper addresses the limitations in trustworthiness and evaluation of current single-cell RNA sequencing dimensionality reduction methods by introducing ViScore, a robust scoring framework, and ViVAE, an interpretable deep learning model designed to preserve multi-scale data structures.

The Gist

Imagine you have a massive, chaotic library containing millions of books. Each book represents a single cell in your body, and the pages inside contain thousands of words (genes) that tell you what that cell is doing. Trying to read every single page to understand the library is impossible.

To make sense of this, scientists use Dimensionality Reduction (DR). Think of this as a magical librarian who creates a 2D map of the library. On this map, similar books are placed close together, and different books are far apart. This map helps researchers see patterns, like "all the biology textbooks are in the north wing" or "the fiction section is a long, winding path."

However, there's a problem: The maps are often misleading.

Some popular mapping tools (like t-SNE and UMAP) are great at grouping similar books together, but they often tear the library apart, creating artificial gaps between sections that should be connected. Others (like PCA) keep the whole library connected but might squash distinct groups so tightly you can't tell them apart. It's like trying to flatten a crumpled piece of paper: if you smooth out the wrinkles, you might stretch the drawing so much that the face looks distorted.

This paper introduces two new tools to fix these issues: ViScore (a quality checker) and ViVAE (a better map-maker).

1. The Problem: "Local" vs. "Global" Confusion

Imagine you are looking at a city map.

• Local Structure: This is like looking at your immediate neighborhood. Are the coffee shops near the park? Are the houses close to the school?
• Global Structure: This is the big picture. Is the city laid out in a grid? Is the downtown area far from the suburbs?

Current tools are great at showing the neighborhood (local) but often mess up the city layout (global). They might put the suburbs right next to downtown just because the streets look similar, even though they are miles apart in reality.

2. The Solution: ViScore (The Quality Control Inspector)

Before you trust a map, you need to know if it's accurate. The authors created ViScore, a scoring system that acts like a rigorous building inspector.

• The Old Way: Inspectors used to just check if the houses were close to each other (local) or if the whole city was connected (global), but they used arbitrary rules that didn't work well for huge datasets.
• The ViScore Way: This tool checks the map at every possible scale simultaneously. It asks: "If I zoom in, are the neighbors right? If I zoom out, is the city layout correct?"
• The Analogy: Imagine a video game map. ViScore doesn't just check if your character is standing next to a tree; it checks if the tree is in the right forest, if the forest is in the right mountain range, and if the mountain range is in the right continent. It gives a single, fair score to tell you how trustworthy the map is.

3. The Solution: ViVAE (The Master Cartographer)

Once they had a way to measure quality, they built a new map-maker called ViVAE.

• How it works: Think of ViVAE as a smart, flexible artist. It uses a technique called a "Variational Autoencoder" (a type of AI) that learns to compress the complex 3D shape of the library into a 2D map.
• The Secret Sauce: Most artists focus on either the details (local) or the big picture (global). ViVAE uses a special "stochastic MDS" technique (a fancy mathematical rule) that forces the artist to get both right at the same time.
• The Result: The maps ViVAE creates show distinct groups of cells (like different cell types) clearly separated, but they also preserve the smooth, continuous paths between them (like how a stem cell slowly turns into a blood cell). It avoids the "island effect" where groups are artificially isolated.

4. The "Distortion Detector" (Encoder Indicatrices)

Even the best maps can have hidden flaws. The authors added a cool feature called Encoder Indicatrices.

• The Analogy: Imagine you have a rubber sheet with a perfect circle drawn on it. When you stretch the sheet to make the map, the circle might turn into an oval, or it might get stretched in one direction.
• What it does: ViVAE draws these "circles" (which are actually tiny mathematical shapes) all over the map. If the circle stays round and the same size, the map is perfect. If the circle turns into a long, thin oval, the tool tells you: "Hey, this part of the map is stretched out! Don't trust the distances here."
• Why it matters: It lets scientists see exactly where the map is lying to them, so they don't make false conclusions about how cells are related.

Summary

In simple terms, this paper says:

1. Current maps of single-cell data are often distorted. They either break things apart or squish them together too much.
2. We built a new ruler (ViScore) to measure exactly how good a map is, checking both the small details and the big picture fairly.
3. We built a new map-maker (ViVAE) that creates maps that are accurate at all scales, keeping both the neighborhoods and the city layout correct.
4. We added a distortion detector so scientists can spot exactly where the map is warped.

This helps biologists trust their data more, ensuring that when they see a pattern in the map, it's a real biological truth and not just an artifact of the software.

Technical Summary

1. Problem Statement

Single-cell RNA sequencing (scRNA-seq) generates high-dimensional (HD) data that requires dimensionality reduction (DR) for visualization and analysis. Current state-of-the-art methods face two primary challenges:

• The Local-Global Trade-off: Neighbor embedding (NE) methods like t-SNE and UMAP excel at preserving local structures (clusters) but often distort global structures (relative positions of clusters, trajectories, and developmental branches) due to "locality bias." Conversely, linear methods (PCA) or global optimizers (MDS) preserve global structure but fail to resolve fine-grained local clusters.
• Lack of Rigorous Evaluation: Existing evaluation metrics for DR are often flawed. They rely on arbitrary thresholds for defining "local" vs. "global," use downsampling that sacrifices rigor, or depend on downstream tasks (clustering/classification) which act as proxies rather than direct measures of embedding fidelity.
• Limited Interpretability: Deep learning models (like VAEs) are often "black boxes," making it difficult to diagnose specific geometric distortions (e.g., artificial stretching or shrinking of cell populations) introduced during the embedding process.

2. Methodology

The authors propose a two-pronged solution: a new evaluation framework (ViScore) and a new deep learning model (ViVAE).

A. ViScore: A Robust Evaluation Framework

ViScore addresses the evaluation gap by introducing scalable, multi-scale metrics:

• RNX Curve Approximation: The authors utilize Rank-based Neighborhood Preservation (RNX) curves to quantify structure preservation across all scales. To overcome the $O(N^2 \log N)$ computational complexity that makes RNX infeasible for large scRNA-seq datasets, they developed a novel approximation algorithm using Vantage-Point (VP) Trees. This allows for efficient sampling of representative points to estimate global and local scores with high precision.
• Dual Scoring Metrics:
  • Global SP: Calculated as the Area Under the Curve (AUC) of the RNX curve with linear scaling, treating all scales equally.
  • Local SP: Calculated using a logarithmic rescaling of the RNX curve, up-weighting small-scale neighborhoods.
• Supervised Metrics (xNPE & NCPs):
  • Extended Neighbourhood-Proportion-Error (xNPE): A supervised metric that compares the distribution of "self" vs. "non-self" neighbors for specific cell populations between HD input and LD embeddings using Earth Mover's Distance (EMD).
  • Neighbourhood Composition Plots (NCPs): Qualitative visualizations showing the composition of neighbors around a population of interest across different neighborhood sizes to diagnose specific mis-embeddings.

B. ViVAE: A Deep Learning Model for Multi-Scale Preservation

ViVAE is a regularized Variational Autoencoder (VAE) designed to balance local and global structure preservation:

• Architecture: A standard VAE encoder-decoder structure (or an encoder-only variant) that learns non-linear transformations.
• Input Denoising: The input data is pre-processed using a mean-shift-based denoising algorithm (k-nearest neighbor graph) to prevent the model from fitting to noise.
• Stochastic-MDS Regularizer: The core innovation is the integration of a Stochastic Multidimensional Scaling (MDS) loss function ($l_{MDS}$).
  • Unlike standard VAEs that only minimize reconstruction error and KL divergence, ViVAE adds a term that minimizes the difference between normalized pairwise distances in the HD input and the LD embedding.
  • This is implemented using a "stochastic quartet" approach (groups of 4 points) to make the loss differentiable and scalable.
  • The total loss is: $L = l_{recon} + l_{KLD} + 100 \cdot l_{MDS}$.
• Encoder Indicatrices (EIs): A novel interpretability tool derived from differential geometry.
  • EIs visualize local distortions in the latent space by mapping a small hypersphere from the HD input space through the encoder.
  • The resulting shape (visualized as ellipses) indicates the direction and magnitude of local expansion, shrinkage, or stretching, allowing researchers to pinpoint artificial artifacts.

3. Key Contributions

1. ViScore Framework: A scalable, fair, and rigorous benchmarking tool that eliminates arbitrary "local/global" cut-offs and enables evaluation on massive datasets (up to ~1 million cells) via VP-tree approximation.
2. ViVAE Model: A parametric DR model that achieves a superior balance between local and global structure preservation compared to t-SNE, UMAP, and standard VAEs, without requiring a specific "scale" hyperparameter.
3. Encoder Indicatrices (EIs): A model-agnostic QC tool for visualizing and quantifying geometric distortions in differentiable DR models, enhancing interpretability.
4. Supervised Evaluation Tools: Introduction of xNPE and NCPs for detailed, population-level error analysis in labeled datasets.

4. Results

The authors evaluated ViVAE against 12 other DR methods (including PCA, t-SNE, UMAP, TriMap, PaCMAP, SQuad-MDS, and standard VAEs) across 8 diverse scRNA-seq datasets (e.g., zebrafish development, human immune cells, mouse retina).

• Structure Preservation:
  • ViVAE consistently placed on the Pareto front of the Local vs. Global SP trade-off.
  • While t-SNE achieved the highest Local SP and SQuad-MDS the highest Global SP, ViVAE provided the best balanced performance, outperforming standard VAEs significantly in Global SP due to the MDS regularizer.
• Case Study 1: Developmental Trajectories (Farrell Dataset):
  • In zebrafish embryogenesis data, NE methods (t-SNE/UMAP) artificially fragmented continuous developmental trajectories into disconnected islands.
  • ViVAE preserved the smooth gradients and bifurcation points of developmental lineages more faithfully, validated by lower xNPE scores and continuous transition distances.
• Case Study 2: Non-Developmental Immune Data (Reed Dataset):
  • On adult human breast immune cells, ViVAE successfully separated major cell compartments while maintaining correct relative positions between populations (e.g., NK cells vs. T cells).
  • NCPs revealed that t-SNE and UMAP exaggerated the proximity of certain populations (e.g., NK T cells to NK cells), whereas ViVAE maintained a more accurate neighborhood composition.
• Distortion Analysis:
  • EIs demonstrated that standard VAEs suffer from "run-away" effects (directed stretching) and scale inconsistencies. ViVAE, particularly the encoder-only variant, significantly reduced these artifacts, producing a more uniform latent space.

5. Significance

This work significantly advances the field of single-cell bioinformatics by:

• Improving Trustworthiness: By providing rigorous, scalable metrics (ViScore), it allows researchers to objectively assess whether an embedding is trustworthy for their specific biological question (local clustering vs. global trajectory).
• Bridging the Gap: ViVAE offers a "best of both worlds" solution, capturing the cluster separation of NE methods while preserving the global topology required for trajectory inference and developmental biology.
• Enhancing Explainability: The introduction of Encoder Indicatrices moves DR from a "black box" to an interpretable process, enabling users to detect and understand specific geometric artifacts before drawing biological conclusions.
• Scalability: The VP-tree approximation makes high-fidelity evaluation feasible for the increasingly large datasets generated by modern atlases (e.g., Tabula Sapiens), addressing a critical bottleneck in current workflows.

The authors have made both ViScore and ViVAE available as open-source Python packages, facilitating immediate adoption in the community.