Resumo Técnico: Sticky Jump Diffusions (SJDs)

Problema

A modelagem generativa baseada em difusão estendeu-se com sucesso de dados contínuos para sequências discretas (ex: texto), mas as abordagens existentes enfrentam limitações fundamentais para unificar dinâmicas discretas e contínuas. Duas famílias dominantes existem:

1. Difusão de Máscara (Masked Diffusion): Opera em espaços de tokens discretos, gerando tokens válidos por construção. No entanto, trata cada posição como um mapeamento para um único token absorvente, apagando informações graduadas sobre a proximidade da corrupção em relação aos candidatos a tokens.
2. Difusão Contínua (Continuous Diffusion): Opera em embeddings de tokens usando aprendizado de score padrão. Seu estado terminal reside fora do suporte discreto, exigindo um passo de projeção externo (ex: um decodificador ou quantização) para recuperar tokens discretos, o que não faz parte das dinâmicas aprendidas.
3. Difusão Híbrida (Hybrid Diffusion): Pareia uma cadeia de máscara com um latente contínuo. Embora eficaz, esses modelos frequentemente montam os amostradores reversos manualmente, tratando taxas de comprometimento, destinos e drifts como escolhas de design separadas, em vez de derivá-los de um processo de avanço unificado.

O desafio central é definir um processo de Markov de tempo contínuo que una naturalmente esses regimes, onde as dinâmicas reversas (comprometimento com tokens discretos e difusão contínua) são matematicamente derivadas do processo de avanço via balanço de fluxo, em vez de serem projetadas heuristicamente.

Metodologia: Sticky Jump Diffusions (SJDs)

Os autores introduzem as Sticky Jump Diffusions (SJDs), uma classe de processos de Markov de tempo contínuo em $\mathbb{R}^d$ com um conjunto finito ou enumerável de âncoras $A$ (identificadas com embeddings de tokens).

1. Dinâmica de Avanço (Unsticking)

No tempo de avanço, os dados começam em uma âncora $a \in A$. O processo evolui da seguinte forma:

• Fase de Adesão (Sticking Phase): A partícula permanece na âncora $a$ por uma duração aleatória governada por uma taxa de risco $\vec{\lambda}_t$.
• Desadesão (Unsticking): No evento de risco, a partícula "descola" para um novo estado $Y$ amostrado de um kernel $R_t(\cdot | a)$ suportado na região contínua $X_A = \mathbb{R}^d \setminus A$.
• Difusão: Uma vez em $X_A$, a trajetória difunde-se segundo uma SDE padrão: $d\vec{X}_t = \vec{b}_t(\vec{X}_t)dt + g_t d\vec{W}_t$.
• Divisão de Massa: A distribuição marginal $p_t$ divide-se em uma densidade absolutamente contínua em $X_A$ e massas atômicas nas âncoras $A$.

Inovação Chave (Mistura de Posições Cruzadas): Ao contrário dos híbridos padrão onde o kernel de desadesão depende apenas da âncora local, as SJDs permitem que o kernel dependa de toda a sequência via uma matriz de mistura $W$. A média da distribuição de desadesão na posição $i$ é uma mistura ponderada dos embeddings dos vizinhos: $\mu_i(\vec{X}_0) = (W E(\vec{X}_0))_i$. Isso codifica vieses indutivos estruturais (ex: localidade espacial ou grafos de restrição) diretamente no processo de corrupção.

2. Dinâmica Reversa (Sticking)

A reversão temporal de uma SJD produz um processo generativo com dois mecanismos acoplados:

• SDE Baseada em Score: Na região contínua $X_A$, o processo segue uma SDE reversa impulsionada pela função de score $\nabla \log p_t(x)$.
• Kernel de Salto Adesivo (Sticky Jump Kernel): O processo salta da região contínua de volta para as âncoras. A taxa de risco reversa (quando saltar) e o destino do salto (para onde saltar) são unicamente determinados pela equação de balanço de fluxo (detalhe de equilíbrio entre os fluxos de avanço e reverso).
  • Crucialmente, a taxa de comprometimento e o destino não são cronogramas desenhados à mão, mas são fixados pela lei de avanço.

3. Aprendizado: Denoising Hazard Matching (DHM)

Para tornar o processo reverso treinável, os autores propõem o Denoising Hazard Matching (DHM), o análogo de risco para o denoising score matching.

• Classificador Único: Um único classificador de denoising $P_\theta(a | y, t)$ é treinado via cross-entropy padrão em pares corrompidos pelo SJD $(\vec{X}_0, \vec{X}_t)$.
• Recalibragem Analítica: A taxa de risco reversa e as probabilidades de alocação são recuperadas analiticamente de $P_\theta$ sem uma segunda cabeça de rede.
  • A taxa de risco reversa $\vec{\lambda}^\star_t(y)$ é estimada como a expectativa condicional de um alvo conhecido no avanço (o alvo DHM) sob a corrupção SJD.
  • A probabilidade de alocação $\pi^\star_t(a|y)$ é derivada reponderando os logits do classificador com o alvo DHM.
