Resumen Técnico: Mitigación del Sesgo en la Decodificación Localmente Restringida mediante Propuestas Tratables

Declaración del Problema

Los Modelos de Lenguaje de Gran Escala (LLMs) a menudo fallan al generar salidas que se adhieren estrictamente a restricciones lógicas o estructurales, tales como esquemas JSON, consultas SQL válidas o formatos específicos de llamada a funciones. Si bien el ajuste de instrucciones y la optimización de preferencias han mejorado el alineamiento general, siguen siendo poco fiables para imponer restricciones "duras".

Los enfoques existentes enfrentan dos limitaciones primarias:

1. Muestreo de Rechazo (Rejection Sampling): Aunque es insesgado, es computacionalmente impráctico para restricciones complejas debido a tasas de aceptación ínfimas.
2. Decodificación Localmente Restringida (LCD): Este enfoque estándar enmascara tokens que violan inmediatamente una restricción en cada paso. Sin embargo, sufre de dos fallas críticas:
   • Sesgo: Distorsiona la distribución subyacente del modelo al enmascarar tokens de forma miope sin considerar la validez futura, lo que conduce a un sesgo de muestreo.
   • Sin Garantía de Terminación: La LCD solo asegura la validez local. Una secuencia puede permanecer localmente válida indefinidamente (por ejemplo, una llave sin cerrar en un JSON) pero fallar en satisfacer la restricción dentro de un presupuesto fijo de $n$ tokens.

Para abordar estos problemas, se ha propuesto el marco de la Simulación de Monte Carlo Secuencial (SMC), utilizando la LCD como una distribución de propuesta y reponderando partículas para converger a la distribución objetivo. La tasa de convergencia de SMC está limitada por la calidad de la propuesta; la LCD es una propuesta débil porque carece de información probabilística sobre completaciones válidas futuras y no garantiza la satisfacción de la restricencia dentro de un presupuesto finito.

Metodología

Los autores proponen un enfoque genérico para construir propuestas tratables y funciones de potencial para el muestreo SMC de la distribución restringida $p_{lm}(\cdot \mid C)$. La innovación central consiste en representar las restricciones como autómatas finitos y aprovechar las operaciones de tensores para una ejecución eficiente en GPU.

1. Decodificación Globalmente Restringida (GCD)

La GCD mejora a la LCD al asegurar que cada token generado pueda conducir a una completación válida dentro del presupuesto de tokens restante $n$.

• Autómatas Tensorizados: Las restricciones codificadas como Autómatas Finitos (FA), específicamente Autómatas Finitos No Deterministas (NFA) y Autómatas Finitos Deterministas (DFA), se compilan en representaciones tensorizadas (matrices de transición $T_{out}, T_{in}$ y matriz de emisión $B_f$). Esto permite una ejecución eficiente en GPUs.
• Inferencia Forward-Backward:
  • Proceso Forward: Computa el conjunto de estados del FA alcanzables tras observar el prefijo $x_{1:t-1}$.
  • Proceso Backward: Computa si un estado puede alcanzar un estado de aceptación en exactamente $n-t$ pasos.
• Enmascaramiento Global: Solo se permite un token si este conduce a un estado que admite una continuación válida dentro del presupuesto restante. Esto garantiza que cualquier secuencia de longitud $n$ generada bajo la GCD satisfaga la restricción $C$.

2. Decodificación Probabilística Globalmente Restringida (P-GCD)

Si bien la GCD garantiza la satisfacción de la restricción, su máscara es binaria (válido/inválido) e ignora la verosimilitud de las completaciones válidas bajo el LLM. La P-GCD aborda esto integrando información probabilística.

• Destilación de HMM: La distribución del LLM es aproximada por un Modelo Oculto de Markov (HMM), siguiendo el trabajo previo (Zhang et al., 2023).
• Multiplicación de Circuitos: El FA tensorizado (visto como un HMM no normalizado sobre aristas) se multiplica con el HMM destilado utilizando algoritmos de multiplicación de circuitos. Esto aprovecha la estructura algebraica compartida de los FA y los HMM.
• Distribución Producto: La distribución producto resultante $p_{prod}$ captura tanto las restricciones lógicas duras del FA como las preferencias probabilísticas del LLM.
• Propuesta y Potencial:
  • Propuesta ($q_{pgcd}$): Repondera la distribución del LM en cada paso utilizando la probabilidad del siguiente token restringida derivada de $p_{prod}$.
  • Potencial ($\phi_{pgcd}$): Pondera las partículas en SMC mediante el producto de la probabilidad del prefijo del LM y la estimación de mirada hacia adelante (look-ahead) del HMM sobre la satisfacción de la restricción.

