HGRN2: Gated Linear RNNs with State Expansion

要約

階層的にゲートされた線形 RNN (HGRN、Qin et al. 2023) は、効率的な推論を提供しながら、言語モデリングにおいて競争力のあるトレーニング速度とパフォーマンスを実証しました。
ただし、HGRN の再帰状態サイズは比較的小さいままであり、その表現力が制限されています。この問題に対処するために、線形注意に触発されて、単純な外積ベースの状態拡張メカニズムを導入します。これにより、再帰状態サイズを大幅に拡大できます。
追加のパラメータを導入します。
リニア アテンション フォームにより、ハードウェア効率の高いトレーニングも可能になります。私たちの広範な実験により、言語モデリング、画像分類、および長距離アリーナにおいて、HGRN1 よりも HGRN2 の利点が検証されています。私たちの最大の 3B HGRN2 モデルは、言語モデリングに関して Mamba および LLaMa Architecture Transformer をわずかに上回ります。
制御された実験環境。
また、ダウンストリーム評価では多くのオープンソース 3B モデルと競合する性能を発揮しながら、使用するトレーニング トークンの総量ははるかに少なくなります。

要約(オリジナル)

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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google