💬 NLP

Training Hybrid Block Diffusion Language Models with Partial Bidirectionality

Este artigo apresenta o BDLM Mamba-H, um modelo de linguagem de difusão híbrido que restringe a varredura Mamba bidirecional aos blocos de denoising ativos para permitir o cache exato, alcançando assim uma perplexidade superior e um throughput de inferência de contexto longo significativamente maior em comparação com as linhas de base de difusão de sequência total e baseadas em atenção existentes.

Autores originais: Pranshu Chaturvedi, Parth Shroff, Tarun Suresh, Hangoo Kang, Kaiyue Wen

Publicado 2026-07-07

📖 4 min de leitura☕ Leitura rápida

CC BY 4.0

Autores originais: Pranshu Chaturvedi, Parth Shroff, Tarun Suresh, Hangoo Kang, Kaiyue Wen

Artigo original sob licença CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta é uma explicação gerada por IA do artigo abaixo. Não foi escrita nem endossada pelos autores. Para precisão técnica, consulte o artigo original. Ler aviso legal completo

Imagine que você esteja tentando escrever uma história muito longa, mas tem uma regra estrita: você só consegue se lembrar das últimas poucas páginas que escreveu. Toda vez que quer escrever a próxima frase, você tem que voltar até o início do seu caderno, ler cada palavra que escreveu até agora e, só então, decidir o que escrever em seguida.

Isso é exatamente como os atuais "Large Language Models" (chatbots de IA) lutam com contextos longos. À medida que a conversa fica mais longa, a IA tem que "reler" todo o histórico toda vez que gera uma única nova palavra. Isso é lento e desperdiça muita energia, como tentar correr uma maratona carregando uma mochila pesada que fica mais pesada a cada passo.

O artigo apresenta uma nova maneira de treinar esses modelos de IA, chamada BDLM Mamaba-H, que resolve esse problema mudando a forma como a IA "lembra" e "pensa".

Aqui está a divisão usando analogias simples:

1. O Problema: A "Mochila Pesada"

Os modelos de IA atuais funcionam como um aluno fazendo uma prova que precisa reler o livro inteiro antes de responder a cada pergunta.

O Problema: À medida que a história fica mais longa, a "mochila" (memória) fica muito pesada. O computador passa mais tempo movimentando dados (largura de banda de memória) do que realmente pensando (computação).
A Solução Antiga: Alguns pesquisadores tentaram tornar a mochila mais leve usando um tipo diferente de memória (chamado "Mamba"). Outros tentaram escrever várias frases de uma vez em vez de uma por uma (chamado "Block Diffusion").
A Falha: Quando os pesquisadores tentaram combinar essas duas ideias, isso quebrou. O sistema de memória "Mamba" precisava olhar para a história inteira de trás para frente para funcionar adequadamente. Mas, se você olha para a história inteira de trás para frente, não consegue salvar um "marcador" para as partes que já terminou. Você tem que reler tudo toda vez.

2. A Solução: O "Escaneamento Reverso Local"

Os autores propõem um truque inteligente: Olhar para trás apenas dentro do parágrafo atual.

Imagine que você está escrevendo um livro capítulo por capítulo.

O Jeito Antigo (Reverso Total): Para escrever uma nova frase, você relê o livro inteiro da última página até a página um, toda vez.
O Novo Jeito (BDLM Mamba-H):
1. A Parte "Limpa": Assim que um capítulo é finalizado, você o coloca em um cofre seguro. Você cria um "cartão de resumo" simples (um cache) para esse capítulo. Você nunca mais precisa abrir o cofre; basta usar o cartão de resumo.
2. A Parte "Ativa": Enquanto você está escrevendo o capítulo atual, você tem permissão para olhar para trás e para frente dentro desse capítulo específico para garantir que as frases fluam bem.
3. A Magia: Como você só olha para trás dentro do capítulo atual, os cartões de resumo dos capítulos finalizados permanecem válidos e não precisam ser atualizados.

3. Os Resultados: Velocidade e Inteligência

