🧬 biology

HDTree: Generative Modeling of Cellular Hierarchies for Robust Lineage Inference

HDTree is a novel generative modeling framework that utilizes a unified hierarchical codebook and quantized diffusion process to robustly infer cellular lineages by capturing tree relationships within a hierarchical latent space, thereby overcoming the scalability, stability, and posterior collapse limitations of existing methods.

Original authors: Zelin Zang, WenZhe Li, Yongjie Xu, Chang Yu, Changxi Chi, Jingbo Zhou, Zhen Lei, Stan Z. Li

Published 2026-05-08

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Zelin Zang, WenZhe Li, Yongjie Xu, Chang Yu, Changxi Chi, Jingbo Zhou, Zhen Lei, Stan Z. Li

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ⚕️ This is an AI-generated explanation of a preprint that has not been peer-reviewed. It is not medical advice. Do not make health decisions based on this content. Read full disclaimer

The Big Picture: Mapping the Family Tree of Cells

Imagine you are trying to understand how a single, tiny seed grows into a massive, complex oak tree with thousands of branches, leaves, and roots. In biology, this is like trying to trace how a single stem cell divides and changes into hundreds of different types of cells (like skin, muscle, or nerve cells).

Scientists call this a "lineage." The problem is that looking at millions of individual cells is like trying to understand the shape of a forest by looking at one leaf at a time. You need a map that shows not just the leaves, but how they connect to branches, which connect to the trunk.

Current computer tools for making these maps are often messy. They are either too slow, they get confused when the tree gets too deep, or they create "ghost branches" that don't actually exist in nature.

HDTree is a new computer program designed to draw this family tree perfectly. It doesn't just guess; it learns the rules of how cells grow and builds a robust, accurate map of their history.

The Problem with Old Tools

The paper explains that previous methods had three main flaws:

The "Specialized Room" Problem: Imagine a hotel where every single branch of the tree has its own unique, custom-built room with its own staff. If a branch is small and has very few guests (cells), that room sits empty and the staff gets confused. The system can't share knowledge between branches.
The "One-Step" Problem: Some old tools try to jump from a stem cell to a final cell type in one giant leap. But in real life, cells change gradually, like a caterpillar slowly turning into a butterfly. Jumping too fast causes the model to hallucinate or collapse.
The "Fragile" Problem: When the data is noisy (which it always is in biology), these old tools tend to break or overfit, meaning they memorize the noise instead of learning the actual pattern.

How HDTree Works: The Three Magic Ingredients

HDTree solves these problems using a clever combination of three ideas, which the authors describe as a unified framework.

1. The Universal "Codebook" (The Shared Dictionary)

Instead of building a custom room for every branch, HDTree uses one giant, shared dictionary (called a hierarchical codebook).

The Analogy: Imagine a library where every book (cell) is assigned a specific set of tags. Instead of writing a new book for every branch of the tree, the system uses the same set of tags for everyone.
Why it helps: Even if a branch of the tree is very small (sparse), it can still use the "knowledge" from the big, popular branches because they all share the same dictionary. This makes the system stable and scalable, no matter how deep the tree gets.

2. The "Diffusion" Process (The Slow Sculptor)

HDTree uses a technique called Diffusion, which is like a sculptor slowly chipping away stone to reveal a statue.

The Analogy: Imagine you have a block of marble covered in noise (static). Instead of trying to guess the final shape instantly, the system starts with pure chaos (noise) and slowly, step-by-step, removes the noise to reveal the cell.
The Connection to Biology: This mimics the Waddington Landscape, a famous biological concept where a cell rolls down a hill, making small, gradual decisions until it settles into a specific fate (like becoming a muscle cell). HDTree models this gradual "rolling down the hill" rather than teleporting the cell to its destination. This prevents the "collapse" issues seen in older models.

3. The "Soft" Contrastive Learning (The Neighborhood Watch)

To make sure the map is accurate, HDTree uses a technique that acts like a neighborhood watch.

The Analogy: It looks at a cell and says, "You look like your neighbors, so you should be close to them on the map." But it's "soft" because it understands that cells on the same branch of the tree are closer than cells on different branches, even if they look similar.
Why it helps: This ensures the computer doesn't just group similar-looking cells together; it respects the actual family tree structure, keeping the hierarchy intact.

What Did They Prove?

The authors tested HDTree on two types of data:

General Data: Standard computer vision datasets (like pictures of handwritten digits or clothes).
Single-Cell Data: Real biological data from human limbs, lung/heart development, and even the tiny worm C. elegans.

The Results:

Better Maps: HDTree drew the family trees more accurately than any previous method (including the current "champion" called TreeVAE).
Stability: It didn't break when the data was messy or when the tree was very deep.
Generative Power: Because it understands the "rules" of the tree, HDTree can actually invent new, fake cells that look real. If you ask it to generate a cell that is halfway between a stem cell and a muscle cell, it can create a plausible "in-between" cell. This acts as a "simulation" to check if the biological theory makes sense.

The Bottom Line

HDTree is like a master cartographer for the microscopic world of cells. By using a shared dictionary, a slow-and-steady sculpting process, and a smart neighborhood system, it creates a reliable, high-definition map of how life develops from a single cell into a complex organism.

The paper concludes that this tool allows scientists to see the "roads" cells take to become what they are, offering a more reliable way to understand the fundamental processes of life without needing massive, pre-trained super-computers.

Technical Summary: HDTree

