Resumo Técnico: Mapas de Fluxo Discretos com Núcleo Latente para Geração de Poucos Passos

Declaração do Problema

Modelos de difusão discreta e de correspondência de fluxo (flow-matching) geram texto ao remover o ruído de uma sequência ao longo de muitos passos. Para manter a eficiência computacional, esses modelos tipicamente fatorizam a probabilidade de transição através das posições dos tokens, decidindo cada token independentemente dado o estado atual da sequência. Embora essa fatorização torne os passos individuais baratos, ela cria um gargalo fundamental para a geração de poucos passos.

Quando uma sequência alvo requer correlações fortes entre posições (ex: concordância sujeito-verbo ou parênteses correspondentes), atualizações independentes comprometem-se com tokens separadamente. Isso frequentemente resulta em combinações localmente consistentes, mas globalmente inconsistentes (ex: "dog barks" vs. "dogs bark"). Corrigir esses desajustes requer muitos passos adicionais de remoção de ruído, negando os ganhos de eficiência da geração de poucos passos.

As soluções existentes para acelerar a geração dividem-se em duas categorias:

1. Modelos de verossimilhança (ex: MDLM, SEDD): Treinam denoisers fatorizados fortes do zero, mas exigem muitos passos na inferência para resolver correlações.
2. Aceleradores de poucos passos (ex: Di4C, SDTT, ReDi): Destilam ou retificam um modelo "professor" treinado para um aluno de poucos passos. No entanto, esses métodos herdam o teto de qualidade do professor e não podem superar o desempenho do modelo que imitam.

O desafio central é se um modelo pode expressar passos correlacionados nativamente sem depender de um professor, rompendo assim a suposição de fatorização em sua raiz, mantendo, ao mesmo tempo, um baixo custo computacional.

Metodologia: Mapas de Fluxo Discretos com Núcleo Latente (LKF)

Os autores propõem os Mapas de Fluxo Discretos com Núcleo Latente (LKF), um modelo de mapa de fluxo construído do zero que introduz correlação intra-passo através de uma mistura de componentes fatorizados vinculados por uma única variável latente compartilhada.

A Arquitetção do Modelo

O núcleo de transição de tempo finito $\Psi_{t \to s}(x_s | x_t)$ é definido como uma mistura de $M$ componentes fatorizados:
$$\Psi_{t \to s}(x_s | x_t) = \sum_{k=1}^{M} w_k(x_t, t, s) \prod_{i=1}^{L} P_i(x_s^i | x_t, k, t, s)$$
Onde:

• $k \in [M]$ é um índice de componente latente compartilhado amostrado uma vez por sequência.
• $w_k$ são pesos do roteador que predizem a probabilidade do componente $k$.
• $P_i$ são distribuições de tokens fatorizadas para a posição $i$ condicionadas ao latente $k$.

Mecanismo Chave:

1. Latente Compartilhado: No início de uma trajetória de remoção de ruído (ou a cada passo), um latente $k$ é extraído do roteador $w(\cdot | x_t)$.
2. Independência Condicional: Dado $k$, os tokens ainda são amostrados independentemente entre as posições. No entanto, como $k$ é compartilhado por toda a sequência, a marginalização sobre $k$ induz correlações entre as posições.
3. Eficiência: O modelo utiliza uma arquitetura de injeção de latente tardia. Um tronco Transformer compartilhado processa a entrada para criar uma representação agnóstica ao latente. Apenas os últimos $L_k$ blocos são replicados para cada um dos $M$ componentes, condicionados a $k$. Isso mantém o custo computacional da marginalização baixo ($O(M \cdot L_k)$ em vez de $O(M \cdot \text{profundidade})$).

Treinamento e Inferência

• Treinamento: O modelo é treinado marginalizando o latente $k$ exatamente. O objetivo é a log-verossimilhança da mistura:
  $$\log \Psi = \log \sum_{k} \exp(\log w_k + \ell_k)$$
  onde $\ell_k$ é a log-probabilidade atribuída pelo componente $k$ aos tokens alvo. Isso força o roteador a aprender uma mistura onde o latente captura correlações de sequência. Um regularizador do roteador garante que os componentes da mistura sejam utilizados (evitando o colapso para um único componente) e que o latente carregue informação suficiente.
