Resumo Técnico: Nemotron-TwoTower

Problema

Modelos de linguagem autorregressivos (AR) dominam a geração de texto, mas sofrem com um gargalo de vazão (throughput) devido ao seu processo de decodificação sequencial, um token por vez. Modelos de linguagem de difusão discreta oferecem uma alternativa promissora ao gerar tokens em paralelo e refiná-los iterativamente, potencialmente proporcionando maior vazão e um controle mais granular. No entanto, as abordagens de difusão existentes geralmente dependem de uma única rede para desempenhar dois papéis distintos simultaneamente: representar os tokens de contexto limpos e realizar o de-noising iterativo dos tokens corrompidos. Esse emaranhamento força um único conjunto de pesos a servir objetivos conflitantes, limitando a capacidade do modelo de se destacar em qualquer um dos papéis. Além disso, permanece incerto se o desacoplamento total desses papéis é viável em larga escala, particularmente em arquiteturas híbridas modernas que combinam camadas Mamba, atenção e Mixture-of-Experts (MoE).

Metodologia

Os autores propõem o TwoTower, um modelo de difusão autorregressivo por blocos que desacopla os papéis de representação de contexto e de de-noising em duas torres separadas, construídas sobre a arquitetura híbrida Nemotron-3-Nano-30B-A3B (que intercala camadas Mamba-2, self-attention e MoE).

Arquitetura

O modelo utiliza duas cópias da rede pré-treinada com papéis complementares:

1. Torre de Contexto AR Congelada: Esta torre permanece congelada e atua como um modelo autorregressivo causal sobre o prompt e os tokens já comprometidos. Ela processa tokens limpos para produzir caches de Chave-Valor (KV) alinhados por camada e estados Mamba. Ela preserva as capacidades autorregressivas originais do backbone e mantém estados reutilizáveis entre os blocos.
2. Torre de De-noiser de Difusão Treinável: Esta torre é treinada para refinar blocos ruidosos. Ela emprega atenção de bloco bidirecional dentro do bloco ruidoso atual (permitindo que os tokens atentem uns aos outros), enquanto mantém uma atenção cruzada causal às camadas correspondentes da torre de contexto congelada. Essa atenção cruzada alinhada por camada fornece ao de-noiser acesso multiescala às representações do backbone, contrastando com métodos anteriores que apenas transmitiam o estado oculto final.

Processo de Geração

A geração ocorre bloco a bloco (Algoritmo 1):

• A torre de contexto processa o prompt e os tokens comprometidos causalmente.
• Um novo bloco é inicializado com $S$ tokens mascarados.
• A torre de de-noiser refina iterativamente este bloco ao longo de $T$ passos. Em cada camada, o de-noiser atenta para o cache KV da torre de contexto para blocos passados e para seu próprio cache KV bidirecional para o bloco atual.
• Uma vez que um bloco é suficientemente refinado (com base em limiares de confiança), ele é comprometido. A torre de contexto então atualiza seus estados com os novos tokens limpos, e o processo se repete para o próximo bloco.

Treinamento e Condicionamento

• Difusão Mascarada: O de-noiser é treinado usando um objetivo de difusão mascarada. Blocos limpos são corrompidos substituindo tokens por [MASK] de acordo com um cronograma linear. O modelo minimiza a log-verossimilhança negativa dos tokens originais nas posições mascaradas.
• Condicionamento de Tempo: O de-noiser é condicionado ao passo de tempo de difusão $t$ usando normalização de camada adaptativa (adaLN), uma técnica comum em difusão de imagem, mas menos utilizada em modelos de linguagem mascarados. Isso adiciona parâmetros mínimos (~1,5M), mas melhora a qualidade da amostragem.
• Currículo: O treinamento segue um currículo de dois estágios. A Fase 1 usa uma mistura de dados diversa para cobertura ampla, enquanto a Fase 2 aumenta o peso de fontes de alta qualidade e focadas em STEM. O tamanho do bloco de treinamento é reduzido gradualmente (de $S=32$ para $S=16$) para equilibrar o paralelismo com a qualidade do condicionamento.
• Amostragem: Os autores empregam uma estratégia de desmascaramento por confiança (Algoritmo 2). Em cada passo, o de-noiser prevê todos os tokens mascarados em paralelo; tokens com confiança acima de um limiar $\gamma$ são comprometidos, enquanto tokens incertos permanecem mascarados para refinamento posterior. Isso torna o número de tokens comprometidos adaptativo.

Principais Resultados

Avaliado no backbone Nemotron-3-Nano-30B-A3B (30B parâmetros, MoE híbrido), o Nemotron-TwoTower demonstra o seguinte:

• Troca Qualidade-Vazão: O modelo retém 98,7% da qualidade agregada de benchmarks da linha de base autorregressiva, enquanto alcança uma melhoria de 2,42× na vazão de geração em tempo real (wall-clock).
• Desempenho em Benchmarks:
  • Conhecimento Geral: O desempenho permanece dentro de ~1 ponto da linha de base AR.
  • Código e Matemática: Apresenta modesta degradação em comparação com a linha de base.
  • Senso Comum e Multilinguismo: O desempenho é recuperado ou ligeiramente melhorado.
• Ablações de Design:
  • Mamba Bidirecional: Testar uma variante Mamba bidirecional mostrou ganhos marginais de qualidade, mas dobrou o computo de SSM; o design final mantém o Mamba causal.
  • Condicionamento de Tempo: O condicionamento explícito de tempo melhorou significativamente as pontuações de geração, código e matemática.
  • Desacoplamento de Torres: Manter a torre de contexto congelada e treinar um de-noiser separado produziu os melhores resultados. Vincular os pesos entre as torres ou continuar o treinamento AR na torre de contexto levou a quedas significativas de precisão.
• Sensibilidade do Tamanho do Bloco: Tamanhos de bloco de treinamento menores ($S=16$) melhoraram a qualidade em relação a tamanhos maiores ($S=128$), mas reduziram a vazão. Os autores selecionaram $S=16$ como padrão para equilibrar a retenção de qualidade com a velocidade. A amostragem com blocos maiores que o bloco de treinamento resultou em severa degradação de desempenho.

Significância e Alegações

O artigo afirma que o Nemotron-TwoTower demonstra com sucesso que representações autorregressivas pré-treinadas podem suportar um de-noiser de difusão treinável sem a necessidade de re-treinar todo o modelo. Ao desacoplar os papéis de contexto e de-noising, a abordagem preserva as amplas capacidades do modelo híbrido MoE de 30B original, ao mesmo tempo em que desbloqueia os benefícios de geração paralela da difusão.

Os autores enfatizam que essa recuperação de qualidade é alcançada através da adaptação em uma fração dos dados usados para o pré-treinamento original (~2,1 trilhões de tokens vs. 25 trilhões de tokens). O trabalho estabelece uma nova fronteira de Pareto para modelos de linguagem, oferecendo um caminho viável para uma geração de maior vazão sem sacrificar a qualidade fundamental dos modelos autorregressivos de larga escala. O código e os pesos do modelo são liberados para facilitar mais pesquisas nesta direção.