Adaptive Layer Selection for Layer-Wise Token Pruning in LLM Inference

1. O Problema

A inferência de Grandes Modelos de Linguagem (LLMs) em contextos longos enfrenta um gargalo crítico: o tamanho da memória do Cache Key-Value (KV). Para gerar tokens eficientemente, os modelos armazenam vetores de chave e valor de todos os tokens anteriores, o que consome uma quantidade massiva de memória à medida que o contexto cresce.

Para mitigar isso, surgiram métodos de poda de tokens por camada (layer-wise token pruning), que selecionam um subconjunto de tokens importantes em camadas específicas para manter no cache KV, descartando os demais. No entanto, a abordagem existente sofre de uma limitação fundamental:

• Rigidez na Seleção de Camadas: Métodos atuais (como FastKV, GemFilter) utilizam camadas de seleção pré-definidas e fixas, independentemente da tarefa.
• Variação de Desempenho: A dificuldade de uma tarefa varia. Em tarefas simples (ex: QA), informações relevantes podem ser identificadas cedo (camadas médias). Em tarefas difíceis (ex: recuperação de pares chave-valor ou "Needle in a Haystack" complexa), a similaridade semântica exige que o modelo processe mais camadas antes de focar nos tokens corretos.
• Consequência: Uma camada fixa pode funcionar bem para tarefas fáceis, mas causa degradação severa de precisão em tarefas difíceis, ou exige aumentar o orçamento de memória (KV budget), anulando os benefícios de compressão.

2. Metodologia: ASL (Adaptive Selection Layer)

Os autores propõem o ASL, um método livre de treinamento (training-free) que determina dinamicamente a camada ideal para realizar a seleção de tokens durante a fase de prefilling (processamento do prompt de entrada).

Princípio Central:
O método baseia-se na observação de que, em camadas iniciais, as pontuações de atenção são distribuídas uniformemente, enquanto em camadas mais profundas, elas se concentram consistentemente em um subconjunto estável de tokens. O ASL monitora a variância dos ranks dos tokens ordenados pela pontuação de atenção.

Funcionamento do Algoritmo:

1. Monitoramento de Variância: A partir de uma camada inicial ($L_{min}$), o ASL calcula a variância dos ranks dos tokens com as maiores pontuações de atenção ao longo de uma janela de camadas observadas ($L_{obs}$).
2. Cálculo de Variância Relativa: Para cada camada $L$, calcula-se a variância dos ranks dos tokens de topo ($top-k$) e normaliza-se pela variância inicial em $L_{min}$.
   $$\text{relative\_variance}(L) = \frac{\text{variance}(L)}{\text{variance}(L_{min})}$$
3. Critério de Parada (Seleção): O processo continua até que a variância relativa caia abaixo de um limiar ($\tau$) especificado pelo usuário.
   • Quando a variância diminui, isso indica que a ordem dos tokens mais importantes estabilizou (as pontuações de atenção focaram consistentemente no mesmo subconjunto).
   • Nesse ponto, a camada atual é escolhida como a camada de seleção.
4. Seleção de Tokens: Uma vez determinada a camada, os tokens são selecionados (top-k) e propagados para todas as camadas subsequentes de forma "one-shot" (única vez).
5. Integração: O ASL pode ser combinado com métodos existentes (como SnapKV para compressão antes da seleção ou GemFilter para uma estratégia de duas passagens) para atender a orçamentos de KV específicos.

3. Principais Contribuições

1. Identificação da Limitação de Métodos Fixos: Demonstraram que a seleção de camadas fixa é inflexível e que a dificuldade da tarefa (similaridade semântica entre pergunta e contexto) impacta diretamente a camada ótima de poda.
2. Método Adaptativo (ASL): Propuseram um algoritmo que ajusta automaticamente a camada de seleção baseado no padrão de atenção observado em tempo de inferência, sem necessidade de ajuste fino (fine-tuning).
3. Integração Modular: O ASL foi projetado para funcionar em conjunto com técnicas de redução de cache KV existentes, otimizando tanto a fase de prefilling quanto a de decoding.
4. Validação Abrangente: Demonstraram através de experimentos que o ASL supera métodos state-of-the-art (SOTA) em tarefas difíceis, mantendo um equilíbrio superior entre precisão e eficiência.

4. Resultados Experimentais

O método foi avaliado em três benchmarks de contexto longo: InfiniteBench, RULER e Needle in a Haystack (NIAH), utilizando modelos como Llama-3.1-8B e Qwen2.5-7B.

• Precisão em Tarefas Difíceis: O ASL superou significativamente métodos como FastKV e GemFilter em tarefas complexas como recuperação de KV (KV retrieval) e NIAH multi-chave. Enquanto métodos fixos falhavam drasticamente nesses cenários, o ASL manteve a precisão próxima à do modelo com cache completo (Full KV).
• Trade-off Precisão/Eficiência:
  • O ASL oferece uma melhor compensação entre velocidade de inferência e precisão. Em tarefas fáceis, ele seleciona camadas mais cedo (similar aos métodos fixos), garantindo velocidade. Em tarefas difíceis, ele seleciona camadas mais profundas, preservando a precisão.
  • Em benchmarks como InfiniteBench, a versão ASL_2pass (integrada ao GemFilter) alcançou a maior pontuação média.
• Eficiência de Memória: O uso de memória do ASL é comparável aos outros métodos de poda (cerca de 0.2GB - 0.3GB para contextos longos), sendo negligenciável o overhead adicional do cálculo de variância.
• Velocidade (TTFT e TPOT):
  • O Time to First Token (TTFT) do ASL é ligeiramente maior que o do FastKV (que usa camadas muito cedo), pois o ASL pode precisar processar mais camadas para tarefas difíceis. No entanto, essa "perda" de tempo inicial é compensada pela alta precisão.
  • O Time Per Output Token (TPOT) é comparável aos outros métodos, pois a decodificação ocorre com o mesmo número reduzido de tokens.

5. Significado e Impacto

O trabalho do ASL é significativo porque introduz adaptabilidade em um domínio onde a maioria das soluções de otimização de LLMs é estática.

• Resolução do Dilema "Precisão vs. Memória": Permite que os usuários especifiquem um orçamento de memória (KV budget) e obtenham a melhor precisão possível para qualquer tipo de tarefa, sem precisar reconfigurar manualmente a camada de poda para cada novo cenário.
• Viabilidade para Aplicações Reais: Ao lidar eficazmente com tarefas de recuperação de informações em contextos extremamente longos (até 256k tokens), o ASL torna viável a execução de LLMs em cenários de produção que exigem alta precisão sem o custo proibitivo de manter o cache KV completo.
• Abordagem Livre de Treinamento: A natureza "training-free" facilita a adoção imediata em modelos existentes sem a necessidade de custos computacionais para ajuste fino.

Em resumo, o ASL representa um avanço na otimização de inferência de LLMs, substituindo heurísticas fixas por uma decisão dinâmica baseada no comportamento interno do modelo, garantindo robustez tanto em tarefas simples quanto nas mais desafiadoras de recuperação de contexto.