HGRN2: 상태 확장을 적용한 게이트 선형 순환 신경망

HGRN2: Gated Linear RNNs with State Expansion

April 11, 2024

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

cs.AI

초록

계층적 게이트 선형 RNN(HGRN, Qin et al. 2023)은 언어 모델링에서 경쟁력 있는 학습 속도와 성능을 보여주며, 효율적인 추론을 제공합니다. 그러나 HGRN의 순환 상태 크기는 상대적으로 작아 표현력이 제한됩니다. 이 문제를 해결하기 위해 선형 어텐션에서 영감을 받아, 추가 매개변수를 도입하지 않고도 순환 상태 크기를 크게 확장할 수 있는 간단한 외적 기반 상태 확장 메커니즘을 제안합니다. 선형 어텐션 형태는 하드웨어 효율적인 학습도 가능하게 합니다. 우리의 광범위한 실험을 통해 HGRN2가 HGRN1보다 언어 모델링, 이미지 분류, Long Range Arena에서 우수함을 검증했습니다. 가장 큰 3B HGRN2 모델은 통제된 실험 설정에서 언어 모델링에 있어 Mamba 및 LLaMa 아키텍처 트랜스포머를 약간 능가하며, 다운스트림 평가에서 많은 오픈소스 3B 모델들과 경쟁력을 보이면서도 훨씬 적은 총 학습 토큰을 사용합니다.

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.