3. Implementación Eficiente

Los autores utilizan operaciones de redes tensoriales (específicamente la identidad Kronecker-vec) para evitar la materialización explícita de grandes matrices de producto de Kronecker. Esto reduce el costo computacional de los pasos de transición de $O((HE)^2)$ a $O(H^2E + HE^2)$, donde $H$ es el tamaño oculto del HMM y $E$ es el recuento de aristas del FA, permitiendo una generación autorregresiva eficiente.

Contribuciones Clave

1. Marcos de Trabajo GCD y P-GCD: La introducción de los marcos de Decodificación Globalmente Restringida y P-GCD, que proporcionan propuestas y potenciales tratables para SMC que garantizan la satisfacción de la restricción dentro de un presupuesto de tokens.
2. Autómatas Tensorizados en GPUs: La primera propuesta de tensorizar autómatas finitos (incluyendo NFAs) para una inferencia eficiente nativa en GPU, superando las limitaciones limitadas por CPU de marcos existentes como Outlines y XGrammar.
3. Soporte para NFA: Demostrar que los NFA pueden ser exponencialmente más sucintos que los DFA para restricciones complejas (por ejemplo, llamadas a funciones con parámetros variables), haciendo que el enfoque sea escalable a aplicaciones agénticas del mundo real donde la construcción de un DFA puede causar desbordamiento de memoria.
4. Integración con SMC: Mostrar que (P-)GCD acelera significativamente la convergencia de SMC en comparación con las propuestas basadas en LCD, requiriendo muchas menos partículas para alcanzar la distribución objetivo.

Resultados Experimentales

Los autores evaluaron (P-)GCD en tres bancos de pruebas: xLAM (llamada a funciones), CommonGen (generación basada en palabras clave) y Spider (text-to-SQL).

• Llamada a Funciones (xLAM):
  • (P-)GCD logró una precisión sustancialmente mayor que LCD con el mismo número de partículas.
  • P-GCD convergió a la distribución real (muestreo de rechazo) más rápido que GCD, resaltando el beneficio de la información probabilística.
  • Satisfacción de Restricciones: GCD y P-GCD lograron tasas de satisfacción de restricciones del 100% en todos los presupuestos de partículas, mientras que LCD falló en satisfacer las restricciones incluso con presupuestos moderados (ej., ~95% para JSON).
  • Eficiencia: LCD y GCD incurrieron en una ralentización modesta (1.08× y 1.10×). P-GCD fue más lento (3.11× a 10.2×) debido a los cálculos de HMM, pero sigue siendo factible.
• Generación de Palabras Clave (CommonGen): El uso de P-GCD tanto como propuesta como potencial mejoró las puntuaciones BLEU y la velocidad de convergencia significativamente sobre LCD.
• Text-to-SQL (Spider): GCD y P-GCD mejoraron la precisión de ejecución sobre LCD a través de varios presupuestos de partículas, particularmente en valores moderados.
• NFA vs. DFA: Los experimentos confirmaron que para esquemas complejos (ej., parámetros de función de longitud variable), los NFA crecen linealmente en tamaño mientras que los DFA crecen exponencialmente. Para la gramática SQL de Spider, la construcción del DFA resultó en errores de falta de memoria (OOM) al nivel de token, mientras que el NFA se mantuvo tratable (~32k estados).

Significado y Reivindicaciones

El artículo afirma que la decodificación localmente restringida introduce un sesgo de muestreo sustancial y no garantiza la satisfacción de las restricciones dentro de presupuestos finitos de tokens. Al proponer (P-)GCD, los autores ofrecen una solución que:

1. Garantiza la Validez: Asegura que cualquier secuencia generada de longitud $n$ satisfaga la restricción, eliminando los casos de falla de la LCD.
2. Reduce el Sesgo: Al incorporar información probabilística (vía P-GCD) y mirada hacia adelante global, las propuestas aproximan mejor la distribución condicional objetivo.
3. Mejora la Eficiencia: En el marco de SMC, (P-)GCD converge a la distribución objetivo con significativamente menos partículas que los enfoques basados en LCD, reduciendo el costo computacional de la generación restringida de alta calidad.

El trabajo se posiciona como un paso fundamental hacia la generación confiable y con satisfacción de restricciones para LLMs en aplicaciones que requieren un estricto cumplimiento estructural, como la llamada a funciones y la generación de código, al unificar restricciones lógicas (vía autómatas) y modelado probabilístico (vía HMMs) de una manera eficiente en GPU.