Resumo Técnico: N-vium – Transformer de Mistura de Saídas para Geração Exata Acelerada

1. Declaração do Problema

Os Modelos de Linguagem Grandes (LLMs) autorregressivos padrão processam tokens sequencialmente através de uma pilha de camadas de decodificador. Esta passagem forward serve a duas funções acopladas: (1) construir representações Chave-Valor (KV) em cada camada para atenção futura, e (2) projetar o estado oculto final para prever o próximo token. Consequentemente, cada token requer uma passagem forward completa através de todas as camadas, independentemente de a distribuição de previsão já ter convergido.

Abordagens existentes para melhorar a eficiência de inferência, como métodos de Saída Antecipada (Early Exit), tentam desacoplar essas funções ao anexar cabeças de previsão em camadas intermediárias. No entanto, esses métodos geralmente tratam previsões intermediárias como aproximações da distribuição da camada final. Isso leva a várias desvantagens:

• Erros de Aproximação: Previsões intermediárias são frequentemente mais fracas do que as de profundidade total, criando uma incompatibilidade entre treino e teste.
• Roteamento Heurístico: Decisões sobre quando sair frequentemente dependem de limites de confiança criados manualmente que exigem calibração.
• Caches KV Incompletos: Para economizar computação, métodos de saída antecipada frequentemente ignoram o cálculo de caches KV para camadas contornadas, exigindo cópia aproximada ou pulo, o que degrada a qualidade.
• Degradação de Qualidade: A velocidade é frequentemente alcançada à custa da perplexidade e da qualidade de geração.

O artigo argumenta que, embora observações empíricas (por exemplo, via "lente de logit") mostrem que as distribuições de previsão frequentemente se estabilizam antes da camada final, os métodos atuais falham em aproveitar isso para geração exata sem perda de qualidade.

2. Metodologia: N-vium

Os autores introduzem o N-vium, um framework de mistura de saídas que generaliza estritamente a arquitetura padrão de transformer. Em vez de tratar saídas antecipadas como aproximações, o N-vium redefine a saída do modelo como uma distribuição de mistura aprendida sobre múltiplas saídas.

Arquitetura

• Mistura de Saídas: O modelo particiona as $L$ camadas da base em $N_V$ blocos. No final de cada bloco (junção $k$), um roteador determina a probabilidade de amostrar o próximo token a partir da distribuição $\pi_k$ daquela camada.
• Roteamento Adaptativo ao Token: Um MLP roteador processa o estado oculto $h^{(\ell_k)}_t$ na junção $k$ para produzir uma probabilidade $w_k$. A probabilidade $p_k$ de sair na junção $k$ é calculada como:
  $$p_k(x_t) = w_k(h^{(\ell_k)}_t) \prod_{j=1}^{k-1} (1 - w_j(h^{(\ell_j)}_t))$$
  onde $w_{N_V} = 1$ por convenção.
• Distribuição de Mistura: A distribuição final do próximo token é uma soma ponderada de todas as distribuições de saída:
  $$\pi_{mix}(x_{t+1} | x_{1:t}) = \sum_{k=1}^{N_V} p_k(x_t) \pi_k(x_{t+1} | x_{1:t})$$
• Cabeça Compartilhada: Todas as saídas compartilham a cabeça de projeção final de Modelagem de Linguagem (LM) $W_{lm}$, com MLPs adaptadores leves transformando estados ocultos intermediários para corresponder à geometria esperada da cabeça.

Treinamento

• Objetivo: O modelo é treinado para minimizar a perda de entropia cruzada da distribuição de mistura conjunta $\pi_{mix}$, em vez de otimizar cabeças individuais isoladamente para corresponder à camada final.
• Penalidade de Computação: Um termo de perda auxiliar $\mathcal{L}_{compute}$ é adicionado para incentivar o roteador a atribuir massa de probabilidade a saídas mais antigas, direcionando o modelo para a eficiência:
  $$\mathcal{L}_{total} = \mathcal{L}_{mix} + \beta \mathcal{L}_{compute}$$
  onde $\mathcal{L}_{compute}$ é a profundidade de camada normalizada esperada.
• Aquecimento (Warm-up): Uma fase de aquecimento do roteador é empregada para evitar o "colapso do roteador" (onde todos os tokens são roteados para a camada final) penalizando desvios de uma distribuição de saída uniforme inicialmente.

Inferência: Aproveitamento (Piggybacking)

Uma inovação crítica é a técnica de Aproveitamento (Piggybacking) para garantir recuperação exata do cache KV sem pular computação:

