Resumen Técnico: Aprendizaje en la Variedad con Ajuste de Flujo Riemanniano

1. Planteamiento del Problema

Los avances recientes en el modelado generativo se han desplazado hacia el aprovechamiento de codificadores de representación semántica de alta dimensión y congelados (p. ej., DiNOv2, SigLIP) para mejorar la fidelidad de la síntesis y la velocidad de convergencia. Sin embargo, los marcos de trabajo estándar de Diffusion Transformers (DiT) y Flow Matching (FM) no logran converger eficazmente cuando se entrenan directamente sobre estas representaciones.

Trabajos previos, específicamente los de Representation Autoencoders (RAE), atribuyeron este fallo a un cuello de botella de capacidad, proponiendo que los modelos requieren un "escalado de ancho" (aumentar el ancho del transformer para igualar la dimensionalidad latente) para resolver el problema. Este artículo desafía esa hipótesis, argumentando que la causa raíz no es una capacidad de modelo insuficiente, sino un desajuste geométrico fundamental, denominado Interferencia Geométrica.

El problema central surge porque los codificadores de representación proyectan las características sobre una variedad hiperesférica ($S^{d-1}$) con una varianza radial cercana a cero (una geometría de "caparazón duro"). En contraste, el Flow Matching euclídeo estándar asume una trayectoria de probabilidad lineal (una cuerda) entre los datos y el ruido Gaussiano. Esta interpolación lineal corta a través del interior de baja densidad de la hiperesfera, obligando al modelo a aprender campos de velocidad en regiones indefinidas del espacio de características. En consecuencia, el modelo desperdicia su capacidad minimizando errores radiales (que son estructuralmente cero en la variedad) en lugar de aprender las transiciones angulares semánticas.

2. Metodología: Ajuste de Flujo Riemanniano con Regularización de Jacobi (RJF)

Para resolver la Interferencia Geométrica sin modificaciones arquitectónicas como el escalado de ancho, los autores proponen RJF, un marco que alinea el proceso generativo con la geometría intrínseca del espacio latente.

2.1 Ajuste de Flujo Riemanniano (Corrección de Trayectoria)

En lugar de la interpolación lineal en el espacio euclídeo, RJF define la trayectoria de probabilidad directamente en la variedad hiperesférica utilizando Interpolación Lineal Esférica (SLERP).

• Trayectorias Geodésicas: La trayectoria entre un punto de datos $x$ y el ruido $\epsilon$ (ambos proyectados sobre la esfera unidad) sigue la geodésica (arco de círculo máximo) en lugar de una cuerda.
• Velocidad del Espacio Tangente: El campo de velocidad objetivo se calcula diferenciando la trayectoria SLERP. Esto asegura que el vector de velocidad predicho resida estrictamente en el espacio tangente de la variedad ($v_t \cdot x_t = 0$), eliminando la necesidad de aprender componentes radiales.
• Muestreo: Durante la inferencia, el modelo utiliza un integrador de Geodésica (Mapa Exponencial) para actualizar el estado latente, asegurando que la trayectoria permanezca en la superficie de la variedad durante todo el proceso de generación.

2.2 Regularización de Jacobi (Corrección de Curvatura)

Los autores observan que corregir simplemente la trayectoria es insuficiente porque el Error Cuadrático Medio (MSE) estándar trata los errores de velocidad de manera uniforme. En una hiperesfera de curvatura positiva, las geodésicas convergen (se enfocan), lo que significa que los errores de velocidad se propagan de forma no lineal.

• Campos de Jacobi: Para tener esto en cuenta, los autores introducen un factor de ponderación derivado de los campos de Jacobi, que cuantifican la separación entre geodésicas causada por las perturbaciones de velocidad.
• Reponderación de la Pérdida: La función de pérdida se multiplica por un factor de ponderación geométrica $\lambda(t, \Omega) = \text{sinc}^2((1-t)\Omega)$, donde $\Omega$ es la distancia geodésica. Esto reduce el peso de los errores cerca de los datos ($t=0$), donde el enfoque mitiga las perturbaciones, y prioriza la precisión cerca del ruido ($t=1$), aplicando un efecto de recocido (annealing) a la señal de aprendizaje para que coincida con la curvatura de la variedad.

