HGRN2: Gated Lineare RNNs mit Zustandserweiterung

HGRN2: Gated Linear RNNs with State Expansion

April 11, 2024

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

cs.AI

Zusammenfassung

Das hierarchisch gesteuerte lineare RNN (HGRN, Qin et al. 2023) hat eine wettbewerbsfähige Trainingsschnelligkeit und Leistung bei der Sprachmodellierung gezeigt, während es eine effiziente Inferenz bietet. Allerdings bleibt die rekurrente Zustandsgröße des HGRN relativ klein, was seine Ausdruckskraft einschränkt. Um dieses Problem anzugehen, führen wir inspiriert von linearer Aufmerksamkeit einen einfachen Mechanismus zur Zustandserweiterung auf Basis von äußeren Produkten ein, so dass die rekurrente Zustandsgröße signifikant vergrößert werden kann, ohne zusätzliche Parameter einzuführen. Die Form der linearen Aufmerksamkeit ermöglicht auch ein hardwareeffizientes Training. Unsere umfangreichen Experimente bestätigen den Vorteil von HGRN2 gegenüber HGRN1 bei der Sprachmodellierung, der Bildklassifizierung und der Long Range Arena. Unser größtes 3B HGRN2-Modell übertrifft leicht Mamba und die LLaMa-Architektur-Transformer bei der Sprachmodellierung in einem kontrollierten Experiment; und konkurriert wettbewerbsfähig mit vielen Open-Source-3B-Modellen bei der nachgelagerten Evaluation, während es deutlich weniger Gesamttrainings-Token verwendet.

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.