Os pesquisadores testaram este novo método contra métodos antigos com dois objetivos principais: manter a IA inteligente e torná-la rápida.

É inteligente? Sim. Quando testaram a IA em um teste padrão de compreensão de leitura (C4-en), o novo modelo (BDLM Mamba-H) na verdade obteve as melhores pontuações entre os modelos menores (87 milhões de parâmetros). Mesmo quando tornaram o modelo maior (350 milhões de parâmetros), ele permaneceu tão inteligente quanto os outros modelos de topo.
É rápido? Diferença enorme.
- Em um comprimento médio (65.000 palavras), o novo modelo foi 19,7 vezes mais rápido que o antigo método de "reverso total".
- Em um comprimento muito longo (262.000 palavras), foi 3,7 vezes mais rápido que o melhor método de "bloco" que não utilizava Mamba.

A Conclusão

Pense neste novo modelo como um escritor que aprendeu a marcar capítulos finalizados para não ter que relê-los, enquanto ainda consegue editar o capítulo atual livremente.

Ao restringir o "olhar para trás" apenas ao bloco atual de texto, a IA pode gerar histórias muito longas sem ficar sobrecarregada pelo tráfego de memória. O artigo afirma que isso torna esta uma arquitetura prática e poderosa para lidar com contextos longos, provando que você pode ter tanto velocidade quanto escrita de alta qualidade sem precisar reprocessar todo o histórico para cada nova palavra.

Resumo Técnico: Treinamento de Modelos de Linguagem de Difusão de Blocos Híbridos com Bidirecionalidade Parcial

Definição do Problema
A geração de contexto longo de alto rendimento em modelos de linguagem de grande escala (LLMs) é limitada primariamente pela largura de banda de memória, e não pelo processamento. Na decodificação autorregressiva padrão, cada etapa exige o streaming do cache de chave/valor (KV) acumulado da memória, fazendo com que a demanda de largura de banda escale linearmente com o comprimento do contexto enquanto emite apenas um token. Embora existam duas abordagens paralelas para mitigar isso — reduzindo o acesso à memória via mixers de tempo linear (ex: Mamba) e aumentando o paralelismo via geração por blocos — combinar ambos apresenta um obstáculo técnico específico.

Modelos de difusão híbridos anteriores, como o DiffuMamba, utilizam varreduras (scans) Mamba totalmente bidirecionais, incluindo uma passagem de direção reversa sobre toda a janela de denoising. Esta varredura reversa depende de tokens ativos, o que significa que seus estados internos não são apenas de prefixo. Consequentemente, esses estados não podem ser precisamente cacheados e reutilizados entre blocos, anulando os ganhos de eficiência da difusão por blocos. O desafio central é projetar uma arquitetura híbrida que suporte a geração por blocos com cache exato, mantendo simultaneamente os benefícios da mistura bidirecional.

Metodologia
Os autores propõem o BDLM Mamba-H, uma arquitetura híbrida que combina Modelos de Linguagem de Difusão por Blocos (BDLMs) com camadas Mamba e de atenção. A inovação central é um mecanismo de bidirecionalidade parcial que restringe a varredura Mamba de direção reversa apenas ao bloco de denoising ativo, em vez de toda a sequência.

Componentes metodológicos principais incluem:

Objetivo de Treinamento Alinhado ao Cache: O modelo emprega um objetivo "all-block" onde os blocos de prefixo completados são armazenados como:
- Estados KV para camadas de atenção.
- Estados de fronteira de apenas avanço (forward-only) para camadas Mamba (estados recorrentes).
  A varredura reversa é confinada ao bloco ativo atual ( $b_t$ ), garantindo que o cache do prefixo ( $C_\ell(c)$ ) permaneça uma função apenas do prefixo limpo, independente do passo de difusão reversa ( $t$ ).
