🤖 machine learning

TIDE: Token-Informed Depth Execution for Per-Token Early Exit in LLM Inference

O TIDE é um sistema pós-treinamento que utiliza roteadores aprendidos para permitir a saída antecipada de tokens em camadas específicas de modelos de linguagem grandes, reduzindo a latência e aumentando o throughput sem necessidade de re-treinamento do modelo.

Autores originais: Jaber Jaber, Osama Jaber

Publicado 2026-03-24

📖 4 min de leitura☕ Leitura rápida

CC BY 4.0

Autores originais: Jaber Jaber, Osama Jaber

Artigo original sob licença CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta é uma explicação gerada por IA do artigo abaixo. Não foi escrita nem endossada pelos autores. Para precisão técnica, consulte o artigo original. Ler aviso legal completo

Imagine que você tem um time de 32 especialistas (camadas de uma Inteligência Artificial) trabalhando juntos para escrever uma história ou resolver um problema.

No modelo tradicional, todos os 32 especialistas precisam ler e processar cada palavra que sai, não importa o quão simples ela seja.

Se a palavra for "o" ou "a" (artigos simples), os 32 especialistas gastam tempo e energia para analisá-la.
Se a palavra for uma etapa complexa de um cálculo matemático, os 32 especialistas também trabalham nela.

Isso é como pedir para um time de 32 engenheiros de foguete analisarem a cor de uma tinta em um desenho de criança. É um desperdício de tempo e energia.

O TIDE (o sistema apresentado neste artigo) é como um gerente inteligente que muda essa regra.

Como o TIDE funciona? (A Analogia do "Checador Rápido")

O TIDE não demite ninguém e não reescreve o manual de instruções do modelo. Em vez disso, ele coloca pequenos "checadores" (chamados de routers) em pontos estratégicos ao longo do processo de 32 etapas.

O Treinamento (A Calibração):
Antes de começar a trabalhar de verdade, o TIDE observa o modelo resolver 2.000 textos curtos. Ele aprende a reconhecer: "Ah, quando a palavra 'o' passa pela camada 8, ela já está perfeita. Não precisa ir até a camada 32." Ele cria um pequeno mapa de quando cada tipo de palavra está "pronto". Isso leva menos de 3 minutos e ocupa pouquíssimo espaço no disco (4 MB).
O Trabalho (A Execução):
Quando você pede para a IA gerar um texto:
- O modelo ainda passa a palavra por todas as 32 camadas (para garantir que a memória de curto prazo, chamada de KV Cache, não se corrompa e que tudo funcione perfeitamente).
- MAS, assim que a palavra passa por uma camada intermediária, o "checador" olha para ela e pergunta: "Ela já está pronta?"
- Se a resposta for SIM (como no caso da palavra "o"), o sistema diz: "Pare de gastar energia calculando o resto para esta palavra específica. Pegue o resultado desta camada e use-o."
- Se a resposta for NÃO (como em uma palavra difícil de matemática), a palavra continua descendo até a camada 32.

A Metáfora da Fábrica de Pães

Imagine uma fábrica de pães com 32 esteiras de controle de qualidade.

Sem TIDE: Cada pão passa por todas as 32 esteiras, mesmo que o pão seja perfeito na esteira 5. O tempo de produção é sempre o máximo.
Com TIDE: O pão passa por todas as esteiras (para manter a linha organizada), mas um sensor inteligente verifica na esteira 5. Se o pão estiver perfeito, o sistema marca: "Este pão já está pronto, não precisa ser reavaliado nas esteiras 6 a 32 para o cálculo final."
- Resultado: A fábrica produz mais pães por hora (mais velocidade) e gasta menos energia elétrica, sem que o pão fique com gosto ruim.

O Que o TIDE Conseguiu?

Os autores testaram isso em modelos reais (como o DeepSeek e o Qwen) e os resultados foram impressionantes:

