HGRN2: Gated Linear RNNs with State Expansion

TL;DR

HGRN2 introduces an outer product-based state expansion to Gated Linear RNNs, significantly increasing their expressiveness and efficiency without extra parameters, leading to improved language modeling performance.

Contribution

The paper proposes a novel state expansion mechanism for HGRN, enhancing its capacity and interpretability while maintaining parameter efficiency.

Findings

HGRN2 outperforms original HGRN across various benchmarks.

The state expansion enables linear attention interpretation.

HGRN2 achieves competitive results with other recurrent models.

Abstract

Hierarchically gated linear RNN (HGRN, \citealt{HGRN}) has demonstrated competitive training speed and performance in language modeling while offering efficient inference. However, the recurrent state size of HGRN remains relatively small, limiting its expressiveness. To address this issue, we introduce a simple outer product-based state expansion mechanism, which significantly enlarges the recurrent state size without introducing any additional parameters. This enhancement also provides a linear attention interpretation for HGRN2, enabling hardware-efficient training. Our extensive experiments verify the advantage of HGRN2 over HGRN consistently across different settings and competitive with other recurrent models.