Trasformatori in Tandem per LLM Efficienti nell'Inferenza

Tandem Transformers for Inference Efficient LLMs

February 13, 2024

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

cs.AI

Abstract

La natura autoregressiva dei convenzionali modelli linguistici di grandi dimensioni (LLM) limita intrinsecamente la velocità di inferenza, poiché i token vengono generati in sequenza. Sebbene le tecniche di decodifica speculativa e parallela tentino di mitigare questo problema, presentano delle limitazioni: o si basano su modelli più piccoli e meno accurati per la generazione, o non riescono a sfruttare appieno le rappresentazioni del modello LLM di base. Introduciamo una nuova architettura, i Tandem Transformers, per affrontare queste problematiche. Questa architettura combina in modo unico (1) un piccolo modello autoregressivo e (2) un modello di grandi dimensioni che opera in modalità a blocchi (elaborando più token simultaneamente). L'accuratezza predittiva del modello piccolo viene notevolmente migliorata consentendogli di prestare attenzione alle rappresentazioni più ricche del modello grande. Sul dataset di pre-addestramento PaLM2, un tandem composto da PaLM2-Bison e PaLM2-Gecko dimostra un miglioramento del 3,3% nell'accuratezza della predizione del token successivo rispetto a un PaLM2-Gecko autonomo, offrendo un'accelerazione di 1,16x rispetto a un modello PaLM2-Otter con prestazioni comparabili nei task downstream. Integriamo ulteriormente il modello tandem all'interno del framework di decodifica speculativa (SPEED), in cui il modello grande convalida i token generati dal modello piccolo. Ciò garantisce che il tandem di PaLM2-Bison e PaLM2-Gecko raggiunga un'accelerazione significativa (circa 1,14x più veloce rispetto all'uso del vanilla PaLM2-Gecko in SPEED) mantenendo al contempo un'accuratezza identica nei task downstream.

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.