Resumen Técnico: Inferencia Discreta Amortizada Dependiente de la Trayectoria

Planteamiento del Problema

El artículo aborda el desafío de muestrear objetos composicionales y discretos (por ejemplo, grafos, secuencias, árboles filogenéticos) a partir de una distribución posterior no normalizada dada. Si bien los avances recientes en Redes de Flujo Generativo (GFlowNets) han enmarcado esto como el aprendizaje de un Proceso de Decisión de Markov (MDP) que construye objetos paso a paso, los autores identifican limitaciones críticas en el supuesto Markoviano estándar. Específicamente, la propiedad de Markov —donde la política depende únicamente del estado actual— puede obstaculizar la propagación de la señal durante el entrenamiento y conducir al aliasing de estado (state aliasing). En el aliasing de estado, estados distintos se vuelven representacionalmente indistinguibles para el agente, lo que provoca que el muestreador aprendido falle al distinguir acciones óptimas para diferentes trayectorias que convergen al mismo estado. Este problema es particularmente agudo en entornos con recompensas dispersas o dependencias estructurales complejas, donde las redes neuronales profundas estándar pueden carecer de la expresividad para representar la política óptima.

Metodología

Para superar estas limitaciones, los autores proponen la Inferencia Discreta Amortizada Dependiente de la Trayectoria. La metodología central consiste en "elevar" (lifting) el espacio de estados original del MDP mediante la introducción de un sistema dinámico latente condicionable al estado y aprendible $\{W_t\}_{t \ge 1}$.

1. DAG Apuntado Elevado

En lugar de un grafo de estados estándar $G = (S, X, E)$, el método opera sobre un DAG apuntado elevado $G^\star$. El espacio de estados se extiende de $S$ a $S^\star \subseteq S \times \Omega$, donde $\Omega$ es un espacio latente (típicamente $\mathbb{R}^d$).

• Dinámica Latente: La variable latente $W_t$ evoluciona de forma determinista basada en el estado actual $s_t$ y el estado latente previo $W_{t-1}$ a través de un campo vectorial aprendible $\phi$:
  $$W_t = W_{t-1} + \phi(s_t, W_{t-1})$$
• Dependencia de la Política: La política hacia adelante $p_F$ se vuelve dependiente de la trayectoria, condicionada al par $(s_t, W_t)$. Esto permite que la política distinga entre diferentes trayectorias que conducen al mismo estado $s_t$, rompiendo efectivamente la restricción de Markov del espacio de estados original.

2. Arquitectura: Matriz de Pesos Autorreferencial (SRWM)

Para implementar el sistema dinámico latente de manera eficiente, los autores utilizan una Matriz de Pesos Autorreferencial (SRWM) modificada.

• La regla de actualización para el estado latente (representado como una matriz de pesos $W_t$) sigue una recurrencia específica:
  $$q_t, k_t, v_t, \beta_t = W \psi(s_t) + b$$
  $$\bar{v}_t = W_{t-1} \zeta(k_t)$$
  $$W_t = W_{t-1} R + \sigma(\beta_t) \cdot (v_t - \bar{v}_t) \otimes \zeta(q_t)$$
• Mecanismo Clave: La inclusión de una matriz de rotación $R$ fija y ergódica es crucial. Esta rota los autovectores de la matriz de pesos a lo largo de la trayectoria de despliegue (rollout). Este mecanismo está diseñado para forzar la diversidad en las representaciones ocultas de la política, mitigando así el aliasing de estado de manera más efectiva que las arquitecturas recurrentes estándar como LSTM o GRU.

3. Objetivos de Aprendizaje

Los autores demuestran que los algoritmos de aprendizaje existentes para GFlowNets pueden extenderse a este entorno elevado. Demuestran que las condiciones de Balance de la Trayectoria (TB), Balance de la Subtrayectoria (SubTB) y Balance Contrastivo (CB) siguen siendo válidas para el DAG elevado.

• El proceso de aprendizaje minimiza el valor esperado del logaritmo del cuadrado de la diferencia entre los lados izquierdo y derecho de estas ecuaciones de balance.
• Crucialmente, el campo vectorial $\phi$ y la red de la política se aprenden conjuntamente mediante el descenso de gradiente estocástico sobre una única función de pérdida.
• Los resultados teóricos (Proposición 4.8) indican que si la política hacia atrás $p_B$ es Markoviana (independiente de $W_t$), la política hacia adelante $p_F$ óptima también será Markoviana con respecto al estado no elevado. Sin embargo, cuando $p_B$ es dependiente de la trayectoria, $p_F$ puede aprovechar la historia completa para lograr una mejor expresividad.

