Resumo Técnico: Mitigando o Viés na Decodificação Localmente Restrita via Propostas Tratáveis

Declaração do Problema

Grandes Modelos de Linguagem (LLMs) frequentemente falham em gerar saídas que adiram estritamente a restrições lógicas ou estruturais, como esquemas JSON, consultas SQL válidas ou formatos específicos de chamada de função. Embora o ajuste de instrução (instruction tuning) e a otimização de preferência tenham melhorado o alinhamento geral, eles permanecem não confiáveis para impor restrições "rígidas".

Abordagens existentes enfrentam duas limitações primárias:

1. Amostragem de Rejeição (Rejection Sampling): Embora seja não enviesada, é computacionalmente impraticável para restrições complexas devido às taxas de aceitação decrescentemente baixas.
2. Decodificação Localmente Restrita (LCD): Esta abordagem padrão mascara tokens que violam imediatamente uma restrição em cada etapa. No entanto, ela sofre de duas falhas críticas:
   • Viés: Ela distorce a distribuição subjacente do modelo ao mascarar tokens de forma míope sem considerar a validade futura, levando ao viés de amostragem.
   • Sem Garantia de Terminação: A LCD garante apenas a validade local. Uma sequência pode permanecer localmente válida indefinidamente (ex: uma chave não fechada em um JSON) mas falhar em satisfazer a restrição dentro de um orçamento fixo de $n$ tokens.

Para abordar esses problemas, foi proposto o framework de Monte Carlo Sequencial (SMC), utilizando a LCD como uma distribuição de proposta e reponderando partículas para convergir para a distribuição alvo. A taxa de convergência do SMC é limitada pela qualidade da proposta; a LCD é uma proposta fraca porque carece de informação probabilística sobre conclusões válidas futuras e não garante a satisfação da restrição dentro de um orçamento finito.

Metodologia

Os autores propõem uma abordagem genérica para construir propostas tratáveis e funções de potencial para a amostragem SMC da distribuição restrita $p_{lm}(\cdot \mid C)$. A inovação central envolve representar restrições como autômatos finitos e aproveitar operações de tensores para execução eficiente em GPU.

1. Decodificação Globalmente Restrita (GCD)

A GCD melhora a LCD ao garantir que cada token gerado possa levar a uma conclusão válida dentro do orçamento de tokens restante $n$.

• Autômatos Tensorizados: Restrições codificadas como Autômatos Finitos (FA), especificamente Autômatos Finitos Não-determinísticos (NFA) e Autômatos Finitos Determinísticos (DFA), são compiladas em representações tensorizadas (matrizes de transição $T_{out}, T_{in}$ e matriz de emissão $B_f$). Isso permite a execução eficiente em GPUs.
• Inferência Forward-Backward:
  • Processo Forward: Computa o conjunto de estados do FA alcançáveis após observar o prefixo $x_{1:t-1}$.
  • Processo Backward: Computa se um estado pode alcançar um estado de aceitação em exatamente $n-t$ passos.
• Mascaramento Global: Um token é permitido apenas se levar a um estado que admita uma continuação válida dentro do orçamento restante. Isso garante que qualquer sequência de comprimento $n$ gerada sob a GCD satisfaça a restrição $C$.

2. Decodificação Globalmente Restrita Probabilística (P-GCD)

Embora a GCD garanta a satisfação da restrição, sua máscara é binária (válido/inválido) e ignora a verossimilhança das conclusões válidas sob o LLM. A P-GCD aborda isso integrando informações probabilísticas.

• Destilação de HMM: A distribuição do LLM é aproximada por um Modelo Oculto de Markov (HMM), seguindo o trabalho anterior (Zhang et al., 2023).
• Multiplicação de Circuitos: O FA tensorizado (visto como um HMM não normalizado sobre arestas) é multiplicado com o HMM destilado usando algoritmos de multiplicação de circuitos. Isso aproveita a estrutura algébrica compartilhada de FAs e HMMs.
• Distribuição Produto: A distribuição produto resultante $p_{prod}$ captura tanto as restrições lógicas rígidas do FA quanto as preferências probabilísticas do LLM.
• Proposta e Potencial:
  • Proposta ($q_{pgcd}$): Repondera a distribuição do LM em cada etapa usando a probabilidade do próximo token restrito derivada de $p_{prod}$.
  • Potencial ($\phi_{pgcd}$): Pondera as partículas no SMC pelo produto da probabilidade do prefixo do LM e a estimativa de look-ahead do HMM para a satisfação da restrição.

