💬 NLP

Prefilling-dLLM: Predictive Prefilling for Long-Context Inference in Diffusion Language Models

O artigo apresenta o Prefilling-dLLM, um framework livre de treinamento que acelera a inferência de contexto longo em modelos de linguagem de difusão ao cachear e selecionar esparsamente chunks de prefixo relevantes para reduzir a complexidade computacional de quadrática em relação ao comprimento total da sequência para quadrática em relação ao comprimento de decodificação, alcançando assim acelerações de estado da arte e mitigando o fenômeno lost-in-the-middle.

Autores originais: Jing Xiong, Qi Han, Shansan Gong, Yunta Hsieh, Chengyue Wu, Chaofan Tao, Chenyang Zhao, Ngai Wong

Publicado 2026-06-10

📖 5 min de leitura🧠 Leitura aprofundada

CC BY 4.0

Autores originais: Jing Xiong, Qi Han, Shansan Gong, Yunta Hsieh, Chengyue Wu, Chaofan Tao, Chenyang Zhao, Ngai Wong

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ê está tentando resolver um quebra-cabeça gigante, mas em vez de olhar para a imagem inteira de uma vez, você tem que reexaminar a caixa inteira de peças toda vez que coloca uma única peça nova na mesa.

Isso é essencialmente como os Modelos de Linguagem de Difusão (dLLMs) funcionam atualmente ao lidar com histórias ou documentos longos. Cada vez que o modelo tenta gerar uma nova palavra, ele relê e reprocessa todo o histórico da conversa desde o início. À medida que a história fica mais longa, esse "reler" torna-se tão lento e caro que é como tentar correr uma maratona carregando uma mochila pesada que fica mais pesada a cada passo.

O artigo apresenta um novo método chamado Prefilling-dLLM para corrigir isso. Veja como funciona, usando algumas analogias do cotidiano:

1. O Problema: A Armadilha do "Re-Ler"

Em modelos de IA tradicionais, uma vez que você lê uma página de um livro, você a memoriza e foca apenas na nova página que está lendo. Mas nesses modelos de difusão, o computador esquece o livro e relê o livro inteiro, da página 1 à página 1.000, toda vez que escreve uma única palavra nova.

O Resultado: Se o livro for curto, tudo bem. Se o livro tiver 32.000 páginas, o computador passará 99% do tempo apenas relendo as páginas antigas e apenas 1% realmente escrevendo novas.

2. A Solução: O Sistema do "Bibliotecário Inteligente"

Os autores propõem um sistema chamado Prefilling-dLLM que atua como um bibliotecário super eficiente. Em vez de reler toda a biblioteca a cada vez, o bibliotecário faz duas coisas:

Passo A: O "Prefill" (Organizando a Biblioteca)

Antes de começar a escrita, o bibliotecário pega o texto de entrada massivo (o "prefixo") e o divide em pedaços (como capítulos de um livro).

O Truque: O bibliotecário lê cada capítulo uma vez, cria um "cartão de resumo" (chamado de cache KV) para ele e o coloca em uma prateleira.
A Inovação: O bibliotecário não lê cada palavra individual de cada capítulo. Eles usam um truque especial para manter apenas as frases mais importantes de cada capítulo, descartando o que é irrelevante. Isso torna os cartões de resumo muito menores e mais rápidos de manipular.

Passo B: O "Decode" (Escrevendo a História)

Agora, quando o modelo precisa escrever a próxima palavra, ele não relê a biblioteca inteira.

A Seleção: O modelo pergunta: "Quais capítulos são realmente relevantes para o que estou escrevendo agora?" Ele usa um sistema de pontuação inteligente para escolher apenas os principais capítulos (os pedaços mais relevantes) da prateleira.
A Eficiência: Ele ignora os outros 90% da biblioteca que não são necessários para aquela frase específica. Ele olha apenas para os "cartões de resumo" relevantes que criou anteriormente.

3. Por que isso é um divisor de águas

O artigo afirma que esta abordagem é um aumento de velocidade massivo porque:

Sem Mais Re-Leitura: O trabalho pesado de ler o texto longo acontece uma vez no início.
Pulo Inteligente: Durante a fase de escrita real, o modelo olha apenas para uma fração minúscula do texto (os principais pedaços relevantes), em vez de todo o conteúdo.
A Velocidade: Em contextos longos (de 8.000 a 32.000 palavras), este método é de 9 a 28 vezes mais rápido do que métodos anteriores, mantendo a precisão das respostas.

4. Resolvendo o Mistério do "Perdido no Meio"

Existe um problema conhecido em modelos de IA de contexto longo chamado "Lost in the Middle" (Perdido no Meio). Imagine uma pessoa lendo uma longa lista de fatos; ela lembra perfeitamente dos primeiros e dos últimos, mas esquece completamente o que está no meio.

