Resumo Técnico: Modelos de Fluxo de Linguagem Mascarada (MLFMs)

Declaração do Problema

Modelos Autoregressivos (ARMs) dominam a geração de linguagem, mas sofrem com gargalos de inferência sequencial, já que cada token depende de todos os anteriores. Modelos de Difusão Mascarada (MDMs) oferecem uma alternativa de decodificação paralela através do desmascaramento iterativo de tokens; no entanto, eles dependem de uma transição reversa fatorizada que assume independência condicional entre posições mascaradas. Essa aproximação é válida no limite de passos infinitesimais, mas degrada significamente em regimes de amostragem de poucos passos, onde a geração paralela é mais eficiente, levando a uma qualidade comprometida devido à negligência das dependências entre tokens.

Inversamente, Modelos de Fluxo de Linguagem (FLMs) operam em espaços de embedding contínuos, aprendendo um fluxo que transporta o ruído para sequências limpas. Embora os FLMs suportem a geração eficiente de passo único via destilação de mapa de fluxo, eles tipicamente colapsam a geração em uma única trajetória. Essa estrutura tem dificuldade com tarefas que exigem raciocínio de múltiplos passos, pois os FLFMs são forçados a decodificar todos os tokens simultaneamente, em vez de utilizar estados intermediários parcialmente completados como contexto para refinamentos subsequentes.

Metodologia: Modelos de Fluxo de Linguagem Mascarada (MLFMs)

Os autores propõem os Modelos de Fluxo de Linguagem Mascarada (MLFMs), um framework que integra o mecanismo de mascaramento dos MDMs à dinâmica de fluxo contínuo dos FLMs.

1. Framework e Interpolante Estocástica

Os MLFMs definem um processo direto no espaço Euclidiano usando um ponte Browniana como um interpolante estocástico. Esta ponte conecta:

• Sequências parcialmente mascaradas: Representadas por embeddings limpos para tokens não mascarados e um embedding [MASK] específico para tokens mascarados.
• Sequências limpas: Representadas por seus embeddings de tokens limpos.

Para uma sequência $X$, no tempo de mascaramento $s$ e tempo de ruído $t$, o estado $z_{s,t}$ é construído de tal forma que:

• Posições não mascaradas ($U_s$) permanecem fixas em seus embeddings limpos $E X_\ell$.
• Posições mascaradas ($M_s$) evoluem via uma ponte Browniana entre o embedding [MASK] ($m$) e o embedding limpo ($E X_\ell$).

Este design permite que o modelo aprenda um fluxo contínuo acoplado para gerar tokens limpos em posições mascaradas, mantendo ao mesmo tempo a estrutura de condicionamento de qualquer posição exata dos MDMs.

2. Treinamento e Adaptação

Uma contribuição fundamental é a capacidade de realizar o warm-start do treinamento do MLFM utilizando MDMs pré-treinados.

• Equivalência Teórica: No ponto final $t=0$, o interpolante estocástico do MLFM reduz-se exatamente à sequência parcialmente mascarada embarcada usada nos MDMs. Consequentemente, a perda do MLFM em $t=0$ coincide com a perda do MDM.
• Estratégia de Adaptação: Os autores adaptam um transformer bidirecional pré-treinado (de um MDM) através de:
  1. Manter a camada de embedding de tokens fixa.
  2. Congelar os blocos do transformer e aumentá-los com AdaLN (para condicionamento no tempo contínuo $t$) e adaptadores LoRA (para adaptação a embeddings continuamente corrompidos).
  3. Adaptar o cabeçalho de saída com um adaptador LoRA.
     Esta abordagem reduz significativamente os custos computacionais em comparação ao treinamento do zero.

3. Estratégia de Amostragem: CCFG com Promoção de Token Online

Para abordar a limitação de raciocínio multi-etapa, os autores introduzem um novo esquema de amostragem combinando dois mecanismos:

• Context-Corrupted Classifier-Free Guidance (CCFG): Uma variante de orientação livre de classificador onde o campo de velocidade de referência é computado corrompendo os tokens de contexto (aplicando a marginal de corrupção direta a eles), enquanto o campo guiado utiliza o contexto limpo. A diferença isola a contribção do contexto limpo para a direção prevista.
• Online Token Promotion (OTP): Durante a amostragem, se o modo posterior para um token mascarado atingir um limiar de confiança ($1-\epsilon$), o token é imediatamente "promovido". Ele é fixado ao seu embedding limpo e tratado como contexto observado para todos os passos de denoising subsequentes. Isso permite que o modelo aproveite previsões de alta confiança como contexto imediato, mimetizando o raciocínio iterativo.

Principais Contribuições

1. Integração de Mascaramento e Fluxo: Os MLFMs unem o mascaramento discreto (MDMs) e os fluxos contínuos (FLMs) usando uma ponte Browniana, permitindo a geração condicional que suporta tanto a eficiência paralela quanto o refinamento iterativo.
2. Adaptação Eficiente: A demonstração de que MDMs pré-treinados podem ser convertidos em MLFMs via adaptação leve (LoRA/AdaLN), aproveitando a equivalência de endpoint em $t=0$.
3. Sampler Inovador: A proposta do CCFG combinado com a Promoção de Token Online, que converte dinamicamente estados latentes em contexto limpo, facilitando o raciocínio de múltiplos passos.
4. Escalabilidade para Tarefas de Downstream: O primeiro escalonamento bem-sucedido de modelos de linguagem baseados em fluxo para tarefas de raciocínio (matemática) e seguimento de instruções.

Resultados Experimentais

Os autores avaliaram um MLFM de 1.028B de parâmetros (adaptado de um SMDM pré-treinado) em GSM8K (raciocínio matemático) e MT-Bench (seguimento de instrução).

• MT-Bench: O MLFM alcançou uma pontuação de primeiro turno de 2.27, superando significativamente o baseline SMDM (1.60) e um baseline AR de tamanho similar (1.57). Notavelmente, isso foi alcançado com apenas 128 passos de amostragem (metade dos 256 do SMDM).
• GSM8K: O MLFM alcançou 31.24% de acurácia. Embora fique atrás do SMDM ajustado para a tarefa específica (58.5%) e do LLaMA-2 (58.6%), os autores atribuem isso ao protocolo de ajuste fino de mistura de instruções mais amplo usado para o MLFM versus o protocolo específico da tarefa usado para os baselines.
• Ablações: Os resultados confirmaram que tanto o CCFG quanto a Promoção de Token Online aumentam significativamente o desempenho. Escalas de orientação maiores e mais passos de amostragem geralmente melhoraram os resultados, com o MLFM superando os baselines no MT-Bench mesmo com 16 passos de amostragem.

Significância e Alegações

O artigo afirma fornecer a primeira evidência de que modelos de linguagem baseados em fluxo podem ser escalonados para resolver tarefas de raciocínio e seguimento de instruções. Ao combinar o potencial de geração paralela dos modelos de fluxo com as capacidades de raciocínio iterativo dos modelos mascarados, os MLFMs oferecem uma arquitetura flexível que suporta geração complexa de múltiplos passos. O trabalho demonstra que os modelos de fluxo não estão limitados à geração incondicional ou destilação simples, mas podem ser adaptados para aplicações práticas e condicionais de downstream via adaptação de MDM eficiente.