Resumo Técnico: Dual-Flow Transformers

Declaração do Problema

O custo de servir Grandes Modelos de Linguagem (LLMs) é cada vez mais dominado pela inferência, em vez do treinamento. No entanto, as duas fases primárias da inferência — prefill (processamento do prompt) e decodificação autorregressiva (geração de tokens) — exibem gargalos de hardware fundamentalmente diferentes:

• Prefill: O processamento paralelo de tokens do prompt é tipicamente limitado por computação (compute-bound), restrito pelo throughput aritmético.
• Decode: A geração sequencial é tipicamente limitada por largura de banda de memória (memory-bandwidth-bound), restrita pela necessidade de ler repetidamente o cache Key-Value (KV) crescente e os pesos do modelo para cada novo token.

Estratégias de escalonamento convencionais (aumento da largura ou profundidade do modelo) elevam os custos para ambas as fases simultaneamente, pois cada camada adicionada é avaliada durante o prefill e o decode. Isso cria uma rigidez onde melhorar a qualidade da geração (que frequentemente se beneficia de computação adicional) aumenta inevitavelmente o custo de processamento do prompt. O artigo questiona se é possível alocar computação aprendida adicional especificamente para a predição de continuação (decode) enquanto preserva a computação ampla do prompt e um único cache KV persistente.

Metodologia: O Dual-Flow Transformer

Os autores propõem o Dual-Flow Transformer, uma arquitetura que desacopla estruturalmente o caminho primário de processamento do prompt do caminho auxiliar de predição de continuação.

Arquitetura Central

O modelo consiste em duas trajetórias de estado residual que compartilham a maioria dos parâmetros do modelo, mas diferem em sua gestão de estado e escopo de ativação:

1. Fluxo Primário (Primary Flow): Um caminho Transformer causal completo e padrão. Ele processa todo o prompt, escreve o cache KV persistente e determina o estado primário. É autocontido e não depende do fluxo auxiliar.
2. Fluxo Auxiliar (Auxiliary Flow): Uma trajetória residual que lê o estado primário e o cache, mas nunca escreve no cache KV persistente e nunca retroalimenta o fluxo primário.
   • Ativação: O fluxo auxiliar é omitido durante o processamento inicial do prompt (prefill). Ele é ativado apenas na posição final do prompt (para prever o primeiro token de continuação) e em cada posição gerada subsequente.
   • Compartilhamento de Pesos: Ambos os fluxos compartilham todas as matrizes principais de atenção, MLP e saída.
   • Separação: Eles utilizam embeddings de tokens distintos e vetores de acoplamento leves e aprendidos.

Mecanismos Chave

• Atenção de KV Compartilhada Assimétrica (Asymmetric Shared-KV Attention): O fluxo primário realiza a autoatenção causal padrão. O fluxo auxiliar gera uma nova query ($Q_2$), que é acoplada à query primária ($Q_1$) via um vetor aprendido ($a_\ell$). Crucialmente, o fluxo auxiliar atende às chaves e valores (Keys e Values) do fluxo primário ($K_1, V_1$). Isso permite que o fluxo auxiliar aproveite o estado em cache sem escrever o seu próprio, permitindo a execução agrupada onde pesos e regiões de KV são carregados uma única vez para servir ambos os fluxos.
• Acoplamento Leve (Lightweight Coupling): A informação flui unidirecionalmente do primário para o fluxo auxiliar via vetores aprendidos ($a_\ell, b_\ell, c_\ell$) nos níveis de query, intermediário de MLP e saída de MLP. Isso enriquece a representação auxiliar sem criar dependências que alterariam o caminho primário.
• Objetivo de Mistura (Mixture Objective): A predição final é uma mistura das distribuições primária e auxiliar: $p_{mix} = \alpha_t p_{primary} + (1-\alpha_t) p_{aux}$. O peso de mistura $\alpha_t$ é ajustado dinamicamente com base na confiança (probabilidade de colisão) da distribuição primária, garantindo que o fluxo primário permaneça dominante (pelo menos 50% de peso) durante a inferência.
• Replay de Roteador (para MoE): Em configurações de Mistura de Especialistas (Mixture-of-Experts - MoE), o fluxo auxiliar reutiliza o mesmo conjunto de especialistas selecionados pelo fluxo primário (determinado pelo roteador primário), mas os aplica com seus próprios pesos de mistura. Isso preserva a lógica de seleção de especialistas do fluxo primário enquanto aumenta a aritmética realizada sobre os especialistas.

