Resumo Técnico: Modelagem de Linguagem com Fluxos Esféricos (S-FLM)

Declaração do Problema

Os modelos autoregressivos (AR) dominam a modelagem de linguagem devido à sua escalabilidade e avaliação rápida de verossimilhança, mas sofrem com a geração sequencial e a falta de contexto bidirecional durante a inferência. Os modelos de difusão discreta oferecem geração paralela e contexto bidirecional, mas são limitados por restrições de tratabilidade que os forçam a amostrar a partir de marginais fatoradas, reduzindo a expressividade em comparação com os modelos AR.

Os Modelos de Fluxo de Linguagem (FLMs) recentes tentam preencher essa lacuna aplicando fluxos contínuos à linguagem, transportando ruído para dados por meio de uma Equação Diferencial Ordinária (ODE) determinística. No entanto, os FLMs existentes enfrentam duas deficiências críticas:

1. Custo Computacional: Eles representam tokens como vetores one-hot de dimensão $|V|$ (tamanho do vocabulário). Para vocabulários grandes (100k–200k), isso exige armazenar e manipular vetores de alta dimensão, tornando o treinamento significativamente mais lento do que em modelos de difusão discreta ou AR.
2. Ambiguidade Semântica: Adicionar ruído gaussiano a vetores one-hot carece de uma interpretação semântica clara. Diferentemente de imagens, onde o ruído degrada progressivamente a estrutura, adicionar ruído gaussiano a embeddings one-hot equidistantes não possui um significado geométrico intuitivo.

Metodologia

Os autores propõem o Modelo de Fluxo de Linguagem Esférico (S-FLM), um modelo de fluxo riemanniano definido na hipersfera unitária $S^{d-1}$, e não no espaço euclidiano ou em espaços de estados discretos.

Mecanismos Principais

• Espaço Latente Esférico: Em vez de vetores one-hot, os tokens são mapeados para embeddings de norma unitária $\hat{e}_v \in S^{d-1}$. O processo forward transporta embeddings limpos em direção a uma prior uniforme na esfera por meio de Interpolação Linear Esférica (SLERP), e não de interpolação linear com ruído gaussiano.
• Correspondência de Fluxo Riemanniano: O modelo aprende um campo de velocidade $u_t$ na variedade $S^{d-1}$ que define uma ODE. A velocidade condicional é derivada usando o mapa logarítmico na esfera:
  $$u_{t|1}(z_t | z_1) = \frac{\dot{\alpha}_t}{1-\alpha_t} \log_{z_t}(z_1)$$
  onde $z_1$ é o embedding limpo e $z_t$ é o latente ruidoso.
• Objetivo de Treinamento: O denoisador $p_\theta^{1|t}$ é treinado usando Entropia Cruzada (CE) para prever o token limpo $x_\ell$ dado o latente ruidoso $z_t$. Isso evita o mínimo trivial onde todos os embeddings colapsam para um ponto, um risco ao regredir velocidades diretamente contra embeddings aprendíveis.
• Amostragem: Na inferência, o modelo integra o campo de velocidade aprendido a partir de ruído uniforme na esfera até a variedade de dados. O campo de velocidade é computado marginalizando sobre vetores tangentes apontando para cada embedding limpo, ponderados pela posterior $p_\theta^{1|t}$.
  • Velocidade Exata: Soma sobre todo o vocabulário (computacionalmente caro).
  • Velocidade Top-k: Restringe a soma aos $k$ tokens mais prováveis. A decodificação Top-1 corresponde à decodificação gulosa.
  • Decodificação Estocástica: Amostra um único token da posterior para definir a direção da velocidade.

Inovações Arquiteturais e de Treinamento

