Resumo Técnico: Amostragem de Posterior Inspirada em Langevin Discreto (∆LPS)

Declaração do Problema

O artigo aborda o desafio de resolver problemas inversos em espaços de estados discretos. Problemas inversos envolvem recuperar um sinal desconhecido $x$ a partir de medições indiretas, esparsas e ruidosas $y$ (modeladas como $y = A(x) + n$). Embora os modelos de difusão contínuos tenham se tornado o padrão para resolver tais problemas via amostragem de posterior, suas contrapartes discretas permanecem pouco exploradas.

Os métodos existentes para amostragem de posterior discreta sofrem de três limitações primárias:

1. Relaxações Contínuas: Métodos como G2D2 relaxam variáveis categóricas em espaços contínuos para usar orientação por gradiente, alterando a geometria de inferência e arriscando que os tokens discretos projetados não decodifiquem consistentemente.
2. Atualizações Sequenciais: Métodos como SGDD dependem de atualizações no estilo de Gibbs, que são sequenciais e lentas para misturar em espaços de alta dimensão.
3. Especialização: Métodos como APS estão fortemente acoplados a processos de corrupção específicos (por exemplo, difusão mascarada) e paradigmas de treinamento, limitando sua generalidade.

O objetivo é desenvolver um amostrador de posterior que permaneça totalmente discreto, atualize todas as dimensões de token em paralelo e seja agnóstico ao paradigma de treinamento específico do prior de difusão discreto.

Metodologia: ∆LPS

Os autores propõem o ∆LPS (Amostrador de Posterior Inspirado em Langevin Discreto), um algoritmo sem treinamento que utiliza informações de gradiente para guiar transições discretas sem sair do espaço de estados discretos.

Formulação Central

O problema é formulado no espaço latente discreto $z_0 \in \{1, \dots, K\}^L$ de um modelo codificador-decodificador. O objetivo é amostrar a partir do posterior $p(z_0|y) \propto p(y|z_0)p_\theta(z_0)$, onde $p_\theta$ é um prior de difusão discreto pré-treinado. O amostrador aproxima a transição do posterior $p(z_s|z_t, y)$ inferindo primeiro um latente limpo $z_0$ consistente com a medição e o prior, e depois re-ruidando-o para o próximo passo de tempo.

Proposta de Langevin Discreto

A inovação chave é adaptar a dinâmica de Langevin para espaços discretos. Em configurações contínuas, uma atualização de Langevin move uma amostra $z_0$ na direção do gradiente do potencial log-posterior $U(z_0) = \log p(y|z_0) + \log p_\theta(z_0; z_t)$. Como $z_0$ é discreto, os autores não adicionam um vetor de gradiente diretamente. Em vez disso, eles usam a distribuição proposta Gaussiana da dinâmica de Langevin contínua para definir probabilidades de transição sobre candidatos discretos válidos $z'_0$:

$$q(z'_0|z_0; z_t, y) \propto \exp\left( -\frac{1}{4\eta} \|z'_0 - z_0 - \eta \nabla_{z_0} U(z_0; z_t, y)\|_2^2 \right)$$

Ao expandir o expoente e normalizar sobre o conjunto finito de tokens válidos, a proposta fatora através das coordenadas do token. Isso permite que o amostrador compute uma distribuição categórica para cada posição de token $\ell$ independentemente:
$$q_\ell(z'_0[\ell]) = \text{Cat}\left( \sigma\left( \frac{1}{2} \nabla_{z_0[\ell]} U \cdot (k - z_0[\ell]) - \frac{(k - z_0[\ell])^2}{4\eta} \right) \right)$$

Este mecanismo permite atualizações totalmente paralelas através de todas as $L$ dimensões de token, evitando o gargalo sequencial dos amostradores de Gibbs.

Generalização e Pré-condicionamento

• Agnosticismo ao Prior: O método requer apenas probabilidades preditivas sobre estados categóricos, tornando-o compatível com priores de difusão mascarada (MDLM) e difusão de estado uniforme (Duo).
• Problemas Inversos Cegos: Para operadores diretos não diferenciáveis ou desconhecidos, o gradiente de verossimilhança é substituído por um surrogado contrastivo aprendido (usando codificadores para dados e medições) para fornecer um sinal de orientação suave.
• Pré-condicionamento Adam: Para lidar com gradientes mal escalados em espaços latentes de alta dimensão, o gradiente bruto é substituído por uma direção pré-condicionada por Adam. Isso reescala adaptativamente as atualizações, amortecendo gradientes grandes e ruidosos enquanto preserva aqueles consistentemente úteis, melhorando a mistura e a convergência.

Contribuições Principais

1. Amostragem de Posterior Totalmente Discreta: O ∆LPS é um amostrador sem treinamento que mantém cada iteração no espaço de estados discreto, evitando otimização de tokens suaves e erros de projeção.
2. Atualizações Paralelas Eficientes: Ao fatorar a proposta através das coordenadas, o método atualiza todas as dimensões de token em paralelo, melhorando significativamente a eficiência sobre amostradores sequenciais no estilo de Gibbs.
3. Design Agnóstico ao Prior: O amostrador não está vinculado a um processo de corrupção específico (mascarado vs. uniforme) e funciona com qualquer prior de difusão discreto pré-treinado que forneça probabilidades preditivas.

Resultados Experimentais

Os autores avaliam o ∆LPS em um conjunto diversificado de problemas inversos em MNIST, CIFAR-10 e FFHQ, cobrindo configurações lineares, não lineares e cegas.

• Restauração de Imagem (FFHQ): Em tarefas incluindo inpainting, desfocagem gaussiana, super-resolução, HDR e desfocagem de movimento, o ∆LPS alcança desempenho comparável ou superior às bases contínuas (por exemplo, DPS, PSLD) e supera as bases discretas existentes (SGDD, G2D2, APS) em qualidade perceptual (LPIPS).
• Tarefas Discretas (MNIST/CIFAR): O método demonstra robustez através de níveis variados de dificuldade (Fácil, Médio, Difícil). Ele consistentemente iguala ou supera o SGDD e o G2D2 em precisão de token e PSNR, mesmo com orçamentos computacionais significativamente menores (menor número de avaliações de função) do que a variante "completa" do SGDD.
• Problemas Inversos Cegos: Em uma tarefa de mapeamento espacial (inferindo plantas baixas a partir de trajetórias de caminhada onde o operador direto é desconhecido), o ∆LPS supera as bases contínuas (CoGuide, variantes DPS) em pontuação F1 e Interseção sobre União (IoU).

Significado e Alegações

O artigo afirma que amostradores de posterior totalmente discretos e informados por gradiente oferecem um caminho escalável e geral para resolver problemas inversos sobre representações discretas.

• Escalabilidade: O mecanismo de atualização paralela aborda as limitações de velocidade de mistura dos amostradores discretos em espaços de alta dimensão.
• Generalidade: O método desacopla o algoritmo de inferência do paradigma de treinamento específico do prior, permitindo alavancar o ecossistema crescente de modelos de difusão discreta (texto, código, imagens, áudio).
• Praticidade: Ao evitar relaxações contínuas, o ∆LPS garante que as amostras finais sejam tokens discretos válidos, prevenindo a "lacuna de projeção" onde uma aproximação contínua falha em decodificar em um sinal discreto consistente.

Os autores concluem que, à medida que os modelos generativos discretos escalam, métodos de inferência principistas como o ∆LPS tornar-se-ão cada vez mais críticos para aplicações em design biomolecular, música simbólica, conclusão de código e outros domínios onde os dados são inerentemente discretos.