Resumen Técnico: Inversión de Transformaciones de Datos mediante Muestreo por Difusión

Planteamiento del Problema

El artículo aborda el problema de la inversión de transformaciones en grupos de Lie generales. En muchos escenarios de aprendizaje automático y modelado científico, los datos observados $\tilde{x}$ son generados aplicando una transformación desconocida $g$ de un grupo de Lie $G$ a una muestra de la distribución original $x$ (es decir, $\tilde{x} = g \cdot x$). El objetivo es recuperar una transformación inversa $g^{-1}$ que mapee la observación fuera de la distribución de vuelta a la distribución de datos original.

Esto se formula como un problema inverso ciego porque la transformación de molestia es desconocida y la solución es generalmente no única. Mientras que los métodos existentes suelen operar en el espacio de los datos o están especializados en familias de transformaciones específicas (por ejemplo, grupos afines o de rotación), este trabajo se dirige a acciones de grupos de Lie generales, incluyendo grupos no compactos y no abelianos sin métricas bi-invariantes (por ejemplo, homografías de imagen $PGL(3, \mathbb{R})$ y simetrías de puntos de Lie para EDP).

Una aplicación clave destacada es la equivarianza en tiempo de prueba (test-time equivariance). Las redes neuronales preentrenadas suelen carecer de robustez ante transformaciones geométricas. Los autores proponen invertir estas transformaciones en el tiempo de prueba para devolver las entradas a la distribución de entrenamiento del modelo, mejorando así la robustez sin requerir arquitecturas equivariantes especializadas o reentrenamiento.

Metodología: Difusión de Energía para la Inversión de Transformaciones (TIED)

El método propuesto, TIED, modela la distribución posterior sobre las transformaciones $p(g | \tilde{x})$ como una distribución de Boltzmann definida por una función de energía $E$ en el espacio de los datos. La energía se deriva de una distribución previa de datos $p_x(x)$, donde $E(x) = -\log p_x(x)$. La densidad de la posterior viene dada por:
$$p(g | \tilde{x}) \propto p_x(g^{-1} \cdot \tilde{x}) |\det J_{g^{-1}}(\tilde{x})|$$

Para muestrear de esta posterior, TIED introduce un proceso de difusión directamente en el grupo de Lie. La metodología supera dos desafíos principales: la naturaleza accidentada y multimodal de los paisajes de energía (a menudo derivados de redes neuronales) y las restricciones geométricas de muestreo en variedades.

1. Difusión Trivializada en Grupos de Lie

En lugar de realizar actualizaciones en los espacios tangentes de elementos de grupo arbitrarios, TIED utiliza la trivialización. Las actualizaciones se realizan en la álgebra de Lie $\mathfrak{g}$ (el espacio tangente en la identidad) y se transportan al elemento de grupo actual mediante el mapa exponencial.

• Proceso hacia adelante (Forward Process): Se define una SDE de explosión de varianza en el grupo mediante el dibujo de perturbaciones infinitesimales en el álgebra de Lie y su transporte vía multiplicación por la derecha:
  $$dg_t = d(L_{g_t})_e [\gamma(t) dw^{\mathfrak{g}}_t]$$
  donde $w^{\mathfrak{g}}_t$ es el movimiento browniano en el álgebra de Lie.
• Proceso inverso (Reverse Process): La SDE de tiempo inverso se deriva para recuperar muestras de la posterior. Crucialmente, el término de deriva involucra el score trivializado $\nabla_g \log p_t(g_t)$, que es un elemento del álgebra de Lie.

2. Identidad de Score Objetivo Trivializado

Una contribución teórica central es una nueva identidad de score objetivo trivializado. El ajuste de score (score matching) estándar requiere acceso a muestras limpias, las cuales no están disponibles en este escenario de inversión. En su lugar, TIED expresa el score de la distribución ruidosa en términos del gradiente de la función de energía limpia.
Los autores demuestran que para grupos de Lie generales (sin requerir una métrica bi-invariante), el score trivializado puede estimarse mediante una expectativa de Monte Carlo sobre el núcleo de ruido $k_t$:
$$\nabla_g \log p_t(g_t) = \nabla_g \log \int_G k_t(w) \exp\left(-E(g_t w^{-1}) - \log \lambda(w)\right) d\mu(w)$$
donde $\lambda(w)$ contabiliza el cambio de la medida de Haar bajo la inversión. Esto permite que el score sea estimado utilizando únicamente la función de energía $E$ y sus gradientes, computados mediante diferenciación automática.

3. Implementación Práctica

El algoritmo de muestreo (Algoritmo 1) discretiza la SDE inversa. En cada paso, este:

