HDTree: Generative Modeling of Cellular Hierarchies for Robust Lineage Inference

Resumen Técnico: HDTree

Enunciado del Problema
En la investigación de células individuales, rastrear y analizar con precisión las trayectorias de diferenciación de alto rendimiento es esencial para comprender los procesos biológicos. Esta tarea requiere un modelado robusto de las estructuras jerárquicas que gobiernan el desarrollo celular. Los métodos tradicionales, que a menudo combinan reducción de dimensionalidad, agrupamiento y regresión, enfrentan limitaciones significativas en cuanto a costo computacional, estabilidad del rendimiento y capacidad generativa. Si bien los enfoques basados en Autoencoders Variacionales (VAE) han mejorado la eficiencia, generalmente dependen de módulos de red especializados para cada rama del árbol. Este diseño conduce a varios problemas:

1. Escalabilidad y Estabilidad: Las ramas profundas con muestras dispersas no pueden aprovechar eficazmente el conocimiento de representación de otras ramas, lo que limita la generalización y la capacidad de preservar la estructura global.
2. Sobreajuste: Los pares independientes de codificador-decodificador en cada nodo de la rama son propensos al sobreajuste debido a la acumulación de ruido, especialmente cuando los datos de entrenamiento son limitados.
3. Colapso Posterior: Los VAE a menudo sufren de colapso posterior, y sus mecanismos de generación de un solo paso no reflejan naturalmente el refinamiento continuo y gradual de los estados celulares descrito por el paisaje de Waddington.

Metodología: HDTree
Para abordar estos desafíos, los autores proponen HDTree (Modelo de Difusión Cuantizado Jerárquico), un marco generativo unificado diseñado para la inferencia robusta de linajes. La innovación central radica en integrar la codificación del espacio latente jerárquico con un proceso de difusión cuantizado, desacoplando la complejidad del árbol del tamaño de la red.

El marco consta de tres módulos sinérgicos:

1. Codificador Unificado y Libro de Códigos de Árbol Jerárquico (HTC): A diferencia de los métodos anteriores donde los parámetros escalan exponencialmente con la profundidad del árbol, HDTree emplea un único codificador compartido ($E_{\theta_E}$) y un libro de códigos jerárquico unificado ($C_W$). El libro de códigos está estructurado como un árbol binario donde cada nodo representa un vector de código aprendible en el espacio latente. Los datos de entrada se mapean a un vector latente, que luego se cuantiza en una secuencia jerárquica de vectores de código ($s_i$) recorriendo el árbol desde la raíz hasta la hoja. Este diseño permite que todas las ramas compartan los mismos vectores del libro de códigos, permitiendo que los nodos profundos dispersos aprovechen el conocimiento de representación de otras ramas mientras preservan la estructura jerárquica global.
2. Decodificador de Difusión Cuantizado: Para modelar las transiciones continuas del estado celular y evitar el colapso posterior, HDTree utiliza un decodificador de difusión ($D_{\theta_D}$) basado en Modelos Probabilísticos de Difusión de Eliminación de Ruido (DDPM). El proceso de generación se condiciona en la secuencia de códigos jerárquicos cuantizados ($s_i$). Comenzando desde ruido gaussiano, el modelo refina iterativamente los datos a través de $T$ pasos de difusión. Esta eliminación de ruido iterativa se alinea con el paisaje de Waddington, modelando la diferenciación como un refinamiento gradual desde progenitores gruesos hasta estados especializados de alta granularidad.
3. Función de Pérdida: El modelo se optimiza utilizando una función de pérdida compuesta:
   • Pérdida de Aprendizaje Contrastivo Suave (SCL): Adapta el aprendizaje contrastivo para ponderar los pares negativos según la distancia del árbol en el espacio hiperbólico, preservando la granularidad jerárquica.
   • Pérdida de Cuantización Jerárquica (HQL): Asegura representaciones robustas con estructura de árbol al alinear las incrustaciones latentes con vectores de código de múltiples niveles y hacer cumplir la consistencia entre los nodos padres e hijos.
   • Pérdida de Difusión: Optimiza el límite inferior variacional (ELBO) para asegurar una reconstrucción fiel de los datos mientras respeta las condiciones jerárquicas.

