Resumo Técnico: Difusão de Blocos Triplet RWKV (B3D-RWKV)

Declaração do Problema

Modelos de linguagem Transformer causais enfrentam duas limitações estruturais primárias: a decodificação estritamente sequencial, que impede a paralelização, e os custos de atenção quadráticos por etapa, que tornam a inferência de contexto longo computacionalmente dispendiosa. Embora modelos causais de tempo linear (por exemplo, RWKV) e modelos de difusão discreta abordem essas fraquezas independentemente, sua integração tem sido inerentemente inconsistente. Modelos de difusão requerem atenção bidirecional para denoizar blocos de tokens em paralelo, ao passo que modelos causais são estritamente unidirecionais. Tentativas existentes de combiná-los frequentemente exigem modificações arquitetônicas que comprometem a eficiência do backbone de tempo linear ou necessitam de treinamento do zero.

Metodologia

Os autores propõem B3D-RWKV, uma variante RWKV no estilo de difusão que unifica a eficiência de inferência $O(L)$ com difusão discreta bidirecional e paralela, sem alterar o backbone causal subjacente. A inovação central é o Layout de Bloco Triplet, uma estratégia de treinamento e inferência projetada para simular acesso bidirecional dentro de uma arquitetura estritamente causal.

1. Layout de Bloco Triplet (Treinamento)

Para habilitar a difusão em um modelo causal, cada bloco lógico de geração de tamanho $B$ é expandido em três blocos físicos contíguos dentro da amostra de treinamento:

• $b_1$ (Cópia Mascada): Uma versão mascarada do bloco lógico.
• $b_2$ (Cópia Mascada Passível de Perda): Uma cópia mascarada idêntica onde a perda de denoização é computada.
• $b_3$ (Cópia Limpa de Verdade Terrena): Uma versão não mascarada que atualiza o estado recorrente para o bloco subsequente.

Acesso Pseudo-Bidirecional:
Como o backbone causal processa tokens estritamente da esquerda para a direita, o estado oculto que chega a qualquer posição mascarada em $b_2$ já processou a totalidade de $b_1$. Como $b_1$ e $b_2$ compartilham o mesmo padrão de máscara, os tokens não mascarados em $b_1$ são efetivamente "absorvidos" no estado oculto antes que o modelo tente prever os tokens mascarados em $b_2$. Isso concede ao modelo acesso pseudo-bidirecional ao seu próprio contexto não mascarado dentro do bloco lógico, satisfazendo os requisitos bidirecionais da difusão enquanto mantém a natureza unidirecional do backbone RNN/SSM.

Objetivo de Treinamento:
O modelo é treinado usando um objetivo de previsão de máscara. A perda consiste em:

• Perda de Entropia Cruzada ($L_{CE}$): Computada em posições supervisionadas onde os tokens estão mascarados.
• Perda Paralela Consciente de Confiança ($L_{CAP}$): Seguindo o LLaDA-2.0, este termo afina a distribuição de probabilidade em posições já previstas corretamente, incentivando o amostrador de inferência a comprometer mais tokens por etapa.

2. Inferência: Denoização Iterativa por Blocos

Durante a inferência, o modelo gera um bloco lógico de cada vez. O processo envolve:

1. Inicializar um bloco com todos os tokens mascarados.
2. Executar etapas iterativas de denoização onde o modelo prevê probabilidades top-1 para posições mascaradas.
3. Comprometer tokens cuja probabilidade top-1 excede um limiar $\tau$.
4. Um mecanismo de fallback garante progresso comprometendo os $k_{min}$ tokens mais confiantes se menos de $k_{min}$ ultrapassarem o limiar.
5. Uma vez que todas as posições são comprometidas, o bloco limpo é anexado ao prefixo e o próximo bloco lógico começa.

Principais Contribuições

1. B3D-RWKV-7.2B: O primeiro modelo de linguagem RNN de tempo linear no estilo de difusão treinado na escala de 7B com um objetivo de previsão de máscara. Ele integra a seleção paralela de tokens no backbone RWKV-7 sem modificar seus parâmetros originais.
2. Framework de Bloco Triplet: Uma receita de treinamento universal que converte backbones estritamente causais (satisfazendo causalidade estrita e propriedades de estado de propagação para frente) em modelos de difusão por blocos. Esta abordagem evita a complexidade arquitetônica de modificar o backbone em um bloco bidirecional.
3. Eficiência e Desempenho: O método alcança melhorias significativas no rendimento (throughput) em comparação com a decodificação causal padrão, mantendo capacidades de raciocínio competitivas.

Resultados Experimentais

O modelo foi avaliado em um conjunto de 8 tarefas, incluindo conhecimento geral (MMLU, ARC, PIQA, RACE) e tarefas de raciocínio/matemática (GSM8K, MATH, GPQA).

• Precisão: O B3D-RWKV-7.2B alcança precisão comparável à linha de base RWKV-7 e a outros modelos de difusão de escala similar (por exemplo, LLaDA-8B, Dream-7B). Ele supera notavelmente as linhas de base em tarefas de raciocínio como ARC-Challenge e RACE, atribuído à percepção pseudo-bidirecional que aprimora o raciocínio.
• Rendimento (Throughput): O modelo demonstra uma aceleração média de 1,6× no rendimento de decodificação em comparação com a linha de base RWKV-7 em comprimentos de contexto variando de 1K a 512K tokens. Com parâmetros de amostragem ajustados, uma aceleração de 2,02× é alcançável com uma pequena troca de qualidade.
• Raciocínio Matemático: O artigo nota uma leve queda no desempenho em tarefas matemáticas complexas (por exemplo, MATH) em comparação com a linha de base causal. Os autores atribuem isso à natureza da decodificação paralela, que luta com a verificação altamente estruturada, passo a passo, necessária para problemas de matemática simbólica onde crédito parcial não é concedido.

Significado e Alegações

O artigo alega oferecer um caminho eficiente para transformar modelos de linguagem causais pré-treinados em modelos de linguagem de difusão sem mudanças arquitetônicas. Ao aproveitar o layout de bloco triplet, o B3D-RWKV unifica com sucesso a eficiência de inferência $O(L)$ de modelos de tempo linear com as capacidades de decodificação paralela da difusão.

Os autores posicionam o trabalho como uma solução estrutural para a incompatibilidade entre os paradigmas causal e de difusão. Eles reconhecem limitações, incluindo a validação atual em um único backbone (RWKV-7), o custo de comprimento de sequência físico 3× durante o treinamento e o uso de um conjunto de dados SFT relativamente pequeno sem alinhamento por RL, o que pode contribuir para regressões de precisão em tarefas específicas. No entanto, a contribuição central permanece a demonstração de que geração de difusão paralela de alto rendimento é alcançável em backbones estritamente causais e de tempo linear através de layout de dados, em vez de reengenharia arquitetônica.