3. Contribuciones Clave

1. Análisis Geométrico del Fallo de Convergencia: El artículo identifica la Interferencia Geométrica como el cuello de botella principal para la difusión estándar en codificadores de representación. Demuestra que el fallo proviene de objetivos euclídeos que obligan a los modelos a aprender a través del interior fuera de la variedad de la hiperesfera, en lugar de una falta de capacidad del modelo.
2. Marco RJF: La propuesta de un marco de entrenamiento geométricamente consistente que combina el Ajuste de Flujo Riemanniano (para el alineamiento de la trayectoria) y la Regularización de Jacobi (para la propagación de error consciente de la curvatura).
3. Eficiencia sin Escalado de Ancho: La demostración de que las arquitecturas DiT estándar (p. ej., DiT-B con 131M parámetros) pueden converger eficazmente en latentes semánticos de alta dimensión sin el costoso escalado de ancho propuesto por trabajos previos.

4. Resultados Experimentales

Los experimentos se realizaron en ImageNet-1K a una resolución de 256×256 utilizando varios escales de modelo (DiT-B, DiT-L, DiT-XL) y codificadores de representación (DiNOv2-B, SigLIP, MAE).

• Convergencia y Rendimiento de FID:
  • DiT-B (131M): El Ajuste de Flujo Euclídeo estándar sobre características de DiNOv2-B falló en converger (FID ~24.21). RJF logró un FID de 6.77 en 80 épocas y 3.37 (con guía) en 200 épocas.
  • DiT-XL (677M): RJF logró un FID de 3.62 en 80 épocas sin guía, superando la línea base euclídea (FID 4.28) y al LightingDiT basado en VAE (FID 4.29).
  • Escalado: El método mostró mejoras consistentes en todos los tamaños de modelo, logrando puntuaciones de Inception (IS 186.2) y Precisión (0.82) de vanguardia.
• Estudios de Ablación:
  • Radial vs. Angular: El enmascaramiento de la pérdida radial proporcionó solo ganancias marginales (FID 21.99), confirmando que la trayectoria lineal (la cuerda) es el problema principal.
  • Trayectoria: Cambiar únicamente al Ajuste de Flujo Riemanniano (SLERP) mejoró drásticamente el rendimiento (FID 7.06), validando la hipótesis del desajuste de trayectoria.
  • Regularización: La adición de la Regularización de Jacobi refinó aún más los resultados (FID 6.77).
• Generalización: El método convergió con éxito en otros codificadores con restricciones hiperesféricas, incluyendo SigLIP (FID 10.39) y MAE (FID 19.82), donde los métodos estándar fallaron por completo.

5. Significado y Reivindicaciones

El artículo afirma que la dificultad de entrenar modelos de difusión en codificadores de representación es fundamentalmente geométrica, no arquitectónica. Al respetar la topología hiperesférica del espacio latente, los Diffusion Transformers estándar pueden aprender semánticas de alta fidelidad de manera eficiente.

Los autores aseveran que RJF desbloquea el potencial de las arquitecturas DiT estándar, permitiéndoles igualar o exceder el rendimiento de modelos mucho más grandes con escalado de ancho o de complejos flujos basados en VAE. El trabajo establece que el modelado generativo eficiente en variedades de representación requiere alineamiento geométrico (trayectorias correctas y pérdidas conscientes de la curvatura) en lugar de un escalado de capacidad por fuerza bruta. Este enfoque permite la síntesis de alta fidelidad con menos parámetros y una convergencia más rápida, como lo demuestra el modelo DiT-B alcanzando un FID de 3.37 donde métodos previos no lograron converger.