HGRN2: Gegate Lineaire RNN's met Staatsexpansie

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

Samenvatting

Hiërarchisch gegate lineaire RNN (HGRN, Qin et al. 2023) heeft een concurrerende trainingssnelheid en prestaties aangetoond in taalmodelering, terwijl het efficiënte inferentie biedt. Echter, de grootte van de recurrente toestand van HGRN blijft relatief klein, wat de expressiviteit beperkt. Om dit probleem aan te pakken, geïnspireerd door lineaire aandacht, introduceren we een eenvoudig op outer-product gebaseerd mechanisme voor toestandsuitbreiding, zodat de grootte van de recurrente toestand aanzienlijk kan worden vergroot zonder extra parameters te introduceren. De lineaire aandachtvorm maakt ook hardware-efficiënte training mogelijk. Onze uitgebreide experimenten bevestigen het voordeel van HGRN2 ten opzichte van HGRN1 in taalmodelering, beeldclassificatie en de Long Range Arena. Ons grootste 3B HGRN2-model presteert lichtjes beter dan Mamba en LLaMa Architecture Transformer voor taalmodelering in een gecontroleerde experimentele setting; en presteert concurrerend met veel open-source 3B-modellen in downstream-evaluatie terwijl het veel minder totale trainings-tokens gebruikt.

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.