Contribuciones Clave

1. Algoritmo Dependiente de la Trayectoria: La introducción de un algoritmo de propósito general para la inferencia amortizada en espacios discretos y composicionales, abordando explícitamente los límites de expresividad de los muestreadores Markovianos.
2. Pruebas de Expresividad Teórica:
   • Políticas Lineales: Los autores demuestran (Proposición 3.1) que para políticas lineales, el espacio de distribuciones aprendibles en un entorno dependiente de la trayectoria es estrictamente mayor que en un entorno Markoviano.
   • Políticas GNN: Extienden esto a las Redes Neuronales de Grafos (Proposición 3.2), mostrando que las parametrizaciones dependientes de la trayectoria pueden distinguir estados que las GNN 1-WL no pueden (por ejemplo, distinguir aristas añadidas en diferentes momentos), mejorando así estrictamente la expresividad.
3. Extensión de la Teoría de Aprendizaje: El artículo proporciona una prueba formal (Proposición 4.5 y Corolarios 4.6, 4.7) de que los objetivos de aprendizaje estándar de GFlowNet (TB, SubTB, CB) son suficientes para asegurar un muestreo correcto en el entorno elevado y dependiente de la trayectoria.
4. Validación Empírica: Experimentos extensos en bancos de pruebas estándar (Generación de Conjuntos, Diseño de Secuencias, Grid World, Modelos de Ising, etc.) demuestran que el enfoque dependiente de la trayectoria logra una convergencia más rápida y una mejor bondad de ajuste (medida por la distancia de Variación Total y el Puntaje de Consistencia de Flujo) en comparación con modelos Markovianos con parámetros similares.

Resultados

El análisis empírico destaca varios hallazgos consistentes:

• Expresividad: En el entorno "LINES" y en tareas estructuradas de grafos, los muestreadores Markovianos fallaron al aprender distribuciones no uniformes o al distinguir entre estados estructuralmente similares, mientras que los modelos dependientes de la trayectoria lograron ajustes casi perfectos.
• Convergencia: Los modelos dependientes de la trayectoria convergieron consistentemente más rápido. Por ejemplo, en el entorno disperso de LINES, una red ReLU Markoviana tardó aproximadamente 100 veces más iteraciones en separar las políticas para estados distintos en comparación con su contraparte dependiente de la trayectoria.
• Eficiencia de la Arquitectura: La SRWM con una matriz de rotación superó a las RNN con compuertas estándar (LSTM, GRU) en términos de velocidad de convergencia y rendimiento final, debido a su capacidad para mitigar el aliasing de estado mediante el mecanismo de rotación.
• Robustez: El método mostró un rendimiento superior en diversas tareas, incluyendo el diseño de afinidad de unión al ADN (secuencias biológicas), el aprendizaje de preferencias y las simulaciones del modelo de Ising.

Significado y Reivindicaciones

El artículo sostiene que elevar el espacio de estados con un sistema dinámico latente es una estrategia necesaria para el muestreo amortizado discreto, de forma análoga a cómo el método de Monte Carlo de Hamiltonianos (HMC) aumenta los espacios continuos con variables de momento para acelerar la convergencia. Los autores argumentan que:

• El Aliasing de Estado es un Cuello de Botella Fundamental: La incapacidad de los muestreadores Markovianos para distinguir estados basados en la historia de la trayectoria es una causa primaria del bajo rendimiento en espacios composicionales complejos.
• La Dependencia de la Trayectoria es Probablemente Superior: Los resultados teóricos confirman que las parametrizaciones dependientes de la trayectoria expanden estrictamente la clase de hipótesis de distribuciones aprendibles.
• Aplicabilidad Práctica: El método propuesto es un reemplazo directo para el entrenamiento estándar de GFlowNets, requiriendo solo la modificación de la representación del estado y el uso de objetivos de aprendizaje específicos que siguen siendo válidos en el espacio elevado.

Los autores concluyen que los algoritmos diseñados para procesos de decisión de Markov parcialmente observables (POMDP) son muy adecuados para el muestreo discreto amortizado. Observan que, aunque se centraron en la dinámica latente determinista, la viabilidad de la dinámica estocástica sigue siendo una pregunta abierta para investigaciones futuras. El trabajo no pretende resolver todos los problemas de muestreo discreto, sino que ofrece un marco de propósito general para mitigar las limitaciones específicas de los enfoques Markovianos actuales.