Trasformatori Ricorsivi Rilassati: Condivisione Efficace dei Parametri con LoRA a Livello di Strato

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

October 28, 2024

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

cs.AI

Abstract

I modelli linguistici di grandi dimensioni (LLM) sono costosi da implementare. La condivisione dei parametri offre un possibile percorso per ridurne dimensioni e costi, ma la sua efficacia nei moderni LLM rimane piuttosto limitata. In questo lavoro, riprendiamo il "layer tying" come forma di condivisione dei parametri nei Transformers e introduciamo nuovi metodi per convertire i LLM esistenti in "Recursive Transformers" più piccoli che condividono i parametri tra i vari layer, con una perdita minima delle prestazioni. Qui, i nostri Recursive Transformers vengono inizializzati in modo efficiente dai Transformers preaddestrati standard, ma utilizzano solo un singolo blocco di layer unici che viene poi ripetuto più volte in un ciclo. Miglioriamo ulteriormente le prestazioni introducendo i Relaxed Recursive Transformers che aggiungono flessibilità al vincolo di condivisione dei layer tramite moduli di adattamento a basso rango (LoRA) in profondità, pur preservando la compattezza del modello complessivo. Dimostriamo che i nostri modelli ricorsivi (ad esempio, Gemma 1B ricorsivo) superano sia modelli preaddestrati simili in dimensioni (come TinyLlama 1.1B e Pythia 1B) che baselines di distillazione della conoscenza, e possono persino recuperare la maggior parte delle prestazioni del modello "a dimensioni complete" originale (ad esempio, Gemma 2B senza parametri condivisi). Infine, proponiamo il Continuous Depth-wise Batching, un nuovo paradigma promettente per l'inferenza reso possibile dal Recursive Transformer abbinato all'uscita anticipata. In un'analisi teorica, mostriamo che questo ha il potenziale per portare a significativi aumenti (2-3 volte) nella capacità di elaborazione dell'inferenza.

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.