Transformadores em Tandem para LLMs Eficientes em Inferência

Tandem Transformers for Inference Efficient LLMs

February 13, 2024

Autores: Aishwarya P S, Pranav Ajit Nair, Yashas Samaga, Toby Boyd, Sanjiv Kumar, Prateek Jain, Praneeth Netrapalli

cs.AI

Resumo

A natureza autoregressiva dos modelos de linguagem grandes (LLMs) convencionais limita inerentemente a velocidade de inferência, uma vez que os tokens são gerados sequencialmente. Embora técnicas de decodificação especulativa e paralela tentem mitigar isso, elas enfrentam limitações: ou dependem de modelos menores e menos precisos para geração ou falham em aproveitar totalmente as representações do LLM base. Introduzimos uma nova arquitetura, os Tandem transformers, para abordar essas questões. Essa arquitetura combina de forma única (1) um pequeno modelo autoregressivo e (2) um modelo grande operando em modo de bloco (processando múltiplos tokens simultaneamente). A precisão preditiva do modelo pequeno é substancialmente aprimorada ao permitir que ele atente para as representações mais ricas do modelo grande. No conjunto de dados de pré-treinamento do PaLM2, um tandem de PaLM2-Bison e PaLM2-Gecko demonstra uma melhoria de 3,3% na precisão de previsão do próximo token em relação a um PaLM2-Gecko autônomo, oferecendo um aumento de velocidade de 1,16x em comparação com um modelo PaLM2-Otter com desempenho comparável em tarefas subsequentes. Além disso, incorporamos o modelo tandem dentro do framework de decodificação especulativa (SPEED), onde o modelo grande valida os tokens do modelo pequeno. Isso garante que o Tandem de PaLM2-Bison e PaLM2-Gecko alcance um aumento substancial de velocidade (cerca de 1,14x mais rápido do que usar o PaLM2-Gecko padrão no SPEED) enquanto mantém a precisão idêntica em tarefas subsequentes.

English

The autoregressive nature of conventional large language models (LLMs) inherently limits inference speed, as tokens are generated sequentially. While speculative and parallel decoding techniques attempt to mitigate this, they face limitations: either relying on less accurate smaller models for generation or failing to fully leverage the base LLM's representations. We introduce a novel architecture, Tandem transformers, to address these issues. This architecture uniquely combines (1) a small autoregressive model and (2) a large model operating in block mode (processing multiple tokens simultaneously). The small model's predictive accuracy is substantially enhanced by granting it attention to the large model's richer representations. On the PaLM2 pretraining dataset, a tandem of PaLM2-Bison and PaLM2-Gecko demonstrates a 3.3% improvement in next-token prediction accuracy over a standalone PaLM2-Gecko, offering a 1.16x speedup compared to a PaLM2-Otter model with comparable downstream performance. We further incorporate the tandem model within the speculative decoding (SPEED) framework where the large model validates tokens from the small model. This ensures that the Tandem of PaLM2-Bison and PaLM2-Gecko achieves substantial speedup (around 1.14x faster than using vanilla PaLM2-Gecko in SPEED) while maintaining identical downstream task accuracy.