Resumo Técnico: BlockBatch

Declaração do Problema

Os Modelos de Linguagem de Difusão (dLLMs) geram texto denoising iterativamente múltiplas posições de tokens em paralelo, oferecendo uma alternativa à decodificação estritamente autoregressiva. No entanto, a inferência eficiente enfrenta um compromisso crítico de granularidade. Os métodos existentes de inferência por blocos (por exemplo, Fast-dLLM) selecionam um único tamanho de bloco por solicitação. Blocos pequenos preservam a condicionamento local, mas exigem muitas etapas de denoising (alto NFE), enquanto blocos grandes aumentam o paralelismo, mas arriscam compromissos prematuros e erros acumulados de cache (deriva do KV-cache). Os métodos atuais de aceleração tratam o tamanho do bloco como um hiperparâmetro fixo, falhando em explorar a natureza complementar de diferentes tamanhos de bloco para diferentes partes da mesma geração.

Metodologia: BlockBatch

Os autores propõem o BlockBatch, um framework de inferência online sem treinamento que trata o tamanho do bloco como uma dimensão de ramificação, e não como um parâmetro fixo. Em vez de selecionar um único tamanho de bloco, o BlockBatch executa múltiplas ramificações com diferentes tamanhos de bloco (por exemplo, 4, 8, 16, 32, 64, 128) para a mesma solicitação dentro de uma única passagem forward em lote.

O framework depende de três mecanismos principais para coordenar essas ramificações:

1. Fusão Controlada por Confiança: As ramificações trocam propostas de tokens apenas quando são compatíveis (isto é, seus prefixos decodificados correspondem exatamente). Uma ramificação de destino aceita um token proposto por uma ramificação de origem apenas se a própria distribuição de probabilidade da ramificação de destino atribuir a esse token uma pontuação de confiança acima de um limiar fixo. Isso permite que ramificações atrasadas progridam sem forçar compromissos de tokens inconsistentes.
2. Sincronização Baseada em Líder: Para evitar que ramificações atrasadas desperdiçem computação em regiões desatualizadas do espaço do KV-cache, aplica-se uma regra de sincronização. Se uma ramificação líder avançar além de uma ramificação atrasada por um número específico de tokens acima de um limiar, a ramificação atrasada copia o estado da sequência e o KV-cache do líder, efetivamente redefinindo sua trajetória para o caminho confiável mais avançado.
3. Atualizações Periódicas de Sequência Completa: O denoising local de blocos atualiza apenas um subconjunto da sequência, o que pode levar a uma deriva acumulada em relação ao estado globalmente consistente do KV-cache. O BlockBatch recalcula periodicamente o KV-cache para toda a sequência das ramificações ativas. Isso ancora novamente as ramificações a um estado globalmente consistente, corrigindo erros acumulados durante atualizações locais.

Otimização do Sistema: Para garantir eficiência, o BlockBatch empacota as posições de consulta não mascaradas de todas as ramificações ativas em um único buffer de comprimento variável. Isso permite que o modelo processe múltiplas ramificações em uma única passagem forward, utilizando o kernel varlen do FlashAttention, garantindo que a latência do denoising de múltiplas ramificações seja comparável à de uma única ramificação.

Caracterizações Principais

O artigo fornece caracterizações teóricas e empíricas de por que essa abordagem funciona:

• Trajetória do KV-Cache: Diferentes tamanhos de bloco induzem trajetórias diferentes no espaço do KV-cache. Blocos pequenos seguem caminhos conservadores com condicionamento de granularidade fina, enquanto blocos grandes fazem progresso local agressivo. Atualizações de sequência completa atuam como correções globais que impedem que essas trajetórias se afastem excessivamente.
• Dinâmica em Nível de Token: As ramificações frequentemente compartilham um prefixo inicial, bifurcam-se em posições semanticamente decisivas e, posteriormente, voltam a concordar em tokens sintaticamente leves (por exemplo, espaços em branco, operadores, dígitos). O framework aproveita isso fundindo ramificações quando concordam em tokens de alta confiança e sincronizando-as quando divergem significativamente.

Resultados Experimentais

O BlockBatch foi avaliado em três dLLMs representativos (LLaDA-1.5-8B, LLaDA-Instruct-8B, Dream-Base-7B) em quatro benchmarks (GSM8K, MATH, HumanEval, MBPP).

• Eficiência: O BlockBatch reduz o número médio de avaliações de função de denoising (NFEs) em 26,6% em comparação com o Fast-dLLM (a linha de base mais avançada).
• Velocidade: Alcança uma aceleração end-to-end média de 1,33× sobre o Fast-dLLM, com acelerações chegando a 5,8× em configurações específicas (por exemplo, LLaDA-Instruct-8B no MBPP).
• Precisão: Apesar da redução agressiva nos NFEs e na latência, o BlockBatch preserva uma precisão comparável ou superior às linhas de base em todas as tarefas.
• Comparação: Supera consistentemente a decodificação de difusão Vanilla, o Fast-dLLM e o LocalLeap em termos de eficiência de NFE e latência em todas as doze configurações de modelo-tarefa testadas.

Significado e Alegações

O artigo alega que a diversidade de tamanhos de bloco é um eixo prático e anteriormente pouco explorado para inferência de dLLM em ramificações paralelas. Ao demonstrar que diferentes tamanhos de bloco fornecem compensações de precisão-eficiência complementares, os autores argumentam que tratar o tamanho do bloco como uma dimensão de ramificação dinâmica permite uma decodificação mais eficiente do que qualquer configuração fixa única.

O significado do BlockBatch reside em sua capacidade de:

1. Explorar a estrutura inerente da inferência de dLLM, onde diferentes granularidades são úteis para diferentes partes de uma geração.
2. Coordenar múltiplas trajetórias de decodificação sem treinamento, usando apenas operações em tempo de inferência (fusão, sincronização, atualização).
3. Fornecer um framework geral que melhora a eficiência enquanto mantém a qualidade do texto gerado, identificando a diversidade de tamanhos de bloco como uma estratégia viável para acelerar modelos de linguagem baseados em difusão.