MARS: Enabling Autoregressive Models Multi-Token Generation

1. O Problema

Os modelos de linguagem autoregressivos (AR) atuais geram texto um token por vez, realizando uma passagem frontal (forward pass) para cada token, independentemente da previsibilidade do conteúdo. Isso resulta em um custo computacional uniforme e ineficiente, especialmente para sequências altamente previsíveis (como fórmulas ou frases comuns).

As abordagens existentes para geração multi-token apresentam desvantagens significativas:

• Decodificação Especulativa: Requer a manutenção de um modelo rascunho separado, dobrando a memória e adicionando complexidade de orquestração.
• Abordagens Multi-cabeça (ex: Medusa, EAGLE): Adicionam cabeças de previsão extras à arquitetura, exigindo parâmetros adicionais e treinamento específico.
• Difusão em Blocos (Block Diffusion): Tenta prever múltiplos tokens simultaneamente, mas frequentemente causa degradação significativa na qualidade, especialmente em tarefas de raciocínio e codificação, devido a incompatibilidades com a atenção causal original do modelo AR.

2. Metodologia: MARS (Mask AutoRegreSsion)

O MARS propõe um método de ajuste fino (fine-tuning) leve que ensina um modelo AR já ajustado para instruções a prever múltiplos tokens por passagem frontal, sem modificar a arquitetura ou adicionar parâmetros. O modelo resultante é um "superconjunto estrito" do original: pode operar como um AR padrão ou acelerar a geração quando confiante.

Principais Componentes Técnicos:

• Análise de Lacunas (Gap Analysis): Os autores identificaram quatro lacunas entre a previsão AR e a previsão mascarada em blocos. Três delas são escolhas de design elimináveis (padrão de atenção bidirecional, alinhamento de logits e ordem de geração), e apenas uma é inerente (o uso de tokens [MASK]). O MARS fecha as três lacunas elimináveis para manter a compatibilidade AR.
• Atenção Causal e Logits Deslocados: Diferente da difusão em blocos que usa atenção bidirecional dentro dos blocos, o MARS mantém atenção causal estrita (esquerda para direita) e logits deslocados para a direita, preservando a função de cabeça de linguagem original.
• Estratégia de Treinamento Duplo:
  • O modelo processa duas cópias da sequência em paralelo: um fluxo "limpo" (tokens originais) e um fluxo "ruidoso" (blocos substituídos por [MASK]).
  • Perda Combinada: O treinamento utiliza uma perda de entropia cruzada para os tokens mascarados (L_mask) mais a perda padrão de previsão do próximo token no fluxo limpo (L_AR).
  • Importância da Perda SFT: A perda no fluxo limpo é crucial. Sem ela, o sinal de treinamento AR decai à medida que o tamanho do bloco aumenta, degradando o desempenho em raciocínio. A perda combinada mantém o sinal AR acima de 50%, independentemente do tamanho do bloco.
• Inferência com Janela Deslizante: Durante a inferência, o modelo anexa B tokens [MASK] ao prefixo atual. Ele gera logits para todos eles e aceita tokens consecutivamente da esquerda para a direita enquanto a confiança (max p(xt)) estiver acima de um limiar (τ). Se a confiança for baixa, ele recai para a geração de um único token.

3. Contribuições Principais

1. Análise Teórica: Demonstração de que a degradação de qualidade em modelos de difusão em blocos não é inerente à previsão multi-token, mas sim resultado de escolhas de design (atenção bidirecional, etc.) que podem ser corrigidas.
2. Método MARS: Uma técnica de ajuste fino que permite geração multi-token adaptativa sem novas cabeças ou parâmetros, usando apenas os dados de SFT existentes.
3. Mecanismo de Preservação de AR: Identificação da perda SFT no fluxo limpo como o ingrediente chave para manter a competência autoregressiva em blocos grandes.
4. Otimização de Inferência: Desenvolvimento de uma estratégia de cache KV em nível de bloco para inferência em lote, permitindo aceleração real de tempo de execução.
5. Controle em Tempo Real: O sistema permite ajustar o limiar de confiança (τ) sob demanda durante o serviço, oferecendo um "botão" prático para equilibrar latência e qualidade sem reiniciar o modelo.

4. Resultados

Os experimentos foram realizados nos modelos Qwen2.5-0.5B e Qwen2.5-7B em seis benchmarks padrão (IFEval, BBH, MMLU-Pro, GPQA, GSM8K, HumanEval).

• Qualidade (Modo 1 Token): Quando configurado para gerar 1 token por passo (τ=1.0), o MARS igualou ou superou a linha de base AR SFT em todos os benchmarks, com ganhos notáveis em GSM8K (+4.5 pontos no 7B) e HumanEval.
• Aceleração (Modo Multi-Token): Ao permitir a aceitação de múltiplos tokens, o MARS alcançou um aumento de 1.5x a 1.7x no throughput (tokens por forward) mantendo a precisão próxima à linha de base.
  • No modelo 7B, a precisão média caiu apenas 1.3 pontos (de 58.1 para 56.8) ao aumentar o throughput médio para 1.68 tokens/forward.
• Aceleração de Tempo de Parede (Wall-Clock): Com a estratégia de cache KV em blocos, o MARS alcançou um speedup de até 1.71x em relação ao AR com cache KV em lotes pequenos (batch size 4) no Qwen2.5-7B. Sem o cache de bloco, o desempenho degradava devido à recomputação da sequência completa.
• Robustez: A perda SFT no fluxo limpo provou ser essencial; sem ela, o aumento do tamanho do bloco (de 4 para 16) causou uma queda drástica de 6.2 pontos na média de precisão.

5. Significado e Impacto

O MARS representa uma mudança de paradigma na aceleração de modelos de linguagem:

• Simplicidade de Implantação: Elimina a necessidade de modelos rascunhos complexos ou modificações arquiteturais, permitindo que o mesmo checkpoint sirva tanto para máxima qualidade quanto para máxima velocidade.
• Eficiência de Recursos: Oferece aceleração significativa sem o custo de memória de manter múltiplos modelos ou cabeças extras.
• Flexibilidade Operacional: A capacidade de ajustar o limiar de confiança (τ) dinamicamente durante o serviço permite que sistemas de produção otimizem a latência sob alta carga sem retrainar ou trocar modelos.
• Validação de Causalidade: Confirma que a geração multi-token pode ser totalmente compatível com a atenção causal estrita, desafiando a suposição anterior de que a difusão bidirecional era necessária para paralelismo.

Em resumo, o MARS transforma modelos AR padrão em modelos "híbridos" adaptativos que podem acelerar a geração de forma segura e controlável, resolvendo o dilema entre qualidade e velocidade de inferência com uma solução elegante e leve.