1. Muestrea ruido $w$ del núcleo de ruido $k_t$.
2. Computa la energía $E$ para candidatos transformados.
3. Estima el score trivializado usando la identidad de Monte Carlo (empleando el truco log-sum-exp para estabilidad).
4. Actualiza el elemento del grupo usando el mapa exponencial.

El método soporta varias definiciones de energía, incluyendo aquellas derivadas de Autoencoders Variacionales (VAEs) o directamente de la confianza de predicción de un clasificador preentrenado.

Contribuciones Clave

1. Muestreador TIED: Un nuevo muestreador de difusión para invertir transformaciones de datos en grupos de Lie generales. Garantiza que las actualizaciones permanezcan en la variedad y maneja acciones no compactas y no lineales utilizando solo el conocimiento del álgebra de Lie.
2. Marco Teórico: Una demostración de que la inversión de transformaciones puede utilizarse para la equivariancia en tiempo de prueba de modelos preentrenados. Los autores demuestran que la posterior de la inversión de la transformación es $G$-equivariante, asegurando la solidez del enfoque.
3. Nueva Identidad de Score: Una generalización de la identidad de score objetivo a grupos de Lie que funciona sin métricas bi-invariantes, permitiendo un muestreo basado en score eficiente utilizando gradientes de energía en lugar de muestras de datos limpios.
4. Validación Empírica: Demostraciones sobre homografías de imagen ($PGL(3, \mathbb{R})$) y simetrías de puntos de EDP ($SL(2, \mathbb{R}) \ltimes H(1, \mathbb{R})$), mostrando mejoras de rendimiento consistentes sobre bases sólidas.

Resultados Experimentales

El artículo evalúa TIED en tres tareas principales:

1. Muestreo Sintético en $SO(10)$: TIED muestrea con éxito de una distribución de energía multimodal en un grupo de Lie de alta dimensión ($dim=45$), convergiendo significativamente más rápido que un muestreador de Langevin cinético trivializado.
2. Clasificación de Imágenes (MNIST):
   • Configuración: Un ResNet18 preentrenado es probado en imágenes MNIST sometidas a transformaciones afines ($Aff(2, \mathbb{R})$) y de homografía ($PGL(3, \mathbb{R})$) desconocidas.
   • Bases (Baselines): Comparado contra métodos basados en optimización (LieLAC, FoCal, ITS) y métodos de muestreo (Langevin cinético).
   • Resultados: TIED supera significativamente a todas las bases. Por ejemplo, bajo transformaciones afines con energía basada en el clasificador, TIED logra una precisión del 82.64% frente al 55.48% de la ResNet18 pura y el 74.72% de Langevin cinético. También supera a redes equivariantes especializadas (affConv) en algunos entornos.
3. Resolución de EDP:
   • Configuración: Un operador neuronal DeepONet es probado en ecuaciones de Calor y Burgers 1D bajo transformaciones de simetría de puntos de Lie desconocidas.
   • Resultados: TIED logra el error relativo $L_2$ más bajo en todos los casos de prueba, superando a LieLAC y FoCal. Para la ecuación de Calor 1D con simetrías $SL(2, \mathbb{R}) \ltimes H(1, \mathbb{R})$, TIED reduce el error a 0.043, comparado con 0.066 para LieLAC y 0.690 para el DeepONet sin corregir.

Significado y Reivindicaciones

Los autores afirman que TIED proporciona un pipeline libre de entrenamiento (training-free) para hacer que cualquier modelo preentrenado sea equivariante a transformaciones de grupos de Lie en el tiempo de prueba. Al aprovechar los modelos basados en energía (que pueden aproximarse por el propio modelo preentrenado), el método evita la necesidad de arquitecturas equivariantes especializadas o reentrenamientos costosos.

La importancia radica en:

• Generalidad: Se aplica a una amplia clase de grupos de Lie, incluyendo aquellos que son no compactos y carecen de métricas bi-invariantes, los cuales han sido un desafío para los métodos de difusión o optimización previos.
• Robustez: Restaura eficazmente las entradas transformadas a la distribución de entrenamiento, mejorando significativamente la generalización de las redes preentrenadas ante transformaciones geométricas fuera de la distribución.
• Eficiencia: El enfoque de difusión, combinado con la identidad de score trivializado, ofrece un proceso de muestreo más eficiente y estable en comparación con el MCMC tradicional o la optimización iterativa en paisajes de energía accidentados.

El artículo concluye que este enfoque representa una forma de escalado de la equivariancia en el tiempo de inferencia, ofreciendo una solución práctica para la robustez en la era de los grandes modelos fundacionales donde el manejo de transformaciones fuera de la distribución es crítico.