HGRN2: RNN Lineari con Cancelli ed Espansione dello Stato

HGRN2: Gated Linear RNNs with State Expansion

April 11, 2024

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

cs.AI

Abstract

La RNN lineare con gate gerarchici (HGRN, Qin et al. 2023) ha dimostrato una velocità di addestramento e prestazioni competitive nella modellazione del linguaggio, offrendo al contempo un'inferenza efficiente. Tuttavia, la dimensione dello stato ricorrente di HGRN rimane relativamente piccola, il che ne limita l'espressività. Per affrontare questo problema, ispirati dall'attenzione lineare, introduciamo un semplice meccanismo di espansione dello stato basato sul prodotto esterno, in modo che la dimensione dello stato ricorrente possa essere significativamente ampliata senza introdurre parametri aggiuntivi. La forma dell'attenzione lineare consente inoltre un addestramento efficiente dal punto di vista hardware. I nostri ampi esperimenti verificano il vantaggio di HGRN2 rispetto a HGRN1 nella modellazione del linguaggio, nella classificazione delle immagini e nel Long Range Arena. Il nostro modello HGRN2 più grande da 3B supera leggermente Mamba e il Transformer dell'architettura LLaMa nella modellazione del linguaggio in un contesto sperimentale controllato; e si comporta in modo competitivo con molti modelli open-source da 3B nella valutazione downstream, pur utilizzando un numero molto inferiore di token totali di addestramento.

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.