Transformadores Recursivos Relajados: Compartición Efectiva de Parámetros con LoRA a Nivel de Capa

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

October 28, 2024

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

cs.AI

Resumen

Los modelos de lenguaje grandes (LLMs) son costosos de implementar. El uso compartido de parámetros ofrece un posible camino para reducir su tamaño y coste, pero su efectividad en los LLMs modernos sigue siendo bastante limitada. En este trabajo, revisitamos el "anudamiento de capas" como una forma de compartir parámetros en Transformadores, e introducimos métodos novedosos para convertir LLMs existentes en "Transformadores Recursivos" más pequeños que comparten parámetros entre capas, con una pérdida mínima de rendimiento. Aquí, nuestros Transformadores Recursivos se inicializan eficientemente a partir de Transformadores preentrenados estándar, pero solo utilizan un bloque único de capas que luego se repite varias veces en un bucle. Mejoramos aún más el rendimiento al introducir Transformadores Recursivos Relajados que añaden flexibilidad a la restricción de anudamiento de capas a través de módulos de adaptación de rango bajo (LoRA) en profundidad, manteniendo la compacidad del modelo en general. Mostramos que nuestros modelos recursivos (por ejemplo, Gemma 1B recursivo) superan tanto a modelos preentrenados similares en tamaño (como TinyLlama 1.1B y Pythia 1B) como a líneas base de destilación de conocimiento, e incluso pueden recuperar la mayor parte del rendimiento del modelo "tamaño completo" original (por ejemplo, Gemma 2B sin parámetros compartidos). Finalmente, proponemos el "Agrupamiento Continuo en Profundidad", un nuevo paradigma prometedor de inferencia habilitado por el Transformador Recursivo cuando se combina con la salida temprana. En un análisis teórico, mostramos que esto tiene el potencial de llevar a ganancias significativas (2-3 veces) en el rendimiento de inferencia.

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.