1. Computação Adiada: Quando um token sai na junção $k$, a computação para as camadas $k+1$ a $L$ é adiada.
2. Agrupamento Paralelo: À medida que tokens futuros são processados, suas passagens forward atingem as camadas adiadas. Nessas profundidades, os estados ocultos dos tokens adiados são concatenados com o estado do token atual e processados em paralelo.
3. Resultado: Isso preenche os estados KV ausentes para todos os tokens adiados em uma única passagem paralela, garantindo que o cache KV esteja completo, exato e consistente com o token amostrado. Nenhuma computação é pulada; ela é meramente reordenada de execução sequencial para execução em lote paralela.

3. Contribuições Principais

1. Amostragem Exata de Mistura: O N-vium é o primeiro método a permitir roteamento adaptativo de profundidade por token com amostragem exata da distribuição treinada, evitando os erros de aproximação inerentes aos métodos padrão de saída antecipada.
2. Recuperação Completa de KV: Ao adiar e agrupar computações de camadas superiores, o N-vium recupera caches KV completos e exatos, uma característica frequentemente perdida em outros métodos de computação adaptativa.
3. Framework de Treinamento: Os autores fornecem um pipeline de treinamento prático que otimiza o objetivo de mistura diretamente, incluindo estratégias para direcionar modelos para a eficiência via uma penalidade de computação aprendível.
4. Escalabilidade: O método é validado em escalas de até 1,5 bilhão de parâmetros, demonstrando que a abordagem não se limita a modelos pequenos.

4. Resultados Experimentais

Os autores pré-treinararam modelos N-vium (especificamente "Quadrivium" com 4 saídas) no conjunto de dados C4 usando a arquitetura LLaMA 2, comparando-os com baselines densas correspondentes em parâmetros e dados.

• Aceleração vs. Perplexidade:
  • O maior modelo (1,5 bilhão de parâmetros, 48 camadas, largura oculta 1536) alcançou uma aceleração de 57,9% no tempo de parede sobre a baseline densa com nenhuma degradação de perplexidade (0% de diferença).
  • Modelos menores mostraram acelerações consistentes variando de 23% a 41% com pequenas melhorias de perplexidade (por exemplo, um modelo de 36 camadas mostrou uma aceleração de 36,4% com uma melhoria de 3,7% na perplexidade).
• Fronteira de Pareto: Em gráficos de qualidade-velocidade, o N-vium ocupa o quadrante "superior-direito" (mais rápido e qualidade igual ou melhor), enquanto métodos como CALM (reimplementado como CALM*) caem no "inferior-direito" (mais rápido, mas qualidade pior) e o LayerSkip fica em uma linha de perplexidade constante abaixo da baseline.
• Desempenho Downstream: Após ajuste fino supervisionado (SFT) na mistura Tulu 3, os modelos N-vium igualaram ou superaram as baselines densas em benchmarks downstream (MMLU, ARC, HellaSwag, etc.), confirmando que os modelos permanecem direcionáveis e não sofrem overfitting no objetivo de entropia cruzada às custas do desempenho da tarefa.
• Robustez de Hardware: Medições de aceleração em GPUs A100, RTX A6000 e RTX 3090 produziram resultados consistentes (~57%), enquanto GPUs mais antigas com arquitetura Turing mostraram ganhos menores devido à falta de suporte nativo a bfloat16.

5. Significado e Alegações

O artigo afirma que o N-vium representa uma mudança de paradigma na eficiência de inferência. Ao contrário da maioria dos métodos que melhoram tokens por segundo reduzindo FLOPs por token (frequentemente via aproximação), o N-vium melhora a eficiência aumentando FLOPs por segundo. Ele alcança isso reorganizando a computação sequencial em passagens em lote paralelas, aproveitando melhor o paralelismo de hardware sem descartar qualquer computação.

Os autores enfatizam que o N-vium é uma generalização estrita do transformer padrão: quando os pesos de roteamento para saídas intermediárias são zero, o modelo se comporta exatamente como um transformer padrão. O método satisfaz quatro critérios críticos simultaneamente, o que os autores afirmam que nenhum outro método existente faz:

1. Inferência Exata: Amostragem da distribuição treinada sem aproximação.
2. Recuperação Completa de Cache KV: Nenhuma camada pulada ou preenchimento aproximado de cache.
3. Roteamento Aprendido: A computação por token é alocada por um roteador treinado, não por heurísticas.
4. Aceleração no Tempo de Decodificação: Reduz a latência no tempo de parede durante a geração autorregressiva.

O artigo conclui que, embora o treinamento do N-vium incorra em um custo maior de FLOPs por token (devido à avaliação de cabeças em múltiplas junções), essa sobrecarga é justificada para implantação em produção onde os custos de inferência dominam, e a sobrecarga relativa diminui à medida que o tamanho do modelo escala.