Resumo Técnico: Decodificação de Exploração de Profundidade (DEX)

1. Definição do Problema

A decodificação autorregressiva (AR) em Grandes Modelos de Linguagem (LLMs) é inerentemente dispendiosa porque cada token gerado deve passar por toda a pilha de camadas antes que o próximo token possa ser produzido. Estudos recentes identificaram uma redundância significativa ao longo da dimensão de profundidade: muitos tokens tornam-se previsíveis em camadas intermediárias, o que significa que não requerem estritamente a profundidade total do modelo para serem precisos.

Métodos existentes de adaptação de profundidade sem perdas (lossless) tentam explorar isso empregando uma estratégia de seleção de profundidade. Esses métodos propõem um token de uma profundidade intermediária pré-selecionada e o verificam contra o modelo de profundidade final. No entanto, os autores identificam um gargalo de seleção crítico:

• Saída Tardia (Late Exit): Se a profundidade de saída escolhida for muito profunda, há desperdício de computação em camadas que já eram redundantes.
• Saída Antecipada (Early Exit): Se a saída escolhida for muito rasa e a previsão não coincidir com a saída da profundidade final, o sistema aciona um fallback. Isso não apenas descarta o token atual, mas também invalida quaisquer rascunhos "descendentes" dependentes, levando a perdas significativas de eficiência.

O artigo postula que a lacuna entre os métodos atuais de seleção e um "oráculo" ideal (que conhece a profundidade exata onde um token se torna estável) é substancial. A questão central é se essa redundância pode ser explorada de forma mais eficaz reduzindo o risco de uma decisão errada de profundidade de saída.

2. Metodologia: Decodificação de Exploração de Profundidade (DEX)

Os autores propõem a Decodificação de Exploração de Profundidade (DEX), um algoritmo sem perdas que substitui a seleção de profundidade única por uma exploração paralela sobre múltiplas profundidades candidatas.

Conceitos Fundamentais

• Profundidade de Disponibilidade Precoce (EAD - Earliest Available Depth): Os autores definem EAD como a camada mais rasa onde a previsão de um token coincide com a saída da profundidade final e permanece estável através das camadas subsequentes.
• Conjunto de Exploração ($X$): Em vez de escolher uma saída, o DEX define um conjunto de profundidades candidatas $\{d_0, d_1, \dots, d_{N-1}\}$, onde $d_{N-1}$ é a profundidade final.
• Resolução: O overhead do DEX é limitado pela "resolução" deste conjunto (a lacuna máxima entre qualquer camada e a próxima profundidade explorada). À medida que o conjunto se torna mais fino, o overhead aproxima-se de zero.

O Algoritmo: Expansão–Compromisso–Colapso (Expand–Commit–Collapse)

O DEX opera em um reticulado de profundidade–posição gerenciado por uma Máquina de Estados Finitos (FSM). O processo consiste em três fases:

1. Expansão do Reticulado (Exploração Paralela):

   • O modelo é particionado em $N$ estágios de profundidade, cada um tratado por um "explorador de profundidade" separado (potencialmente em diferentes unidades de hardware).
   • Em cada posição de compromisso (commit), o sistema expande ramos candidatos em paralelo. Um token de exploração pode sofrer expansão de profundidade (passando seu estado oculto para o próximo explorador de camada) ou expansão de posição (usando o candidato atual como um prefixo para gerar a próxima posição de token).
   • Isso cria uma árvore de futuros potenciais, onde diferentes ramos representam previsões feitas em diferentes profundidades.

2. Compromisso (Validação):

   • Uma vez que o "caminho âncora" (o ramo que atravessa todas as camadas) atinge a profundidade final, ele produz o token de referência de profundidade final.
   • O DEX valida todos os tokens candidatos não finais do caminho âncora contra esta referência.
   • Ele seleciona o ramo mais raso cujo token coincida com a referência de profundidade final. Se nenhum match intermediário for encontrado, ele recorre ao ramo de profundidade final.
   • Crucialmente, o token anexado à saída é sempre o produzido pela referência de profundidade final, garantindo equivalência sem perdas à decodificação AR padrão.

3. Colapso (Poda de Estado):

   • Após o compromisso de um token, o sistema colapsa o reticulado. Ramos que são inconsistentes com o novo prefixo comprometido são descartados.
   • Apenas os estados (caches KV, estados ocultos) do ramo selecionado são preservados para o próximo passo.
   • Isso garante que, embora o sistema explore múltiplas profundidades, ele mantenha o estado exato necessário para o próximo passo AR.