Custos de Inferência vs. Treinamento

• Inferência: A fase de prefill incorre em nenhum custo adicional em comparação com um Transformer padrão; o fluxo auxiliar é ignorado. A fase de decode incorre em computação adicional (aproximadamente o dobro da aritmética por token), mas não aumenta o tamanho do cache KV nem o conjunto de pesos únicos lidos, pois os dois fluxos são agrupados.
• Treinamento: Ambos os fluxos são avaliados ao longo de toda a sequência, resultando em aproximadamente 2x mais FLOPs de treinamento do que o modelo primário. No entanto, o aumento no número de parâmetros é mínimo (apenas embeddings separados e vetores de acoplamento).

Contribuições Principais

1. Arquitetura: Introdução de um design de KV compartilhado assimétrico que desacopla a computação ampla do prompt da computação de continuação. O fluxo primário lida com a persistência de estado, enquanto o fluxo auxiliar fornece uma trajetória "residual" para uma predição aprimorada sem alterar o estado primário.
2. Qualidade Empírica: Demonstrou reduções consistentes na perda de validação em três eixos experimentais:
   • Escalonamento de Dados (NanoGPT): O Dual-Flow supera os Transformers padrão em orçamentos de tokens equivalentes, com um piso de perda ajustado menor.
   • Escalonamento de Tamanho de Modelo (Dense LLaMA): Os ganhos persistem conforme o tamanho do modelo aumenta.
   • MoE Esparsa: Validado em configurações MoE do tipo LLaMA esparsa com replay de roteador.
3. Alocação de MoE Específica por Fase: Diferente de designs de fluxo paralelo anteriores que aplicam alocação de computação uniforme, o Dual-Flow permite que os fan-outs de especialistas primário e auxiliar sejam configurados de forma independente. Isso permite um espaço de trade-off entre custo de prefill, custo de decode e qualidade preditiva, permitindo otimizar sistemas para cargas de trabalho específicas (ex: fixar o custo de prefill enquanto aumenta o compute de decode).

Resultados

• Perda de Validação: Em comparações de orçamento de tokens iguais nos experimentos NanoGPT, o Dual-Flow alcança consistentemente uma perda de validação menor que os Transformers padrão. O piso de perda assintótica ajustada diminuiu de 2.9416 (Transformer) para 2.9013 (Dual-Flow).
• Eficiência de Computação: Em uma comparação de computação de treinamento pareada (onde um modelo Dual-Flow com metade dos passos de dados realiza aproximadamente os mesmos FLOPs de backbone compartilhado que um Transformer com o dobro de passos), o Dual-Flow ainda alcançou menor perda de validação, sugerindo que alocar computação entre fluxos interagentes é mais eficaz do que simplesmente aumentar os dados para um único fluxo.
• Ablações: A análise de componentes mostrou que a construção mínima de Dual (embeddings separados, KV compartilhado) já supera a arte anterior mais próxima (PHD-2). Adicionar o acoplamento entre fluxos e a mistura de probabilidade melhorou ainda mais o desempenho.

Significância e Alegações

O artigo alega que o Dual-Flow estabelece uma oportunidade estrutural para alocação de especialistas específica por fase. Ao desacoplar o caminho do prompt do caminho de continuação, a arquitetura permite que sistemas de serviço tratem a computação de tempo de geração como um recurso escalável sem aumentar proporcionalmente o custo do processamento do prompt.

Especificamente, os autores destacam que, em modelos MoE, os fluxos primário e auxiliar podem ter fan-outs de especialistas independentes. Isso cria um controle de três vias:

1. Custo de Prefill: Determinado pelo uso de especialistas do fluxo primário.
2. Custo de Decode: Determinado pela soma do uso de especialistas primários e auxiliares.
3. Qualidade Preditiva: Melhorada ao aumentar a computação total alocada para a fase de continuação.

Isso permite flexibilidade de implantação específica, onde um sistema pode preservar os custos de computação do prompt enquanto gasta o orçamento de decode disponível em mais especialistas auxiliares, ou vice-versa, sem alterar o pool de parâmetros compartilhado ou a estrutura de cache primário único. O trabalho apoia a premissa de que a computação do próximo token é um recurso de escalonamento útil que pode ser alocado internamente dentro do modelo, em vez de externamente via rastros de raciocínio mais longos ou múltiplas amostragens.