Transformateurs en Tandem pour des Modèles de Langage à Inférence Efficace

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

Résumé

La nature autorégressive des grands modèles de langage (LLM) conventionnels limite intrinsèquement la vitesse d'inférence, car les tokens sont générés séquentiellement. Bien que les techniques de décodage spéculatif et parallèle tentent d'atténuer ce problème, elles rencontrent des limitations : soit elles s'appuient sur des modèles plus petits et moins précis pour la génération, soit elles ne parviennent pas à exploiter pleinement les représentations du LLM de base. Nous introduisons une architecture novatrice, les Tandem Transformers, pour résoudre ces problèmes. Cette architecture combine de manière unique (1) un petit modèle autorégressif et (2) un grand modèle fonctionnant en mode bloc (traitant plusieurs tokens simultanément). La précision prédictive du petit modèle est considérablement améliorée en lui accordant un accès attentionnel aux représentations plus riches du grand modèle. Sur le jeu de données de pré-entraînement PaLM2, un tandem composé de PaLM2-Bison et PaLM2-Gecko démontre une amélioration de 3,3 % dans la précision de prédiction du token suivant par rapport à un PaLM2-Gecko autonome, offrant une accélération de 1,16x par rapport à un modèle PaLM2-Otter avec des performances comparables en aval. Nous intégrons en outre le modèle tandem dans le cadre de décodage spéculatif (SPEED), où le grand modèle valide les tokens générés par le petit modèle. Cela garantit que le tandem de PaLM2-Bison et PaLM2-Gecko atteint une accélération substantielle (environ 1,14x plus rapide que l'utilisation de PaLM2-Gecko standard dans SPEED) tout en maintenant une précision identique sur les tâches en aval.

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.