A Correção do Artigo: Os autores descobriram que, ao dividir o texto em pedaços e adicionar um "token de âncora" especial (como um título de capítulo) ao início de cada pedaço, o modelo consegue encontrar informações em qualquer lugar do texto, mesmo no meio. É como ter um índice que funciona perfeitamente, para que o modelo nunca se perca, não importa o quão longa seja a história.

5. O Equilíbrio do "Tamanho do Pedaço"

O artigo também testou o quão grandes esses "capítulos" deveriam ser.

Muito Grandes: Se os pedaços forem enormes, o modelo pode perder detalhes importantes no meio do texto.
Muito Pequenos: Se os pedaços forem minúsculos, o computador gasta tempo demais gerenciando a lista de pedaços.
O Ponto Ideal: Eles descobriram que, para textos muito longos, usar muitos pedaços menores funciona melhor para manter a precisão alta.

Resumo

Prefilling-dLLM é como mudar de um estudante que relê um livro de 500 páginas toda vez que responde a uma pergunta de um quiz, para um estudante que cria um guia de estudo condensado e inteligente uma vez, e depois consulta apenas as páginas específicas de que precisa ao responder.

O resultado? A IA pode lidar com conversas e documentos muito mais longos sem ficar lenta, e não esquece os detalhes importantes escondidos no meio do texto.

Resumo Técnico: Prefilling-dLLM

Definição do Problema

Os Modelos de Linguagem de Grande Escala de Difusão (dLLMs) oferecem uma alternativa promissora aos modelos autorregressivos (AR) ao permitir a geração paralela de tokens por meio de denoising iterativo. No entanto, os dLLMs enfrentam uma ineficiência crítica em cenários de contexto longo: em cada etapa de denoising, o modelo deve reencodar todo o prefixo de entrada para computar a autoatenção. Isso resulta em uma complexidade computacional que escala quadraticamente com o comprimento total da sequência ( $O((L_p + L_d)^2 \cdot T)$ , onde $L_p$ é o comprimento do prefixo, $L_d$ é o comprimento de decodificação e $T$ é o número de etapas de denoising). Diferente dos modelos AR, onde a fase de prefill é limitada por computação e a fase de decodificação é limitada por memória, a inferência do dLLM permanece limitada por computação durante todas as iterações porque a sequência completa (prefixo + resposta) é processada conjuntamente em cada etapa. Essa recomputação repetida do prefixo torna-se proibitiva para aplicações de contexto longo (ex: 8K–32K+ tokens).

Metodologia: PREFILLING-DLLM

Os autores propõem o PREFILLING-DLLM, um framework de desagregação de prefill-decode livre de treinamento, projetado para eliminar a recomputação do prefixo e introduzir esparsidade. O método opera em três estágios:

1. Prefill em Blocos (Chunked Prefill) e Cache de KV

Em vez de processar todo o prefixo como um único bloco, o prefixo de entrada é particionado em $N$ blocos não sobrepostos de tamanho fixo $C$ .

Processamento Independente: Cada bloco é processado independentemente com atenção bidirecional intra-bloco. Um token especial de Início de Sequência (BOS) é concatenado ao início de cada bloco para atuar como um delimitador.
Pontuação Preditiva: Para determinar quais blocos são relevantes sem alvos de verdade fundamental (ground-truth), o framework gera rótulos pseudo via uma passagem de difusão preliminar. Ele então pontua cada bloco usando dois sinais:
- Pontuação de Auto-Informação ( $s_I$ ): A log-verossimilhança negativa da janela de consulta condicionada ao bloco.
- Pontuação de Pseudo-Rótulo ( $s_P$ ): A verossimilhança dos pseudo-rótulos condicionada ao bloco e à consulta.
Seleção: Os top- $K$ blocos mais relevantes são selecionados com base na pontuação combinada ( $s = s_I + s_P$ ). As representações KV desses blocos selecionados são armazenadas em cache.

2. Esparsidade de Dois Níveis

O framework introduz esparsidade tanto no nível de bloco quanto no nível de token para reduzir a complexidade:

Esparsidade Inter-bloco: Apenas os top- $K$ blocos participam do mecanismo de atenção durante a decodificação, em vez do conjunto completo de $N$ blocos.
Esparsidade Intra-bloco: Dentro de cada bloco selecionado, os tokens são classificados por seu peso de atenção bidirecional cumulativo com os tokens de consulta. Apenas os top- $B$ tokens mais relevantes por bloco são retidos no cache KV, descartando tokens irrelevantes.

3. Decodificação com Reuso de KV Estático

Durante o processo de denoising iterativo:

