Ontspannen Recursieve Transformers: Effectieve Parameterdeling met Laagsgewijze LoRA

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

Samenvatting

Grote taalmodellen (LLM's) zijn duur om in te zetten. Parameterdeling biedt een mogelijke weg om hun omvang en kosten te verlagen, maar de effectiviteit ervan in moderne LLM's blijft vrij beperkt. In dit werk herzien we "laagbinding" als een vorm van parameterdeling in Transformers, en introduceren we nieuwe methoden om bestaande LLM's om te zetten in kleinere "Recursieve Transformers" die parameters delen over lagen, met minimaal verlies aan prestaties. Hierbij worden onze Recursieve Transformers efficiënt geïnitialiseerd vanuit standaard voorgeleerde Transformers, maar maken ze slechts gebruik van een enkel blok unieke lagen dat vervolgens meerdere malen in een lus wordt herhaald. We verbeteren de prestaties verder door het introduceren van Ontspannen Recursieve Transformers die flexibiliteit toevoegen aan de laagbindingbeperking via diepte-gewijze lage-rang aanpassingsmodules (LoRA), maar toch de compactheid van het totale model behouden. We tonen aan dat onze recursieve modellen (bijv. recursieve Gemma 1B) zowel vergelijkbare in omvang zijnde standaard voorgeleerde modellen (zoals TinyLlama 1.1B en Pythia 1B) als baselines voor kennisoverdracht overtreffen - en zelfs het grootste deel van de prestaties van het oorspronkelijke "volledige" model (bijv. Gemma 2B zonder gedeelde parameters) kunnen herstellen. Tot slot stellen we Continue Diepte-gewijze Batching voor, een veelbelovend nieuw inferentieparadigma mogelijk gemaakt door de Recursieve Transformer in combinatie met vroegtijdig verlaten. In een theoretische analyse tonen we aan dat dit het potentieel heeft om te leiden tot aanzienlijke (2-3x) winst in inferentiedoorvoer.

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.