Transformadores Recursivos Relaxados: Compartilhamento Eficaz de Parâmetros com LoRA por Camada

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

Resumo

Modelos de linguagem de grande escala (LLMs) são caros de implantar. O compartilhamento de parâmetros oferece um caminho possível para reduzir seu tamanho e custo, mas sua eficácia em LLMs modernos ainda é bastante limitada. Neste trabalho, revisitamos o "layer tying" como uma forma de compartilhamento de parâmetros em Transformers e introduzimos métodos inovadores para converter LLMs existentes em "Transformers Recursivos" menores que compartilham parâmetros entre camadas, com perda mínima de desempenho. Aqui, nossos Transformers Recursivos são inicializados de forma eficiente a partir de Transformers pré-treinados padrão, mas usam apenas um bloco de camadas exclusivas que é então repetido várias vezes em um loop. Melhoramos ainda mais o desempenho ao introduzir Transformers Recursivos Relaxados que adicionam flexibilidade à restrição de "layer tying" por meio de módulos de adaptação de baixa classificação em profundidade (LoRA), mas ainda preservam a compacidade do modelo geral. Mostramos que nossos modelos recursivos (por exemplo, Gemma 1B recursivo) superam tanto modelos pré-treinados semelhantes em tamanho (como TinyLlama 1.1B e Pythia 1B) quanto baselines de destilação de conhecimento - e podem até recuperar a maior parte do desempenho do modelo "tamanho completo" original (por exemplo, Gemma 2B sem parâmetros compartilhados). Finalmente, propomos o Agrupamento Contínuo em Profundidade, um novo paradigma promissor de inferência habilitado pelo Transformer Recursivo quando combinado com saídas antecipadas. Em uma análise teórica, mostramos que isso tem o potencial de levar a ganhos significativos (2-3x) na taxa de inferência.

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.