• Treinamento Livre de Simulação: Sob condições padrão (sobrevivência de forma fechada, fechamento Gaussiano de convolução), o treinamento não requer simulação do caminho SJD. A densidade de mistura necessária para o alvo é avaliada via quadratura de forma fechada, tornando o pipeline de treinamento eficiente e livre de simulação.

Principais Contribuições

1. A Classe de Processo SJD: Um framework unificado para jump-diffusions onde a marginal de avanço se divide em uma densidade contínua e átomos discretos. Isso generaliza as clássicas difusões de fronteira adesiva para conjuntos de âncoras de alta dimensão.
2. Dinâmica Reversa a partir de um Único Classificador: Os autores derivam que tanto o score quanto as intensidades reversas por âncora podem ser estimados de um único classificador de denoising via DHM. Isso elimina a necessidade de cabeças de risco separadas e baseia-se em um reponderamento de Rao–Blackwellized da saída do classificador.
3. Visão Unificadora de Famílias de Difusão: As SJDs recuperam modelos existentes como limites específicos:
   • Difusão de Máscara (Masked Diffusion): Surge quando o kernel de desadesão colapsa para um único ponto absorvente. O "vazio de informação" na difusão de máscara é explicado estruturalmente: o kernel de desadesão é independente da âncora, de modo que o estado corrompido não carrega evidência do token de origem.
   • Difusão Contínua (Continuous Diffusion): Surge quando a taxa de salto é zero. A necessidade de uma projeção terminal é explicada pela ausência de átomos na marginal de avanço, o que não gera saltos reversos.
   • Difusão Híbrida (Hybrid Diffusion): Surge quando o kernel de desadesão coincide com o kernel de perturbação de avanço. Mostra-se que a taxa de comprometimento, o destino e o drift seguem naturalmente do balanço de fluxo, em vez de serem designs separados.
4. O Kernel de Salto como Espaço de Design: Ao variar a matriz de mistura $W$, os autores introduzem um novo espaço de design. A mistura de posições cruzadas corrompe cada posição em direção a uma mistura dos valores limpos de seus vizinhos, transformando estruturas de dependência (espaciais, sequenciais ou baseadas em restrições) em um viés indutivo da própria corrupção.

Resultados Experimentais

Os autores validam empiricamente o eixo de mistura de posições cruzadas em três datasets onde as dependências estruturais são conhecidas:

• CIFAR-10 (Geração de Imagem): Usando um kernel de desfoque de espaço de valores para codificar localidade espacial, a SJD com um kernel não-local ($\sigma_B = 2.5$) alcança um FID menor (14.57) que o híbrido de kernel identidade (CADD, 15.88) e outros baselines (MDLM, Bit Diffusion) a 256 NFE.
• Text8 (Geração de Texto): Usando uma convolução Gaussiana 1D para codificar correlação sequencial local, a SJD com $\sigma_W \ge 1.0$ supera o CANDI e o MDLM em contagens de geração de palavras válidas, particularmente em sequências longas ($\ge 6$) e em vários orçamentos de NFE.
• Sudoku (Satisfação de Restrições): Usando um kernel de grafo de restrição para codificar dependências de linha/coluna/bloco, a SJD supera significativamente o híbrido de kernel identidade (CADD). Enquanto o CADD exibe desempenho bimodal (frequentemente colapsando para próximo do acaso), a SJD resolve os tabuleiros consistentemente, alcançando 95.65% de precisão e convergindo 4x mais rápido em etapas de treinamento.

Significância e Alegações

O artigo afirma que as Sticky Jump Diffusions fornecem um relato ao nível de processo que unifica as difusões de máscara, contínuas e híbridas.

• Unificação Teórica: Demonstra que as coordenadas discretas e contínuas na modelagem generativa não são entidades separadas que devem ser acopladas manualmente, mas sim um único estado revertido onde as taxas de comprometimento e destinos são fixados pelo balanço de fluxo.
• Expansão do Espaço de Design: Ao tratar o kernel de desadesão como um espaço de design explícito, as SJDs permitem que pesquisadores codifiquem priors estruturais (como localidade espacial ou restrições lógicas) diretamente nas dinâmicas de corrupção de avanço. Isso melhora a qualidade da geração e a consistência lógica sem alterar a arquitetura fundamental do amostrador reverso.
• Eficiência Prática: O framework treina e amostra com o mesmo custo que os híbridos existentes, pois a estrutura específica da SJD entra apenas através de reponderagem analítica e quadraturas de forma fechada, não exigindo parâmetros de rede adicionais ou overhead de simulação.

Os autores concluem que esta visão unificada torna mais fácil estender os modelos de difusão: novas corrupções, cronogramas e regras de comprometimento podem ser especificados como dinâmicas de avanço e obtidos em reverso gratuitamente via teorema de reversão de tempo.