Tandem Transformers voor Inferentie-efficiënte LLM's

Tandem Transformers for Inference Efficient LLMs

February 13, 2024

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

cs.AI

Samenvatting

De autoregressieve aard van conventionele grote taalmodellen (LLM's) beperkt inherent de inferentiesnelheid, omdat tokens sequentieel worden gegenereerd. Hoewel speculatieve en parallelle decodeertechnieken proberen dit te verlichten, hebben ze beperkingen: ze vertrouwen ofwel op minder nauwkeurige kleinere modellen voor generatie, of benutten de representaties van het basis-LLM niet volledig. We introduceren een nieuwe architectuur, Tandem-transformers, om deze problemen aan te pakken. Deze architectuur combineert uniek (1) een klein autoregressief model en (2) een groot model dat in blokmodus werkt (waarbij meerdere tokens tegelijkertijd worden verwerkt). De voorspellingsnauwkeurigheid van het kleine model wordt aanzienlijk verbeterd door het aandacht te geven aan de rijkere representaties van het grote model. Op de PaLM2-pre-trainingsdataset laat een tandem van PaLM2-Bison en PaLM2-Gecko een verbetering van 3,3% zien in de nauwkeurigheid van next-token-voorspellingen ten opzichte van een standalone PaLM2-Gecko, wat een versnelling van 1,16x biedt in vergelijking met een PaLM2-Otter-model met vergelijkbare downstreamprestaties. We integreren het tandemmodel verder binnen het speculatieve decodeerframework (SPEED), waarbij het grote model tokens van het kleine model valideert. Dit zorgt ervoor dat de Tandem van PaLM2-Bison en PaLM2-Gecko een aanzienlijke versnelling bereikt (ongeveer 1,14x sneller dan het gebruik van standaard PaLM2-Gecko in SPEED) terwijl de nauwkeurigheid van downstreamtaken identiek blijft.

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.