1. S-arch (Backbone Esférico): Uma variante do Transformer onde as ativações intermediárias são restritas a $S^{d-1}$. Ela substitui resíduos aditivos por interpolações normalizadas e parametriza as camadas de atenção/MLP como rotações. Isso alinha a arquitetura com a variedade de entrada, melhorando a qualidade das amostras.
2. Truncamento do Cronograma de Ruído: Os autores derivam um limite teórico $\alpha^\star(\delta)$ baseado na geometria de alta dimensão. Eles mostram que, em altas dimensões, a posterior colapsa para uma distribuição one-hot assim que o nível de ruído cai abaixo de um certo limiar. O treinamento é truncado para níveis de ruído acima desse limite para evitar aprender a partir de tarefas de denoising triviais.
3. Cronograma de Ruído Adaptativo: Inspirado pelo InfoNoise, o cronograma de treinamento é adaptado para alocar mais amostras aos níveis de ruído onde o gradiente de perda é mais íngreme, melhorando a eficiência do treinamento.

Contribuições Principais

1. Estrutura S-FLM: Um modelo de linguagem baseado em fluxo que opera em embeddings de $d$ dimensões em uma hipersfera, eliminando a necessidade de materializar vetores one-hot de dimensão $|V|$. Isso reduz os custos de treinamento para serem comparáveis aos modelos de difusão discreta.
2. Backbone S-arch: Uma arquitetura inovadora que mantém as ativações na hipersfera, resultando em melhor desempenho em tarefas de raciocínio (GSM8K) e modelagem de linguagem (OpenWebText) em comparação com backbones DiT padrão.
3. Truncamento Teórico de Ruído: Uma heurística principial para truncar o cronograma de ruído com base no tamanho do vocabulário e na dimensão do embedding, o que melhora significativamente o desempenho em tarefas de raciocínio.
4. Ganhos de Desempenho: O S-FLM melhora substancialmente os modelos de fluxo contínuo anteriores. No GSM8K, alcança ~18% de precisão com decodificação top-1 (comparado a <1% para FLMs anteriores), fechando a lacuna para modelos de difusão mascarada (MDLM, Duo) na temperatura padrão ($T=1$).

Resultados Experimentais

• Raciocínio de Sudoku: O S-FLM alcança desempenho comparável a FLMs anteriores e modelos de difusão discreta (Duo), com os melhores resultados de difusão contínua atingindo 94,8% de precisão em quebra-cabeças fáceis e 45,0% em quebra-cabeças difíceis ao usar truncamento e cronogramas adaptativos.
• GSM8K (Raciocínio Matemático):
  • Modelos contínuos anteriores (FLM, CANDI) alcançaram <1% de precisão.
  • O S-FLM com decodificação top-1 atinge 18,0% de precisão, igualando o desempenho do MDLM e do Duo em $T=1$.
  • Uma lacuna de desempenho permanece em temperatura baixa ($T=0,1$), onde modelos discretos (Duo, MDLM) atingem 33–36%, enquanto o S-FLM estabiliza em torno de 18%.
• OpenWebText (Modelagem de Linguagem): O S-FLM iguala a fronteira de Perplexidade Generativa (Gen. PPL) / Entropia de FLMs anteriores em alto Número de Avaliações de Função (NFE) e supera-os em baixo NFE, onde modelos contínuos anteriores exibiam fronteiras instáveis.
• Eficiência de Treinamento: O S-FLM treina mais rápido que FLMs (que usam vetores one-hot) e é comparável em velocidade aos modelos de difusão discreta.

Significado e Alegações

O artigo afirma que o S-FLM representa um passo significativo para frente para modelos de linguagem de fluxo contínuo, abordando as questões de escalabilidade e interpretação semântica das abordagens anteriores. Ao operar na hipersfera, o modelo evita o gargalo computacional de vetores one-hot de alta dimensão e fornece um mecanismo de injeção de ruído geometricamente coerente.

Os autores observam que, embora o S-FLM feche a lacuna para modelos de difusão discreta e AR em termos de Perplexidade Generativa e precisão de raciocínio em temperaturas padrão, uma lacuna substancial permanece na decodificação em temperatura baixa. O trabalho é apresentado como um avanço metodológico que torna os modelos de fluxo contínuo mais viáveis para modelagem de linguagem com vocabulário grande, embora os autores permaneçam modestos sobre a lacuna de desempenho restante em comparação com modelos autoregressivos de última geração em domínios verificáveis.