HGRN2: RNNs Lineares com Portas e Expansão de Estado

HGRN2: Gated Linear RNNs with State Expansion

April 11, 2024

Autores: Zhen Qin, Songlin Yang, Weixuan Sun, Xuyang Shen, Dong Li, Weigao Sun, Yiran Zhong

cs.AI

Resumo

A Rede Neural Recorrente Linear Hierarquicamente Controlada (HGRN, Qin et al. 2023) demonstrou velocidade de treinamento e desempenho competitivos em modelagem de linguagem, ao mesmo tempo em que oferece inferência eficiente. No entanto, o tamanho do estado recorrente da HGRN permanece relativamente pequeno, o que limita sua expressividade. Para resolver esse problema, inspirados pela atenção linear, introduzimos um mecanismo simples de expansão de estado baseado em produto externo, de modo que o tamanho do estado recorrente possa ser significativamente ampliado sem a introdução de parâmetros adicionais. A forma de atenção linear também permite um treinamento eficiente em termos de hardware. Nossos extensos experimentos verificam a vantagem da HGRN2 sobre a HGRN1 em modelagem de linguagem, classificação de imagens e no Long Range Arena. Nosso maior modelo HGRN2 de 3B supera ligeiramente o Mamba e o Transformer da Arquitetura LLaMa em modelagem de linguagem em um ambiente de experimento controlado; e desempenha-se de forma competitiva com muitos modelos de 3B de código aberto em avaliações downstream, enquanto utiliza um número muito menor de tokens de treinamento totais.

English

Hierarchically gated linear RNN (HGRN,Qin et al. 2023) has demonstrated competitive training speed and performance in language modeling, while offering efficient inference. However, the recurrent state size of HGRN remains relatively small, which limits its expressiveness.To address this issue, inspired by linear attention, we introduce a simple outer-product-based state expansion mechanism so that the recurrent state size can be significantly enlarged without introducing any additional parameters. The linear attention form also allows for hardware-efficient training.Our extensive experiments verify the advantage of HGRN2 over HGRN1 in language modeling, image classification, and Long Range Arena.Our largest 3B HGRN2 model slightly outperforms Mamba and LLaMa Architecture Transformer for language modeling in a controlled experiment setting; and performs competitively with many open-source 3B models in downstream evaluation while using much fewer total training tokens.