Task-Centric Acceleration of Small-Language Models

Título: Aceleração Centrada em Tarefas de Modelos de Linguagem Pequenos (SLMs)

Autores: Dor Tsur, Sharon Adar, Ran Levy (Amazon)

1. Problema e Motivação

Os Modelos de Linguagem Pequenos (SLMs) tornaram-se alternativas viáveis e eficientes aos Grandes Modelos de Linguagem (LLMs) para aplicações específicas, como sistemas multi-agente e implantação em dispositivos com recursos limitados (edge computing). No entanto, em cenários de alto volume e baixa latência, a eficiência da inferência é crítica.

O gargalo principal na geração autoregressiva é o número de passos de decodificação necessários para produzir uma resposta. A maioria das pesquisas anteriores foca em reduzir requisitos de memória (quantização) ou acelerar o processamento de entrada (KV caching). Este trabalho propõe uma abordagem diferente: focar na redução do número de passos de geração explorando a baixa variabilidade de saída inerente a tarefas específicas.

A premissa central é que tarefas bem definidas (como classificação, QA médica, extração de slots) possuem uma distribuição de saída altamente concentrada, com n-gramas (sequências de palavras) de alta frequência que se repetem consistentemente. Isso resulta em uma entropia empírica de saída significativamente menor em comparação com a geração de texto aberto.

2. Metodologia: Framework TASC

Os autores propõem o TASC (Task-Adaptive Sequence Compression), um framework que oferece duas soluções distintas dependendo do cenário de uso:

A. TASC-ft (Fine-Tuning Adaptativo)

Utilizado quando é possível realizar o fine-tuning do modelo para a tarefa específica.

• Enriquecimento do Tokenizador: O método identifica n-gramas de alta frequência no corpus de saída esperado da tarefa.
• Algoritmo de Mesclagem: Utiliza um esquema inspirado no Byte Pair Encoding (BPE) para iterativamente fundir esses n-gramas frequentes em novos tokens únicos.
• Critérios de Seleção:
  • Recompensa de Mesclagem ($f_{merge}$): Baseada na frequência do n-grama e na redução esperada no comprimento da sequência.
  • Pontuação de Colisão de Prefixo (PCS): Um mecanismo de segurança que evita adicionar tokens que colidiriam com prefixos de outros tokens existentes, prevenindo viés na decodificação gulosa.
• Ajuste Fino: O vocabulário expandido é integrado ao modelo. As embeddings dos novos tokens são inicializadas como a média das embeddings de seus tokens constituintes. O modelo é então ajustado (fine-tuned) para utilizar esse novo vocabulário, otimizando as embeddings dos novos tokens junto com os adaptadores (ex: LoRA).
• Objetivo: Reduzir o número de tokens necessários para representar a saída, diminuindo diretamente os passos de inferência.

B. TASC-spec (Decodificação Especulativa sem Treinamento)

Utilizado quando o fine-tuning não é viável ou quando se deseja acelerar um modelo já ajustado.

• Modelo de Rascunho (Drafter) Baseado em N-gramas: Em vez de treinar um modelo auxiliar paramétrico (que consome memória e tempo), o TASC-spec constrói um modelo de rascunho leve e sem treinamento diretamente a partir do corpus de saída da tarefa.
• Mistura de Informações: O modelo de rascunho combina:
  1. Informação Global ($p_{corp}$): N-gramas derivados do corpus de treinamento da tarefa.
  2. Informação Local ($p_{prompt}$): N-gramas derivados do contexto imediato (prompt) da instância atual.
  • A distribuição final é uma mistura ponderada: $p_{mixed} = \lambda p_{corp} + (1-\lambda) p_{prompt}$.
• Decodificação Especulativa: O modelo de rascunho gera uma sequência de tokens candidatos. O modelo alvo (SLM) verifica esses tokens em paralelo em uma única passagem de forward pass.
• Vantagens: Não requer treinamento adicional, é agnóstico ao tokenizador (não exige que o rascunho e o alvo compartilhem o mesmo vocabulário) e tem custo computacional de rascunho quase nulo.

3. Contribuições Principais

1. Caracterização Teórica: Formalização da relação entre variabilidade de entrada/saída e a eficiência da geração através da perspectiva da entropia de n-gramas típicos.
2. TASC-ft: Um algoritmo que enriquece o vocabulário do tokenizador com n-gramas específicos da tarefa, integrando-se perfeitamente ao fine-tuning eficiente de parâmetros.
3. Indicador Preditivo: Demonstração de que a Entropia de Rényi de ordem 2 da distribuição de tokens pode atuar como um proxy confiável para prever a redução de tempo de execução após o fine-tuning.
4. TASC-spec: Um método de decodificação especulativa leve, sem treinamento, que utiliza modelos de rascunho baseados em n-gramas, alcançando acelerações significativas sem sobrecarga de memória.

4. Resultados Experimentais

Os métodos foram avaliados em múltiplos conjuntos de dados (EUR-LEX, Asclepius, MASSIVE, MedQuaD) com modelos como Qwen2.5 e Llama.

• TASC-ft:

  • Aceleração: Redução de tempo de execução de até 2.1x (em alguns casos, redução de 2.65x no comprimento médio da saída).
  • Qualidade: A qualidade da geração (medida por F1, BERT-score e avaliação de LLM) foi preservada ou ligeiramente melhorada, mesmo com a adição de tokens.
  • Eficiência: O uso de um orçamento de tokens moderado (ex: 500 novos tokens) já oferece ganhos substanciais, com retornos marginais decrescentes além desse ponto.
  • Correlação: A Entropia de Rényi mostrou uma correlação inversa forte ($\tau < -0.9$) com o tempo de execução, validando sua utilidade como métrica preditiva.

• TASC-spec:

  • Aceleração: Atingiu acelerações de até 3.15x em tarefas de QA médica e 2.17x em classificação massiva.
  • Comparação: Superou métodos existentes como n-grammys, especialmente quando combinado com informações locais do prompt.
  • Taxa de Aceitação: Mantém taxas de aceitação altas nas posições iniciais da sequência (ex: 57% no primeiro token em Asclepius), permitindo longos horizontes de previsão.

5. Significado e Conclusão

O trabalho TASC oferece uma mudança de paradigma na aceleração de LLMs/SLMs: em vez de focar apenas na infraestrutura ou na arquitetura do modelo, ele explora a estrutura estatística das tarefas específicas.

• Eficiência Operacional: Permite que SLMs operem com latência extremamente baixa em cenários de produção, tornando-os viáveis para aplicações em tempo real e em dispositivos de borda.
• Flexibilidade: Oferece duas vias (com e sem fine-tuning), adaptando-se a diferentes restrições de recursos e desenvolvimento.
• Simplicidade: A abordagem baseada em n-gramas e tokenização adaptativa é computacionalmente barata e fácil de implementar, evitando a complexidade de treinar modelos de rascunho paramétricos complexos.

Em resumo, o TASC demonstra que alinhar a representação do modelo (tokenização) e a estratégia de inferência (especulação) com a baixa variabilidade de saída de tarefas específicas pode gerar ganhos de eficiência dramáticos sem sacrificar a qualidade da resposta.