Otimizações

• Amostragem Acoplada à Profundidade (DECS): Para evitar que diferentes profundidades no mesmo fluxo de posição proponham o mesmo token (o que desperdiça exploração), o DEX mascara tokens já propostos por profundidades mais rasas ao decodificar candidatos mais profundos. Isso incentiva a diversidade nas propostas intermediárias sem alterar a distribuição de amostragem final.
• Adaptador Indutor: Para LLMs padrão onde as camadas intermediárias não estão naturalmente alinhadas com a cabeça de saída, os autores propõem anexar pequenos adaptadores residuais às profundidades intermediárias. Estes são treinados via autodestilação para alinhar os estados ocultos intermediários com a cabeça LM final, melhorando a "prontidão" dos tokens em profundidades mais rasas.

3. Principais Contribuições

1. Formalização da Prontidão do Token: O artigo define a Profundidade de Disponibilidade Precoce (EAD) para quantificar quando um token se torna previsível. Utiliza isso para demonstrar matematicamente o "gargalo de seleção" nos métodos existentes, mostrando que o overhead deles é impulsionado pelo risco de uma seleção errada, e não apenas pelo custo de computação.
2. O Algoritmo DEX: Um algoritmo de decodificação sem perdas que instancia a exploração de profundidade através de um procedimento de expansão–compromisso–colapso. Ele preserva a equivalência exata à decodificação AR padrão enquanto reduz o custo de compromisso dos tokens.
3. Validação Empírica: Os autores demonstram que o DEX supera as linhas de base de seleção de profundidade representativas (LayerSkip, AdaDecode, DEL) e alcança um throughput competitivo contra métodos de decodificação especulativa e distribuída (Lookahead, PEARL, EAGLE).

4. Resultados Experimentais

Os autores avaliaram o DEX tanto em modelos treinados para saída antecipada (LayerSkip-treinados CodeLlama-34B e Llama-2-70B) quanto em LLMs padrão (CodeLlama-34B, Llama-2-70B, Qwen3-32B) em tarefas incluindo raciocínio matemático (GSM8K), geração de código (HumanEval) e sumarização (XSum).

• Ganhos de Throughput: O DEX alcançou consistentemente maior tokens por segundo (TPS) do que as linhas de base de seleção de profundidade. Nos modelos LayerSkip, o DEX realizou uma fração maior do headroom teórico do "oráculo" do que os métodos de seleção.
• Escalabilidade com a Resolução: O throughput melhorou à medida que a resolução de exploração se tornava mais fina (ou seja, mais exploradores eram adicionados). Por exemplo, no Llama-2-70B, aumentar os exploradores de 3 para 8 reduziu significativamente a lacuna em relação ao ganho de velocidade do oráculo.
• Competitividade: Em comparações de ponta a ponta, o DEX (com 3–4 exploradores) foi competitivo com o paralelismo de tensor multi-GPU (TP) e a decodificação especulativa distribuída (PEARL, EAGLE) em hardware correspondente.
• Ablação: A remoção de ramos intermediários (variante de saída única) mostrou que o ganho de desempenho vem especificamente da exploração de múltiplas profundidades, e não apenas do paralelismo das camadas.

5. Significância e Alegações

O artigo afirma que a exploração de profundidade paralela fornece uma maneira escalável de explorar o eixo de profundidade subutilizado da decodificação de LLMs.

• Mudança de Paradigma: Ao contrário dos métodos anteriores que apostam em um único ponto de saída, o DEX trata a adaptação de profundidade como um problema de controle de resolução. Ao explorar um conjunto de profundidades, ele converte a incerteza de "quando sair" em um overhead limitado determinado pelo espaçamento do conjunto de exploração.
• Eficiência Sem Perdas: O DEX alcança esses ganhos sem comprometer a correção da saída (lossless) e sem exigir mudanças na arquitetura do modelo (exceto para adaptadores opcionais em modelos padrão).
• Complementaridade: Os autores observam que o DEX opera no "eixo da profundidade" (quão cedo um token é determinado), o que é complementar aos métodos de aceleração do "eixo do token" (quantos tokens futuros são previstos). Eles sugerem que, embora a combinação seja possível, ela requer coordenação complexa.

Em conclusão, o artigo estabelece que, ao substituir a seleção de profundidade única pela exploração paralela, é possível reduzir significativamente a latência de decodificação mantendo o comportamento exato dos modelos autorregressivos padrão.