Resumo Técnico: Sumi – Modelo de Linguagem de Difusão Uniforme Aberto Treinado do Zero

Declaração do Problema

Embora os modelos de difusão tenham emergido como uma alternativa promissora aos modelos de linguagem autorregressivos (AR), existe uma lacuna significativa no desenvolvimento de Modelos de Linguagem de Difusão Uniforme (UDLMs) em escala. Diferente dos Modelos de Linguagem de Difusão Mascarada (MDLMs), que preenchem tokens mascarados de forma irreversível, os UDLMs permitem que qualquer token seja atualizado em qualquer etapa, possibilitando teoricamente uma geração e autocorreção mais flexíveis.

No entanto, nenhum UDLM foi pré-treinado do zero tanto em escalas de parâmetros grandes (bilhões de parâmetros) quanto em grandes orçamentos de tokens (trilhões de tokens). Os UDLMs existentes em larga escala são ou checkpoints computacionalmente otimizados treinados com pequenos orçamentos de tokens, ou modelos adaptados de arquiteturas AR pré-treinadas (ex: DiffusionGemma). Consequentemente, o comportamento da difusão uniforme nativa em larga escala no regime rico em dados permanece inexplorado, carecendo de um ponto de referência limpo para estudar leis de escala, dinâmicas de geração e compensações (trade-offs) contra baselines estabelecidos de AR e MDLM.

Metodologia

Arquitetura do Modelo e Objetivo de Treinamento

Os autores apresentam o Sumi, um UDLM de 7 bilhões de parâmetros totalmente aberto.

• Arquitetura: O Sumi utiliza um Transformer bidirecional independente de tempo com 36 camadas, tamanho oculto de 4.096 e MLPs SwiGLU. Emprega atenção de consulta agrupada (grouped-query attention) (32 cabeças, 8 grupos de KV), RMSNorm e embeddings de entrada/saída não vinculados (untied). Para mitigar sumidouros de atenção (attention sinks), um softmax "off-by-one" é aplicado na atenção.
• Objetivo: O modelo é treinado usando o objetivo de Difusão Discreta de Interpolação Generalizada (GIDD), especificamente em sua forma reparametrizada por Razão Sinal-Ruído (SNR).
• Esquema de Ruído: O treinamento utiliza ruído uniforme puro com o log-SNR restrito a $\lambda \in [-9, 9]$.
• Tokenização: É utilizado o tokenizador OLMo 3 com um vocabulário de 100.278.

Regime de Treinamento

O Sumi foi pré-treinado do zero em 1,5 trilhão de tokens usando 288 GPUs NVIDIA H100 (totalizando ~43.308 horas-GPU).

1. Pré-treinamento (1,3T tokens): Realizado em uma sequência de comprimento 1.184 com um tamanho de lote global de 4.608. A mistura de dados é fortemente ponderada para conteúdo educacional, filtrado e reclassificado por pontuações educacionais derivadas de um classificador Qwen3-32B destilado.
2. Treinamento Intermediário (250B tokens): Duas etapas seguem o pré-treinamento:
   • 130B tokens no comprimento original (1.184).
   • 120B tokens em um comprimento estendido (4.864).
   • A mistura foca em código, matemática e raciocínio, compreendendo 32,5% de dados de codificação, 29,7% de matemática, 21,0% de geral e 16,8% de raciocínio.

Estratégia de Inferência

• Comprimento da Tela (Canvas Length): O modelo utiliza um comprimento de tela fixo de 2.048 tokens para todas as avaliações. Este comprimento foi escolhido porque o desempenho degrada significativamente em telas substancialmente menores ou maiores.
• Amostragem: Os autores exploram amostragem adaptativa baseada em confiança versus amostragem ancestral.
• Protocolo de Avaliação: Todas as avaliações utilizam o lm-evaluation-harness modificado para pontuação de difusão. Para tarefas de verossimilhança (likelihood), as posições além da resposta são preenchidas com tokens aleatórios; para tarefas de geração, as atualizações são restritas ao comprimento da geração.

Principais Contribuições