As representações KV em cache dos $K$ blocos selecionados permanecem estáticas e são reutilizadas em todas as $T$ etapas de denoising.
O modelo apenas recomputa as representações KV para a consulta e os tokens de resposta em evolução.
Isso reduz a complexidade por etapa de $O((L_p + L_d)^2 \cdot T)$ para $O(N \cdot C^2 + (L_d^2 + K \cdot C) \cdot T)$ , efetivamente desacoplando o custo de decodificação do comprimento total do prefixo.

Otimização de Kernel de Atenção

Para lidar com os gargalos de memória causados pelo carregamento de caches KV de blocos não contíguos, os autores implementam um kernel Split-S FlexAttention. Este kernel paraleliza a decodificação sobre os KVs de blocos armazenados de forma não contígua, particionando a dimensão KV e processando os blocos independentemente antes de mesclar os resultados via redução log-sum-exp, evitando a necessidade de reunir os dados em memória contígua.

Principais Descobertas e Resultados

Desempenho e Eficiência

Aceleração: No LongBench e InfiniteBench, o PREFILLING-DLLM alcança um ganho de velocidade de 9,1× a 28,0× em comprimentos de contexto de 8K a 32K em comparação com baselines de Transformers padrão.
Qualidade: O método alcança qualidade de estado da arte entre métodos de aceleração de dLLM. No LongBench, atinge uma pontuação média de 34,59 (Dream-7B) e 37,02 (UltraLLaDA), superando baselines como Fast-dLLM, Sparse-dLLM e dKV-Cache.
Escalabilidade de Contexto Longo: No InfiniteBench (contextos >128K), o método alcança 43,62% de precisão média no Dream-7B, superando significativamente o baseline mais forte (Fast-dLLM v2 com 30,32%). Notavelmente, alcança 95,42% de precisão na tarefa de recuperação de Passkey e 70,00% na recuperação de Números.

Abordagem do "Lost-in-the-Middle"

Fenômeno: O artigo observa que os dLLMs exibem um fenômeno de "perda no meio" (lost-in-the-middle) sob extrapolação de contexto, onde a precisão degrada para informações localizadas no meio do contexto, mostrando um viés de decaimento monotônico em direção à resposta, em vez da curva em forma de U vista em modelos AR.
Mitigação: Ao selecionar blocos baseados em relevância preditiva em vez de proximidade posicional, o PREFILLING-DLLM elimina esse viés. O método mantém alta precisão de correspondência exata em todas as posições do contexto, mesmo em comprimentos de 32K e 128K.

Insights de Mecanismo

Âncoras de Atenção: Os tokens BOS concatenados a cada bloco atuam como âncoras de atenção periódicas. A análise mostra que esses tokens distribuem a massa de atenção uniformemente pelo contexto (formando cristas estáveis de baixa magnitude), em vez de degenerar em "sumidouros de atenção" (attention sinks) que absorvem massa de atenção desproporcional. Esse mecanismo permite a recuperação invariante à posição.
Trade-off de Tamanho de Bloco: Existe um compromisso entre o tamanho do bloco e a precisão da recuperação. Blocos menores (ex: 256 tokens) mantêm alta precisão mesmo em contextos de 128K, enquanto blocos maiores (ex: 4096 tokens) degradam bruscamente além de 32K. A seleção de múltiplos blocos é essencial para a inferência de contexto longo.

Significância e Alegações

O artigo afirma que o PREFILLING-dLLM preenche a lacuna entre a natureza limitada por computação da inferência de dLLM e a eficiência das arquiteturas de desagregação de prefill-decode usadas em modelos AR. Suas principais contribuições são:

Desagregação Livre de Treinamento: Introduz um framework que armazena em cache as representações KV do prefixo uma única vez e as reutiliza, eliminando o custo de recomputação quadrática do prefixo em cada etapa de denoising sem exigir o retreinamento do modelo.
Prefill Esparso: Demonstra que o prefill esparso (seleção de blocos e tokens relevantes) pode superar a atenção densa em dLLMs, reduzindo a complexidade enquanto melhora ou mantém a qualidade de geração.
Otimização de Sistema: Fornece um kernel de atenção especializado (Split-S FlexAttention) que aproveita a estrutura em blocos do prefixo para alcançar reduções significativas de latência (até 10,2× sobre o vanilla FlexAttention) durante a fase de decodificação.
Robustez a Contextos Longos: Mostra que esta abordagem mitiga o fenômeno "lost-in-the-middle" e permite que os dLLMs escalem efetivamente para 128K+ tokens com alta precisão de recuperação, uma capacidade anteriormente inexplorada na aceleração de dLLM.

Os autores observam limitações, incluindo a natureza estática da seleção de blocos (sem reseleção dinâmica durante a decodificação), a necessidade de ajuste de tamanho de bloco e $K$ para tarefas específicas, e o fato de a avaliação atual ser limitada a benchmarks em inglês em arquiteturas de modelos específicas (Dream-7B e UltraLLaDA).