Problem Statement
In single-cell research, accurately tracing and analyzing high-throughput differentiation trajectories is essential for understanding biological processes. This task requires robust modeling of the hierarchical structures governing cellular development. Traditional methods, which often combine dimensionality reduction, clustering, and regression, face significant limitations regarding computational cost, performance stability, and generative capacity. While Variational Autoencoder (VAE)-based approaches have improved efficiency, they typically rely on specialized network modules for each tree branch. This design leads to several issues:

Scalability and Stability: Deep branches with sparse samples cannot effectively leverage representation knowledge from other branches, limiting generalization and the ability to preserve global structure.
Overfitting: Independent encoder-decoder pairs at each branch node are prone to overfitting due to noise accumulation, especially when training data is limited.
Posterior Collapse: VAEs often suffer from posterior collapse, and their single-step generation mechanisms fail to naturally mirror the continuous, gradual refinement of cell states described by the Waddington landscape.

Methodology: HDTree
To address these challenges, the authors propose HDTree (Hierarchical vector quantized Diffusion Model), a unified generative framework designed for robust lineage inference. The core innovation lies in integrating hierarchical latent space encoding with a quantized diffusion process, decoupling tree complexity from network size.

The framework consists of three synergistic modules:

Unified Encoder & Hierarchical Tree Codebook (HTC): Unlike prior methods where parameters scale exponentially with tree depth, HDTree employs a single, shared encoder ( $E_{\theta_E}$ ) and a unified hierarchical codebook ( $C_W$ ). The codebook is structured as a binary tree where each node represents a learnable code vector in the latent space. Input data is mapped to a latent vector, which is then quantized into a hierarchical sequence of code vectors ( $s_i$ ) by traversing the tree from root to leaf. This design allows all branches to share the same codebook vectors, enabling sparse deep nodes to leverage representation knowledge from other branches while preserving global hierarchical structure.
Quantized Diffusion Decoder: To model continuous cell state transitions and avoid posterior collapse, HDTree utilizes a diffusion decoder ( $D_{\theta_D}$ ) based on Denoising Diffusion Probabilistic Models (DDPM). The generation process is conditioned on the quantized hierarchical code sequence ( $s_i$ ). Starting from Gaussian noise, the model iteratively refines the data through $T$ diffusion steps. This iterative denoising aligns with the Waddington landscape, modeling differentiation as a gradual refinement from coarse progenitors to fine-grained specialized states.
Loss Function: The model is optimized using a composite loss function:
- Soft Contrastive Learning Loss (SCL): Adapts contrastive learning to weight negative pairs by tree distance in hyperbolic space, preserving hierarchical granularity.
- Hierarchical Quantization Loss (HQL): Ensures robust tree-structured representations by aligning latent embeddings with multi-level code vectors and enforcing consistency between parent and child nodes.
- Diffusion Loss: Optimizes the variational lower bound (ELBO) to ensure faithful reconstruction of data while respecting the hierarchical conditions.

Lineage Inference
For trajectory analysis, the learned hierarchical tree is transformed into a weighted graph. Additional edges are added within the same depth using a k-nearest neighbors (KNN) approach to capture local semantic relationships. Developmental trajectories are inferred by finding the shortest path between a designated origin (e.g., stem cell) and destination (e.g., differentiated cell) on this graph, minimizing total edge weights.

Key Contributions

Methodological Innovation: HDTree introduces a unified framework that integrates hierarchical vector quantization with diffusion processes. By utilizing a single, shared hierarchical codebook, it effectively decouples tree complexity from network size, ensuring scalability to deep biological hierarchies without the instability of branch-specific modules.
Application to Lineage Dynamics: The authors propose a lineage inference module that leverages the learned hierarchical latent space. By performing shortest-path analysis over the learned tree-conditioned graph, HDTree recovers robust developmental trajectories even in sparse data regimes.
Empirical Superiority: Extensive evaluations on general-purpose datasets (MNIST, Fashion-MNIST, 20News-Groups, CIFAR-10) and complex single-cell datasets (Limb, LHCO, Weinreb, ECL) demonstrate that HDTree consistently outperforms state-of-the-art baselines (including TreeVAE, LadderVAE, and foundation models like Geneformer and CellPLM) in reconstruction quality, clustering accuracy, and lineage consistency.

Results and Significance
The paper claims that HDTree offers a reliable and efficient approach for modeling cellular differentiation paths.

Performance: On single-cell datasets, HDTree achieved superior tree structure quality (measured by Dendrogram Purity and Leaf Purity) and clustering accuracy compared to both traditional clustering methods and deep learning-based baselines. For instance, on the ECL dataset, it improved the average performance metric by 10.3% over the next best method.
Stability and Generalization: The unified codebook design significantly reduced variance compared to methods using branch-specific modules, particularly in deep hierarchies.
Biological Plausibility: By aligning the generative process with the Waddington landscape, the model produces biologically plausible transitions. Case studies on C. elegans and MNIST demonstrated that HDTree captures more accurate lineage relationships and generates smoother, continuous transitions compared to TreeVAE.
Efficiency: While the diffusion-based inference is inherently slower than VAE-based sampling, HDTree demonstrated competitive training efficiency, particularly on large-scale datasets where traditional methods scale poorly.

The authors conclude that HDTree enables accurate and efficient modeling of cellular differentiation paths, offering reliable insights for biological discovery. They acknowledge limitations regarding the computational cost of diffusion-based sampling for large-scale datasets and suggest future work on accelerating generation via distillation techniques. The work is positioned as a contribution to hierarchical representation learning and AI for Science, facilitating data-driven discoveries in genomics.