Entspannte rekursive Transformer: Effektives Parameter-Sharing mit schichtweisem LoRA

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

October 28, 2024

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

cs.AI

Zusammenfassung

Große Sprachmodelle (LLMs) sind teuer in der Bereitstellung. Das Teilen von Parametern bietet einen möglichen Weg zur Reduzierung ihrer Größe und Kosten, aber seine Wirksamkeit in modernen LLMs bleibt recht begrenzt. In dieser Arbeit überdenken wir "Layer Tying" als Form des Parameterteilens in Transformern und führen neue Methoden ein, um bestehende LLMs in kleinere "Rekursive Transformer" umzuwandeln, die Parameter über Schichten hinweg teilen, mit minimalem Leistungsverlust. Hier werden unsere Rekursiven Transformer effizient von standardmäßig vorab trainierten Transformern initialisiert, verwenden jedoch nur einen einzigen Block einzigartiger Schichten, der dann mehrmals in einer Schleife wiederholt wird. Wir verbessern die Leistung weiter, indem wir "Relaxed Recursive Transformers" einführen, die über Tiefen-weise Low-Rank-Anpassung (LoRA) Module Flexibilität zur Schichtbindung hinzufügen, aber dennoch die Kompaktheit des Gesamtmodells bewahren. Wir zeigen, dass unsere rekursiven Modelle (z. B. rekursive Gemma 1B) sowohl ähnlich großen vorab trainierten Modellen (wie TinyLlama 1.1B und Pythia 1B) und Wissensverdichtungs-Baselines übertreffen - und sogar die meisten der Leistung des ursprünglichen "Vollgrößen"-Modells (z. B. Gemma 2B ohne gemeinsame Parameter) wiederherstellen können. Abschließend schlagen wir kontinuierliches Tiefen-Weises Batching vor, ein vielversprechendes neues Inferenz-Paradigma, das durch den Rekursiven Transformer in Verbindung mit frühzeitigem Beenden ermöglicht wird. In einer theoretischen Analyse zeigen wir, dass dies das Potenzial hat, zu signifikanten (2-3-fachen) Gewinnen bei der Inferenzdurchsatz zu führen.

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.