Dual Filter: A Transformer-like Inference Architecture for Hidden Markov Models

Resumen Técnico: Dual Filter

1. Planteamiento del Problema

El objetivo central del artículo es abordar el problema de la predicción causal no lineal en entornos donde las observaciones se generan a partir de un Modelo Oculto de Markov (HMM). Específicamente, se busca calcular la probabilidad condicional del siguiente token ($Z_{T+1}$) dada una secuencia finita de observaciones pasadas ($Z_1, \dots, Z_T$).

El problema se inspira en la arquitectura de los transformadores de solo decodificador (decoder-only transformers), que mapean secuencias de tokens a probabilidades condicionales. Sin embargo, a diferencia de los enfoques de aprendizaje profundo que aprenden estas funciones de datos, este trabajo busca derivar arquitecturas similares a los transformadores desde primeros principios, asumiendo un modelo generativo subyacente (HMM) y enfocándose exclusivamente en la inferencia, no en el aprendizaje.

El desafío consiste en encontrar una representación para el predictor no lineal que sea causal (dependa solo de observaciones pasadas) y que pueda ser computada eficientemente, estableciendo un puente teórico entre el filtrado estocástico clásico y las arquitecturas modernas de IA.

2. Metodología y Marco Teórico

Los autores desarrollan un marco matemático basado en la teoría de control óptimo y la dualidad entre el filtrado no lineal y el control estocástico.

• Formulación del Predictor No Lineal:
  Se propone una representación para la medida condicional (posterior) $\pi_T(F)$ en la forma:
  $$\pi_T(F) = (\text{constante}) - \sum_{t=1}^{T} U_{t-1}^T e(Z_t)$$
  Donde $e(Z_t)$ es una codificación de los tokens (similar al one-hot encoding pero con una restricción de suma cero para garantizar la unicidad) y $U_t$ es un proceso estocástico adaptado a las observaciones. Esta estructura es análoga a un predictor lineal (filtro de Wiener), pero con pesos aleatorios ($U_t$) que dependen de la historia de observaciones, lo que introduce la no linealidad.

• Dualidad y Control Óptimo:
  El problema de minimizar el error cuadrático medio (MMSE) para la estimación se reformula como un problema de control óptimo.

  • Se introduce una Ecuación Diferencial Estocástica hacia Atrás (BS$\Delta$E) como restricción dinámica.
  • Se define un objetivo de control que incluye un costo de ejecución y una varianza terminal.
  • Se establece un Principio de Dualidad (Teorema 9) que conecta el error de estimación con el valor óptimo del problema de control.

• Fórmula del Control Óptimo:
  Mediante la resolución del problema de control, se deriva una fórmula explícita para el control óptimo $U_t^{(opt)}$ en forma de retroalimentación:
  $$U_t^{(opt)} = \phi(Y_t, V_t; \pi_t)$$
  Donde $\phi$ es una función que depende de la medida condicional actual $\pi_t$. Esto revela que la medida condicional satisface una ecuación de punto fijo.

• El Algoritmo Dual Filter:
  Para resolver la ecuación de punto fijo de manera práctica, los autores proponen el algoritmo Dual Filter. Este algoritmo define una transformación de capa $N^{(dfltr)}$ en el espacio de medidas de probabilidad.

  • Algoritmo Iterativo: Simula la transformación de capas de manera similar a un transformador, convergiendo hacia la medida condicional exacta.
  • Algoritmo de Disparo Único (Single-shot): Una versión más eficiente que calcula la solución en una sola pasada de retroceso, con una complejidad de $O(d^2 T)$, independiente del tamaño del vocabulario $m$.

3. Contribuciones Clave

1. Nueva Representación del Predictor: Se introduce una representación causal no lineal para el predictor basada en una suma ponderada de observaciones pasadas con pesos estocásticos (Ecuación 1), generalizando el filtro de Wiener al caso no lineal.
2. Teoría de Dualidad para HMMs: Se presenta una teoría de dualidad original entre el filtrado no lineal y el control óptimo para HMMs con observaciones discretas. Esto resuelve un problema abierto en la teoría de control, extendiendo la dualidad de mínima varianza (tradicionalmente limitada a modelos lineales gaussianos) a sistemas discretos.
3. Arquitectura de Punto Fijo: Se define una transformación de capa $N^{(dfltr)}$ que actúa como un punto fijo para la medida condicional, proporcionando una base teórica para arquitecturas de inferencia.
4. Algoritmo Escalable: El algoritmo Dual Filter tiene una complejidad computacional que no depende del tamaño del vocabulario ($m$), lo que lo hace escalable para vocabularios grandes típicos en modelos de lenguaje, a diferencia de métodos que requieren $O(m^3)$.
5. Correspondencia con Transformadores: Se establece una correspondencia estructural detallada entre las capas del Dual Filter y las capas de atención de un transformador, interpretando la atención como un predictor no lineal y la dinámica de las capas como un transporte en el espacio de medidas de probabilidad.

4. Resultados Numéricos

Los autores validan el algoritmo mediante experimentos numéricos utilizando configuraciones de parámetros que imitan modelos de transformadores a escala de investigación (dimensiones $d=384$, vocabulario $m=65$, longitud de secuencia $T=256$).

• Precisión: El algoritmo de "disparo único" y el iterativo logran una precisión casi perfecta al compararse con el filtro no lineal estándar (algoritmo hacia adelante) en un HMM conocido.
• Análisis del Control Óptimo: Se visualiza el control óptimo $U_t$ para diferentes dinámicas del sistema (controladas por un parámetro de homotopía $\alpha$):
  • En dinámicas rápidas (mezcla rápida), el control es significativo solo en las observaciones más recientes.
  • En dinámicas deterministas, el control es activo al inicio para resolver la incertidumbre y luego se vuelve cero.
  • En dinámicas intermedias, el control refleja correlaciones a largo plazo, comportándose de manera similar a lo observado en predicciones basadas en transformadores.
• Comparación con NanoGPT: Se entrenó un modelo nanoGPT con datos generados por el HMM. Los resultados mostraron que el transformador puede aproximar bien el filtro no lineal en casos "favorables" (donde la distribución condicional es dispersa), pero falla en casos "desfavorables" donde la complejidad de la tarea de predicción es alta. Esto sugiere que el rendimiento del transformador depende de la estructura subyacente de los datos.

5. Significado e Impacto

Este trabajo es significativo por varias razones:

• Fundamentación Teórica de los Transformadores: Ofrece una explicación matemática rigurosa de por qué las arquitecturas tipo transformador son efectivas para la predicción causal, vinculándolas con la teoría clásica de filtrado y control estocástico, en lugar de tratarlas como "cajas negras" empíricas.
• Puente entre Control y Aprendizaje Automático: Extiende la dualidad de Kalman (mínima varianza) a contextos no lineales y discretos, un avance teórico importante en la teoría de control.
• Inferencia sin Aprendizaje: Proporciona un método de inferencia exacto y eficiente para HMMs sin necesidad de entrenamiento, lo cual es útil en escenarios donde el modelo generativo es conocido o puede ser estimado.
• Nuevas Direcciones de Investigación: Abre la puerta para desarrollar algoritmos de aprendizaje que aprendan la transformación de capas $N^{(dfltr)}$ directamente, posiblemente superando las limitaciones de los transformadores actuales en tareas de control parcial (POMDPs) y mejorando la estabilidad de los filtros.

En conclusión, el artículo propone el Dual Filter no solo como un algoritmo de inferencia eficiente, sino como un marco teórico unificador que revela la naturaleza de los transformadores como arquitecturas de transporte de medidas de probabilidad para la predicción causal.