HGRN2 : Réseaux Récurrents Linéaires à Porte avec Expansion d'État

HGRN2: Gated Linear RNNs with State Expansion

April 11, 2024

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

cs.AI

Résumé

Le RNN linéaire à portes hiérarchiques (HGRN, Qin et al. 2023) a démontré une vitesse d'entraînement et des performances compétitives en modélisation du langage, tout en offrant une inférence efficace. Cependant, la taille de l'état récurrent du HGRN reste relativement petite, ce qui limite son expressivité. Pour résoudre ce problème, inspirés par l'attention linéaire, nous introduisons un mécanisme simple d'expansion de l'état basé sur un produit extérieur, permettant ainsi d'augmenter significativement la taille de l'état récurrent sans introduire de paramètres supplémentaires. La forme de l'attention linéaire permet également un entraînement optimisé pour le matériel. Nos expériences approfondies confirment l'avantage du HGRN2 par rapport au HGRN1 en modélisation du langage, classification d'images et sur le Long Range Arena. Notre plus grand modèle HGRN2 de 3 milliards de paramètres surpasse légèrement Mamba et l'architecture Transformer LLaMa en modélisation du langage dans un cadre expérimental contrôlé ; et se montre compétitif avec de nombreux modèles open-source de 3 milliards de paramètres en évaluation en aval, tout en utilisant beaucoup moins de tokens d'entraînement au total.

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.