Resumen Técnico: Decodificador Generativo Riemanniano

Enunciado del Problema

Los datos del mundo real a menudo poseen estructuras intrínsecas no euclidianas (por ejemplo, árboles evolutivos, redes sociales, señales periódicas) que están mal representadas por espacios latentes euclidianos estándar ($\mathbb{R}^d$). Forzar tales datos en configuraciones euclidianas distorsiona sus relaciones geométricas, lo que conduce a visualizaciones deficientes y a la pérdida de señales.

Los enfoques existentes para el aprendizaje de representaciones riemannianas, principalmente los Autoencoders Variacionales (VAE), intentan incrustar datos en variedades imponiendo priores específicos (por ejemplo, von Mises–Fisher en esferas o normales riemannianas en espacios hiperbólicos). Sin embargo, estos métodos enfrentan desafíos significativos:

1. Fragilidad Numérica: Requieren optimizar objetivos complejos que involucran constantes de normalización intratables y estimaciones de Monte Carlo para densidades en variedades arbitrarias.
2. Escalabilidad: Estas aproximaciones de densidad a menudo limitan la escalabilidad a variedades generales o espacios latentes de alta dimensión.
3. Complejidad: Imponer restricciones de variedad generalmente depende de un codificador para estimar densidades, añadiendo sobrecarga arquitectónica y computacional.

Metodología

Los autores proponen el Decodificador Generativo Riemanniano (RGD), un marco unificador que descarta el codificador por completo. En lugar de amortizar la inferencia, el modelo trata los códigos latentes como parámetros libres para ser optimizados directamente en una variedad riemanniana elegida.

Formulación Central

El problema de aprendizaje se formula como una estimación de Máxima A Posteriori (MAP). Dado un conjunto de datos $X = \{x_i\}_{i=1}^N$ y un decodificador diferenciable $f_\theta: \mathcal{M} \to \mathcal{X}$, el modelo optimiza conjuntamente los parámetros del decodificador $\theta$ y los códigos latentes $Z = \{z_i\}_{i=1}^N$ minimizando el log-posterior negativo:
$$\mathcal{L}(\theta, Z) = \sum_{i=1}^N \left( -\log p_\theta(x_i|z_i) - \log p(z_i) \right) - \log p(\theta)$$

• Verosimilitud: Modelada mediante pérdida de reconstrucción (por ejemplo, error cuadrático para ruido gaussiano).
• Prior: Para variedades compactas, se utiliza un prior uniforme; para variedades no compactas, se emplean distribuciones envueltas o normales riemannianas.
• Optimización: Los códigos latentes se actualizan utilizando descenso de gradiente riemanniano (específicamente Adam riemanniano), asegurando que las iteraciones permanezcan en la variedad $\mathcal{M}$ mediante mapas de retracción. Los parámetros del decodificador se actualizan mediante Adam euclidiano estándar.

Regularización Geométrica

Para alinear la suavidad del decodificador con la geometría de la variedad sin requerir estimación explícita de densidades, los autores introducen una regularización consciente de la geometría mediante ruido de entrada.