Inferencia de Linaje
Para el análisis de trayectorias, el árbol jerárquico aprendido se transforma en un grafo ponderado. Se agregan aristas adicionales dentro de la misma profundidad utilizando un enfoque de k-vecinos más cercanos (KNN) para capturar relaciones semánticas locales. Las trayectorias de desarrollo se infieren encontrando el camino más corto entre un origen designado (por ejemplo, célula madre) y un destino (por ejemplo, célula diferenciada) en este grafo, minimizando los pesos totales de las aristas.

Contribuciones Clave

1. Innovación Metodológica: HDTree introduce un marco unificado que integra la cuantización vectorial jerárquica con procesos de difusión. Al utilizar un único libro de códigos jerárquico compartido, desacopla eficazmente la complejidad del árbol del tamaño de la red, asegurando la escalabilidad hacia jerarquías biológicas profundas sin la inestabilidad de los módulos específicos de cada rama.
2. Aplicación a la Dinámica de Linajes: Los autores proponen un módulo de inferencia de linaje que aprovecha el espacio latente jerárquico aprendido. Al realizar un análisis de camino más corto sobre el grafo condicionado al árbol aprendido, HDTree recupera trayectorias de desarrollo robustas incluso en regímenes de datos dispersos.
3. Superioridad Empírica: Evaluaciones extensas en conjuntos de datos de propósito general (MNIST, Fashion-MNIST, 20News-Groups, CIFAR-10) y conjuntos de datos complejos de células individuales (Limb, LHCO, Weinreb, ECL) demuestran que HDTree supera consistentemente a las líneas base más avanzadas (incluyendo TreeVAE, LadderVAE y modelos fundacionales como Geneformer y CellPLM) en calidad de reconstrucción, precisión de agrupamiento y consistencia de linaje.

Resultados y Significado
El artículo afirma que HDTree ofrece un enfoque confiable y eficiente para modelar las rutas de diferenciación celular.

• Rendimiento: En conjuntos de datos de células individuales, HDTree logró una calidad superior de estructura de árbol (medida por Pureza de Dendrograma y Pureza de Hojas) y precisión de agrupamiento en comparación tanto con los métodos de agrupamiento tradicionales como con las líneas base basadas en aprendizaje profundo. Por ejemplo, en el conjunto de datos ECL, mejoró la métrica de rendimiento promedio en un 10,3% sobre el siguiente mejor método.
• Estabilidad y Generalización: El diseño unificado del libro de códigos redujo significativamente la varianza en comparación con los métodos que utilizan módulos específicos de cada rama, particularmente en jerarquías profundas.
• Plausibilidad Biológica: Al alinear el proceso generativo con el paisaje de Waddington, el modelo produce transiciones biológicamente plausibles. Los estudios de caso en C. elegans y MNIST demostraron que HDTree captura relaciones de linaje más precisas y genera transiciones más suaves y continuas en comparación con TreeVAE.
• Eficiencia: Si bien la inferencia basada en difusión es intrínsecamente más lenta que el muestreo basado en VAE, HDTree demostró una eficiencia de entrenamiento competitiva, particularmente en conjuntos de datos a gran escala donde los métodos tradicionales escalan mal.

Los autores concluyen que HDTree permite un modelado preciso y eficiente de las rutas de diferenciación celular, ofreciendo perspectivas confiables para el descubrimiento biológico. Reconocen las limitaciones relacionadas con el costo computacional del muestreo basado en difusión para conjuntos de datos a gran escala y sugieren trabajo futuro sobre la aceleración de la generación mediante técnicas de destilación. El trabajo se posiciona como una contribución al aprendizaje de representaciones jerárquicas y a la IA para la Ciencia, facilitando descubrimientos impulsados por datos en genómica.