3. Implementação Eficiente

Os autores utilizam operações de redes de tensores (especificamente a identidade Kronecker-vec) para evitar a materialização explícita de grandes matrizes de produto de Kronecker. Isso reduz o custo computacional das etapas de transição de $O((HE)^2)$ para $O(H^2E + HE^2)$, onde $H$ é o tamanho oculto do HMM e $E$ é a contagem de arestas do FA, permitindo a geração autorregressiva eficiente.

Principais Contribuições

1. Frameworks GCD e P-GCD: A introdução da Decodificação Globalmente Restrita e da GCD Probabilística, que fornecem propostas e potenciais tratáveis para SMC que garantem a satisfação da restrição dentro de um orçamento de tokens.
2. Autômatos Tensorizados em GPUs: A primeira proposta de tensorizar autômatos finitos (incluindo NFAs) para inferência eficiente nativa em GPU, superando as limitações limitadas à CPU de frameworks existentes como Outlines e XGrammar.
3. Suporte a NFA: Demonstrar que os NFAs podem ser exponencialmente mais sucintos que os DFAs para restrições complexas (ex: chamada de função com parâmetros variáveis), tornando a abordagem escalável para aplicações de agentes do mundo real onde a construção de DFA pode causar estouro de memória.
4. Integração com SMC: Mostrar que a (P-)GCD acelera significativamente a convergência do SMC em comparação com as propostas baseadas em LCD, exigindo muito menos partículas para atingir a distribuição alvo.

Resultados Experimentais

Os autores avaliaram a (P-)GCD em três benchmarks: xLAM (chamada de função), CommonGen (geração baseada em palavras-chave) e Spider (text-to-SQL).

• Chamada de Função (xLAM):
  • A (P-)GCD alcançou uma precisão substancialmente maior que a LCD para o mesmo número de partículas.
  • A P-GCD convergiu para a distribuição real (amostragem de rejeição) mais rapidamente que a GCD, destacando o benefício da informação probabilística.
  • Satisfação de Restrição: GCD e P-GCD alcançaram taxas de satisfação de restrição de 100% em todos os orçamentos de partículas, enquanto a LCD falhou em satisfazer as restrições mesmo em orçamentos moderados (ex: ~95% para JSON).
  • Eficiência: LCD e GCD incorreram em um overhead modesto (lentidão de 1.08× e 1.10×). A P-GCD foi mais lenta (3.11× a 10.2×) devido aos cálculos de HMM, mas permaneceu viável.
• Geração de Palavras-Chave (CommonGen): O uso da P-GCD como proposta e potencial melhorou as pontuações BLEU e a velocidade de convergência significativamente em relação à LCD.
• Text-to-SQL (Spider): A GCD e a P-GCD melhoraram a precisão de execução sobre a LCD em vários orçamentos de partículas, particularmente em valores moderados.
• NFA vs. DFA: Experimentos confirmaram que, para esquemas complexos (ex: parâmetros de função de comprimento variável), os NFAs crescem linearmente em tamanho enquanto os DFAs crescem exponencialmente. Para a gramática SQL do Spider, a construção do DFA resultou em erros de Out-of-Memory (OOM) ao nível do token, enquanto o NFA permaneceu tratável (~32k estados).

Significância e Alegações

O artigo afirma que a decodificação localmente restrita introduz um viés de amostragem substancial e falha em garantir a satisfação da restrição dentro de orçamentos finitos de tokens. Ao propor a (P-)GCD, os autores oferecem uma solução que:

1. Garante Validez: Garante que qualquer sequência gerada de comprimento $n$ satisfaça a restrição, eliminando os casos de falha da LCD.
2. Reduz o Viés: Ao incorporar informações probabilísticas (via P-GCD) e lookahead global, as propostas aproximam-se melhor da distribuição condicional alvo.
3. Melhora a Eficiência: No framework SMC, a (P-)GCD converge para a distribuição alvo com significativamente menos partículas do que as abordagens baseadas em LCD, reduzindo o custo computacional da geração restrita de alta qualidade.

O trabalho posiciona-se como um passo fundamental para a geração confiável e que satisfaz restrições para LLMs em aplicações que exigem aderência estrutural estrita, como chamadas de função e geração de código, ao unificar restrições lógicas (via autômatos) e modelagem probabilística (via HMMs) de maneira eficiente em GPU.