1. Primeiro UDLM Nativo de Larga Escala: O Sumi é o primeiro UDLM pré-treinado do zero em uma grande escala de parâmetros (7B) e grande orçamento de tokens (1,5T), fornecendo uma base para o estudo da difusão uniforme nativa.
2. Lançamento Aberto: Os autores liberam os pesos do modelo, checkpoints e a receita completa de treinamento, incluindo uma especificação completa da mistura de dados sobre corpora publicamente disponíveis.
3. Sondas de Inferência Exploratórias: O artigo apresenta análises direcionais sobre dinâmicas de geração, incluindo sensibilidade ao comprimento da tela, ordem de comprometimento de tokens, capacidades de decodificação paralela e potencial de autocorreção.

Resultados

Desempenho em Benchmarks

O Sumi foi avaliado em 13 benchmarks em quatro categorias e comparado com modelos AR abertos (Falcon-7B, Llama 2-7B, OLMo-7B) treinados em orçamentos de tokens comparáveis.

• Conhecimento Geral e Codificação: O Sumi apresenta desempenho competitivo, alcançando as melhores pontuações entre os modelos avaliados sob o mesmo protocolo em MMLU, TruthfulQA, HumanEval e MBPP. Isso é atribuído à mistura de dados pesada em educação e código.
• Raciocínio e Matemática: O Sumi é competitivo com o Llama 2-7B e supera o Falcon-7B em tarefas como GSM8K e ARC, embora fique atrás de modelos maiores ou treinados de forma diferente (ex: LLaDA-8B, Llama 3-8B).
• Senso Comum: O Sumi tem um desempenho significativamente inferior em benchmarks de senso comum (PIQA, HellaSwag, WinoGrande). Os autores atribuem essa lacuna à mistura de dados pesada em educação, que tem sido observada por degradar o desempenho de senso comum, embora notem que a lacuna é provavelmente maior do que a composição de dados pode explicar sozinha.

Observações de Inferência

• Sensibilidade da Tela: A fluência da geração é altamente sensível ao comprimento da tela. O Sumi gera de forma fluente dentro de uma banda dependente da tarefa (centrada em 2.048), mas degrada bruscamente fora desse intervalo, particularmente em telas curtas.
• Ordem de Comprometimento: Embora o objetivo de treinamento seja independente de ordem, a amostragem baseada em confiança induz uma ordem de comprometimento auto-organizada. O modelo tende a comprometer tokens determináveis cedo e adiar os incertos, uma estrutura não presente na amostragem ancestral.
• Decodificação Paralela: Em tarefas de codificação (HumanEval, MBPP), o modelo suporta uma modesta decodificação paralela (comprometendo até 4 tokens por etapa) com perda mínima de precisão. No entanto, aritmética de múltiplos passos (GSM8K) permanece sensível à ordem, degradando imediatamente com passos paralelos.
• Autocorreção: Um orçamento de revisão explícito (permitindo passagens de denoising extras) não produz autocorreção. Embora o modelo sobrescreva tokens comprometidos, essas edições são predominantamente "viagens de ida e volta" ($A \to B \to A$) em vez de melhorias direcionadas, deixando a precisão final inalterada.

Significância e Alegações

O artigo posiciona o Sumi como um recurso fundamental para a comunidade estudar a difusão uniforme nativa em escala. Os autores enfatizam que suas sondas de inferência são "direcionais e não conclusivas", destinadas a motivar pesquisas futuras em aspectos pouco compreendidos dos UDLMs, tais como:

• Os comportamentos intrínsecos de geração da difusão uniforme versus a difusão mascarada.
• As condições sob as quais a autocorreção pode ser alcançada.
• O impacto da composição de dados nas forças e fraquezas específicas dos modelos de difusão.

Os autores declaram explicitamente que o Sumi é lançado como um modelo base pré-treinado sem ajuste de instrução (instruction tuning) ou filtragem de segurança, destinado a apoiar a pesquisa em vez de implantação direta. Eles alertam que o modelo herda riscos padrão associados a modelos de linguagem base, incluindo potencial para saídas prejudiciais e imprecisões fatuais.