Fatoração de Passo de Tempo (Timestep Factorization): Para manter a validade do cache, a arquitetura fatia o cálculo da camada em um operador de prefixo livre de passo de tempo e uma continuação do bloco ativo condicionada ao passo de tempo. Isso evita que o cache do prefixo limpo se torne dependente do passo de tempo específico da difusão reversa, um requisito para a reutilização válida durante a inferência.
Protocolo de Treinamento: O modelo é treinado usando um objetivo de difusão discreta sobre blocos. O cache de prefixo é construído através da varredura do prefixo limpo, e as perdas dos blocos subsequentes treinam diretamente as representações armazenadas nesses estados de prefixo, alinhando o cálculo de treinamento com a interface de cache do amostrador em tempo de inferência.

Principais Contribuições

Treinamento Alinhado ao Cache: Os autores formulam um objetivo BDLM "all-block" onde os objetos em cache (estados KV e estados de prefixo Mamba) correspondem exatamente aos caches de bloco em tempo de inferência. Isso permite que as perdas dos blocos subsequentes treinem a varredura de prefixo usada para a reutilização de estado recorrente.
Validação Controlada e Escalonamento: Os autores realizam um levantamento abrangente com 87 milhões de parâmetros no dataset DCLM e escalam modelos selecionados para 350 milhões de parâmetros usando uma regra de transferência de taxa de aprendizado do tipo Complete(d)P.
Ganhos de Rendimento em Contexto Longo: O estudo avalia o rendimento (throughput) de inferência em comprimentos de contexto estendidos, demonstrando que a abordagem híbrida supera significativamente tanto os baselines de sequência completa quanto os modelos de difusão de blocos baseados puramente em atenção conforme o comprimento do contexto aumenta.

Resultos

Qualidade de Validação (87M de Parâmetros): O BDLM Mamba-H alcançou a melhor perplexidade de validação C4-en (61.6) entre os modelos de 87M, superando o BDLM atenção (76.5), o DiffuMamba-H de sequência completa (83.7) e a atenção de sequência completa (87.6).
Qualidade de Validação (350M de Parâmetros): Na escala de 350M, o BDLM Mamba-H permaneceu competitivo com o BDLM atenção, alcançando uma perplexidade C4-en de 38.2 comparado a 36.97 do BDLM atenção.
Desempenho de Rendimento (Throughput):
- Em 65K tokens, o BDLM Mamba-H alcançou 19.7x o rendimento do baseline de sequência completa DiffuMamba-H.
- Em 262K tokens, o BDLM Mamba-H alcançou 3.7x o rendimento do BDLM atenção.
- Embora os modelos de sequência completa sejam mais rápidos em contextos curtos (8192 tokens) devido ao overhead do objetivo "all-block", os modelos BDLM nativos com caching demonstram uma escalabilidade superior em regimes de contexto longo.

Significância e Alegações
O artigo afirma que o BDLM Mamba-H representa uma arquitetura viável para difusão de contexto longo. Ao restringir a varredura Mamba reversa ao bloco ativo, o modelo permite o cache exato entre blocos, resolvendo o problema de reutilização de estado inerente aos híbridos de difusão bidirecional anteriores. Os autores afirmam que esta abordagem permite que modelos híbridos Mamba-atenção mantenam uma qualidade de validação competitiva enquanto oferecem melhorias substanciais de rendimento para inferência de contexto longo, posicionando os híbridos de parte-reversa por blocos como uma direção prática para escalar modelos de linguagem de difusão.

O artigo permanece modesto quanto à otimalidade, observando que as regras de transferência de taxa de aprendimento utilizadas para o escalonamento são candidatas fortes baseadas em leis de parametrização existentes, mas não comprovadas como ótimas para modelos de linguagem de difusão discreta (dLLMs). Adicionalmente, os autores reconhecem que, embora o objetivo por blocos melhore a geração de contexto longo, ele atualmente incorre em uma penalidade de rendimento de treinamento (aproximadamente 25–34% mais lento) em comparação com os baselines de sequência completa em configurações de hardware atuais.

1. O Problema: A "Mochila Pesada"

2. A Solução: O "Escaneamento Reverso Local"

3. Os Resultados: Velocidade e Inteligência

A Conclusão

Mais como este