Velocidade: O tempo para processar a primeira parte do texto (o "preenchimento") caiu em até 7,2%.
Eficiência: Em modelos de 32 camadas, cerca de 99% das palavras (especialmente as comuns) pararam de ser processadas profundamente na camada 31 ou até antes. Apenas as palavras realmente difíceis foram até o fim.
Qualidade: A IA continua resolvendo problemas de matemática complexos perfeitamente. O TIDE não "corta" o raciocínio; ele apenas corta o trabalho desnecessário em palavras simples.
Facilidade: Funciona com qualquer modelo de IA popular (HuggingFace) sem precisar reprogramar o modelo original. É como instalar um "plug-in" que melhora o desempenho.

Resumo em uma frase

O TIDE é um sistema inteligente que ensina a Inteligência Artificial a parar de trabalhar em palavras simples assim que elas estão prontas, economizando tempo e energia, mas mantendo a inteligência total para as palavras difíceis. É como ter um funcionário que sabe exatamente quando pode relaxar sem prejudicar o resultado final.

Resumo Técnico: TIDE

1. O Problema

Os Grandes Modelos de Linguagem (LLMs) baseados em Transformers atribuem a mesma quantidade de computação a cada token, independentemente da sua complexidade semântica. Palavras funcionais simples (como "o" ou "a") recebem o mesmo tratamento de 32 camadas que passos de raciocínio complexos em derivações matemáticas.

Ineficiência: Estudos mostram que, para uma grande fração de tokens, os estados ocultos intermediários convergem para o estado final muito antes da última camada. As camadas subsequentes produzem mudanças negligenciáveis na distribuição de saída, representando um desperdício de recursos computacionais.
Limitações das Soluções Atuais:
- Métodos baseados em encoder (ex: DeeBERT) não generalizam para geração autoregressiva com caches KV.
- Métodos que exigem re-treinamento (ex: LayerSkip, EE-LLM) são impraticáveis para usuários que utilizam checkpoints pré-treinados de hubs de modelos.
- Heurísticas baseadas em confiança (ex: entropia do softmax) são pouco confiáveis para geração, onde a entropia é naturalmente alta.

2. Metodologia

O TIDE é um sistema de saída antecipada pós-treinamento (post-training) que não requer modificação do modelo base nem re-treinamento. Ele opera em duas fases principais:

A. Calibração (Offline - Uma única vez por modelo)

Coleta de Estados Ocultos: O modelo congelado é executado em 2.000 amostras de texto (WikiText-103), coletando vetores de estados ocultos em camadas de checkpoint periódicas e na camada final.
Rótulos de Convergência: Calcula-se a similaridade de cosseno entre o estado oculto de um token em uma camada intermediária $k$ e o estado final $L$ . Se a similaridade superar um limiar $\tau$ (padrão 0.98), o token é marcado como "convergido".
Treinamento de Roteadores: Para cada camada de checkpoint, treina-se um classificador binário leve (um MLP de duas camadas com dimensão de gargalo de 128). O objetivo é prever se o token já convergiu naquela camada específica. O treinamento leva menos de 3 minutos e gera um checkpoint de apenas ~4 MB.

B. Inferência (Online - Por solicitação)

Passo Forward Completo: Durante a geração autoregressiva, o modelo executa o forward pass completo (todas as camadas) com output_hidden_states=True. Isso garante que o KV Cache permaneça íntegro e consistente.
Avaliação Pós-Hoc: Após o forward pass, os roteadores treinados avaliam os estados ocultos de cada camada de checkpoint.
Seleção de Saída: O sistema identifica a camada mais antiga onde o roteador indicou convergência para todos os tokens do lote e utiliza os logits dessa camada (após normalização) para a próxima previsão, ignorando as camadas subsequentes para o cálculo da saída, embora elas tenham sido computadas.

Tecnologias de Implementação:

Adapter Universal: Um módulo que detecta automaticamente a estrutura de 17 caminhos de atributos, suportando arquiteturas como LLaMA, Mistral, Qwen, GPT-2, Phi, Falcon, OPT, etc., sem código específico por modelo.
Kernels CUDA Fundidos: Implementação de kernels em CUDA que fundem RMSNorm, projeção, ativação SiLU e sigmoid em uma única chamada, com suporte nativo a fp16, bf16 e fp32.
Compatibilidade: Funciona com qualquer versão da biblioteca transformers (incluindo v5.3+), evitando problemas de corrupção de cache comuns em métodos baseados em exceções ou hooks.

3. Principais Contribuições

Sistema Pós-Treinamento: Permite saída antecipada por token em LLMs autoregressivos sem modificar o modelo ou re-treiná-lo.
Universalidade: Suporte automático a qualquer modelo causal do HuggingFace através de um adaptador universal.
Otimização de Hardware: Kernels CUDA fundidos para avaliação de roteadores e normalização, com detecção automática de arquitetura GPU (de V100 a Blackwell).
Integridade do KV Cache: Diferente de métodos que pulam camadas fisicamente (o que quebra o cache), o TIDE mantém o cache intacto ao executar todas as camadas, mas seleciona a saída da camada correta para a próxima iteração.
Código Aberto: Liberação completa com 3.097 linhas de código (Python + CUDA/C++), 74 testes passando e instalação via PyPI.

4. Resultados Experimentais

Os testes foram realizados em uma GPU NVIDIA A100 com modelos como DeepSeek R1 Distill 8B (32 camadas) e Qwen3 8B (36 camadas).

Taxas de Saída (Prefill):
- 100% dos tokens encontraram um ponto de saída.
- No DeepSeek R1, 5% dos tokens saíram na camada 11 (apenas 1/3 do caminho), e o restante na camada 31.
Latência e Vazão (Throughput):
- Latência de Prefill: Redução de 5,5% a 7,2% no DeepSeek R1 e 5,7% no Qwen3.
- Vazão (Throughput): Aumento de 6,6% (DeepSeek R1, BS=1) e 8,1% (Qwen3, BS=8).
- Nota: Em lotes grandes (BS=8) no DeepSeek R1, houve uma redução de vazão (-16,3%) devido ao custo superlinear de coletar estados ocultos de todas as camadas, indicando uma área para otimização futura.
Qualidade de Geração:
- Em problemas matemáticos de múltiplos passos, 98–99% dos tokens de decodificação saíram antecipadamente (geralmente na camada 31).
- A precisão matemática foi mantida (mesmo número de tokens únicos e raciocínio correto) comparado ao modelo base de profundidade total.
Calibração: Leva menos de 3 minutos em uma única GPU e gera um checkpoint de roteador de ~4 MB.

5. Significado e Limitações

Significado:
O TIDE demonstra que a execução de profundidade variável por token é viável em modelos pré-treinados existentes sem re-treinamento. Ele oferece uma via prática para reduzir latência e consumo de energia em inferência de LLMs, mantendo a integridade do cache e a compatibilidade com o ecossistema atual de bibliotecas.

Limitações e Trabalhos Futuros:

Modo "Pós-Hoc": O método atual executa fisicamente todas as camadas para preservar o KV Cache. Isso reduz a latência de saída (pela seleção de logits) e a vazão, mas não elimina o custo computacional das camadas subsequentes (não é um "pulo" físico de camadas). Um modo de "pulo real" exigiria gerenciar descontinuidades no cache.
Limiar Conservador: O limiar de convergência ( $\tau=0.98$ ) faz com que a maioria das saídas ocorra na penúltima camada. Ajustes dinâmicos desse limiar poderiam permitir saídas mais precoces.
Overhead de Batch: A coleta de estados ocultos de todas as camadas torna-se um gargalo em lotes grandes, sugerindo a necessidade de uma abordagem híbrida que colete estados apenas nas camadas de checkpoint.

Em suma, o TIDE representa um avanço significativo na eficiência de inferência de LLMs, equilibrando ganhos de desempenho com a simplicidade de implementação em modelos pré-treinados.