Transformateurs récursifs relaxés : Partage efficace des paramètres avec LoRA par couche.

Relaxed Recursive Transformers: Effective Parameter Sharing with Layer-wise LoRA

October 28, 2024

Auteurs: Sangmin Bae, Adam Fisch, Hrayr Harutyunyan, Ziwei Ji, Seungyeon Kim, Tal Schuster

cs.AI

Résumé

Les grands modèles de langage (LLM) sont coûteux à déployer. Le partage de paramètres offre une voie possible pour réduire leur taille et leur coût, mais son efficacité dans les LLM modernes reste assez limitée. Dans ce travail, nous revisitons le "layer tying" en tant que forme de partage de paramètres dans les Transformers, et introduisons de nouvelles méthodes pour convertir les LLM existants en "Recursive Transformers" plus petits qui partagent des paramètres entre les couches, avec une perte minimale de performance. Ici, nos Recursive Transformers sont initialisés de manière efficace à partir des Transformers pré-entraînés standard, mais n'utilisent qu'un seul bloc de couches uniques qui est ensuite répété plusieurs fois en boucle. Nous améliorons encore les performances en introduisant des Relaxed Recursive Transformers qui ajoutent de la flexibilité à la contrainte de layer tying via des modules d'adaptation à faible rang en profondeur (LoRA), tout en préservant la compacité du modèle global. Nous montrons que nos modèles récursifs (par exemple, Gemma 1B récursif) surpassent à la fois des modèles pré-entraînés de taille similaire (comme TinyLlama 1.1B et Pythia 1B) et des références de distillation de connaissances - et peuvent même retrouver la plupart des performances du modèle "full-size" original (par exemple, Gemma 2B sans paramètres partagés). Enfin, nous proposons le Batching en profondeur continue, un nouveau paradigme prometteur d'inférence rendu possible par le Recursive Transformer lorsqu'il est associé à des sorties anticipées. Dans une analyse théorique, nous montrons que cela a le potentiel de conduire à des gains significatifs (2-3x) dans le débit d'inférence.

English

Large language models (LLMs) are expensive to deploy. Parameter sharing offers a possible path towards reducing their size and cost, but its effectiveness in modern LLMs remains fairly limited. In this work, we revisit "layer tying" as form of parameter sharing in Transformers, and introduce novel methods for converting existing LLMs into smaller "Recursive Transformers" that share parameters across layers, with minimal loss of performance. Here, our Recursive Transformers are efficiently initialized from standard pretrained Transformers, but only use a single block of unique layers that is then repeated multiple times in a loop. We further improve performance by introducing Relaxed Recursive Transformers that add flexibility to the layer tying constraint via depth-wise low-rank adaptation (LoRA) modules, yet still preserve the compactness of the overall model. We show that our recursive models (e.g., recursive Gemma 1B) outperform both similar-sized vanilla pretrained models (such as TinyLlama 1.1B and Pythia 1B) and knowledge distillation baselines -- and can even recover most of the performance of the original "full-size" model (e.g., Gemma 2B with no shared parameters). Finally, we propose Continuous Depth-wise Batching, a promising new inference paradigm enabled by the Recursive Transformer when paired with early exiting. In a theoretical analysis, we show that this has the potential to lead to significant (2-3x) gains in inference throughput.