• Durante el entrenamiento, cada latente $z$ se perturba con ruido gaussiano $\epsilon \sim \mathcal{N}(0, \sigma^2 G^{-1}(z))$, donde $G(z)$ es el tensor métrico riemanniano.
• Este ruido se inyecta mediante el mapa exponencial (aproximado como identidad para ruido pequeño).
• Derivación Teórica: Basado en Bishop (1995) y An (1996), la pérdida esperada bajo esta perturbación de ruido produce un regularizador aditivo:
  $$\mathbb{E}_\epsilon[\mathcal{L}(z')] \approx \mathcal{L}(z) + \sigma^2 \text{Tr}\left( J(z)^\top G^{-1}(z) J(z) \right)$$
  donde $J(z)$ es el jacobiano del decodificador. Este término penaliza grandes gradientes de salida ponderados por la métrica inversa, fomentando efectivamente que el decodificador respete la curvatura local de la variedad.

Flexibilidad

El marco soporta cualquier variedad riemanniana implementada en la biblioteca geoopt, incluyendo esferas, bolas de Poincaré, modelos de Lorentz y ProductManifolds (que permiten geometrías latentes heterogéneas donde diferentes dimensiones del espacio latente residen en variedades diferentes).

Contribuciones Clave

1. Marco Unificador: Un método para el aprendizaje de representaciones en cualquier variedad riemanniana combinando optimización riemanniana con un modelo generativo sin codificador, eliminando la necesidad de aproximaciones variacionales de densidad.
2. Regularización Geométrica Escalable: Una técnica de regularización basada en ruido que promueve la coherencia entre la función del decodificador y la métrica de la variedad, evitando penalizaciones de curvatura de segundo orden computacionalmente costosas.
3. Validación Empírica: Demostración del método en diversos conjuntos de datos, mostrando una mejor alineación con la geometría intrínseca y estabilidad en dimensiones altas donde las líneas base variacionales fallan.

Resultados y Evaluación

El método fue validado en tres estudios de caso:

1. Proceso de Difusión Ramificado Sintético:

   • Tarea: Recuperar datos con estructura de árbol.
   • Hallazgos: Las proyecciones euclidianas y UMAP fallaron en revelar la topología del árbol. Las incrustaciones hiperbólicas regularizadas (disco de Poincaré) recuperaron con éxito la jerarquía.
   • Ablación: Aumentar la escala de ruido $\sigma$ mejoró la correlación entre las distancias geodésicas latentes y las distancias en el espacio de datos hasta un punto óptimo ($\sigma \approx 0.9$), después del cual el ruido saturó la señal.

2. ADN Mitocondrial Humano (hmtDNA):

   • Tarea: Inferir patrones de migración y linaje a partir de secuencias genéticas.
   • Hallazgos: Los modelos hiperbólicos (Lorentz/Poincaré) recuperaron jerarquías de haplogrupos y correlaciones geográficas significativamente mejor que las líneas base euclidianas o UMAP.
   • Utilidad Posterior: En tareas de clasificación (predecir regiones geográficas y haplogrupos), el RGD hiperbólico superó a las líneas base euclidianas y P-VAE, particularmente en configuraciones de alta dimensión.

3. Fases del Ciclo Celular (scRNA-seq):

   • Tarea: Modelar la progresión cíclica de la división celular.
   • Hallazgos: Las incrustaciones esféricas ($S^2$) y toroidales ($S^1 \times S^1$) capturaron la naturaleza cíclica de los datos. Mientras que los espacios euclidianos ofrecieron una fidelidad de reconstrucción ligeramente superior (debido a mayores grados de libertad), los modelos riemannianos mostraron una correlación superior con la geometría de la fase biológica.
   • Escalabilidad: El RGD demostró una escalabilidad estable a dimensiones latentes altas (hasta 500D), mientras que las líneas base variacionales (P-VAE, $\Delta$VAE) se volvieron numéricamente inestables o inviables.

Significado y Afirmaciones

El artículo afirma que el Decodificador Generativo Riemanniano ofrece una alternativa más simple y estable a los métodos existentes de aprendizaje de variedades. Al eliminar el codificador, el método evita las funciones objetivo "numéricamente frágiles" asociadas con la estimación de densidades en espacios curvos.

• Interpretabilidad: Los espacios latentes aprendidos están explícitamente alineados con la geometría prescrita, facilitando la exploración guiada por hipótesis en dominios científicos.
• Generalidad: El enfoque no se limita a variedades específicas (como esferas o espacios hiperbólicos) sino que se extiende a productos de variedades heterogéneas.
• Practicidad: El método requiere únicamente un decodificador y es compatible con arquitecturas de redes neuronales existentes.

Los autores reconocen limitaciones, señalando que la parametrización por muestra (almacenar cada latente explícitamente) conduce a un crecimiento lineal de la memoria con el tamaño del conjunto de datos, lo cual puede ser prohibitivo para millones de puntos. Sugieren esquemas híbridos (por ejemplo, inferencia amortizada) como una dirección futura. El trabajo se posiciona como una herramienta para el análisis exploratorio y la prueba de hipótesis en campos científicos, particularmente donde la topología de los datos es conocida o hipotetizada.