• Inferência (Best-of-M): Para aproveitar a heterogeneidade da mistura, os autores empregam uma estratégia de decodificação Best-of-M. $M$ execuções paralelas são geradas, cada uma congelando um componente diferente $k$. A saída final é selecionada com base no próprio limite inferior negativo da ELBO de sequência (NELBO) do modelo. Esta seleção não requer um modelo de recompensa externo ou um professor.

Principais Contribuições

1. Passos Correlacionados Nativos: Introdução do LKF, um núcleo de mapa de fluxo que expressa passos correlacionados nativamente via uma mistura de $M$ componentes fatorizados vinculados por um único latente compartilhado, eliminando a necessidade de um modelo professor.
2. Análise Teórica: Uma decomposição de informação exata mostrando que a correlação total (TC) que um passo de rank-$M$ pode carregar é limitada por $(L-1)H(k|xt)$. Os autores provam que o erro de aproximação do modelo é limitado inferiormente pela correlação total não capturada da distribuição alvo.
3. Validação Sintética: Experimentos em conjuntos de dados sintéticos (acordo oculto e paridade) demonstram que o LKF captura informação e melhora a verossimilhança monotonicamente com $M$ para correlações representáveis, enquanto falha corretamente em capturar correlações que exigem um rank de mistura exponencial (paridade), alinhando-se com resultados teóricos de dificuldade.
4. Desempenho Empírico: Nos benchmarks LM1B e WikiText-103, o LKF alcança perplexidade generativa de estado da arte. Com $M=8$, o LKF melhora a perplexidade generativa em 2,1× a 3,3× em relação aos baselines de verossimilhança (MDLM, SEDD) e supera os aceleradores destilados/retificados (Di4C, ReDi) sem utilizar qualquer professor.

Resultados Experimentais

• Benchmarks: Avaliado em LM1B e WikiText-103.
• Desempenho:
  • LM1B: O modelo LKF com $M=8$ alcança uma perplexidade generativa (gen-PPL) de 105 com 32 avaliações de função (NFE), comparado a 304 para o MDLM e 117 para o melhor acelerador retificado (ReDi).
  • WikiText-103: O modelo $M=8$ atinge uma gen-PPL de 75, superando significativamente o MDLM (228) e o SEDD (247).
• Diversidade: Ao contrário dos aceleradores retificados que frequentemente sofrem de colapso de entropia (produzindo texto de baixa diversidade), o LKF mantém alta entropia de amostragem enquanto melhora a verossimilhança.
• Ablações:
  • O ganho é atribuído à própria estrutura de mistura, não apenas ao reranking. Substituir os $M$ componentes por $M$ desenhos de ruído de um único componente resulta em uma queda significativa de desempenho (ex: 35% de PPL maior em $M=8$).
  • A seleção Best-of-M é crucial; uma política de compromisso único (commit-M) produz resultados semelhantes ao baseline, confirmando que a mistura fornece diversidade de candidatos que a seleção explora.
  • Eficiência do Roteador: O roteador atua como um peso de mistura, não como um sinal de qualidade; a seleção é mais eficaz quando baseada na própria verossimilhança do modelo (evidência em execução) em vez da previsão inicial do roteador.

Significância e Alegações

O artigo afirma que o LKF representa uma mudança da aceleração "dependente de professor" para a geração de poucos passos "do zero". Ao modelar o núcleo de transição como uma mistura de baixo rank, o LKF permite que um único passo coloque massa em conclusões correlacionadas (ex: concordância sujeito-verbo) com a mesma complexidade de amostragem de um modelo fatorizado.

Os autores enfatizam que:

• O método não herda o teto de qualidade de um professor.
• O ganho escala com o tamanho da mistura $M$.
• A abordagem é teoricamente fundamentada, com limites para a capacidade de correlação e erro de aproximação.
• Ele alcança a melhor perplexidade generativa de poucos passos entre os métodos de baseline em benchmarks padrão de modelagem de linguagem, oferecendo um caminho viável para a geração discreta eficiente e de alta qualidade sem o overhead de destilação ou retificação.