HGRN2: RNNs Lineales con Compuertas y Expansión de Estado

HGRN2: Gated Linear RNNs with State Expansion

April 11, 2024

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

cs.AI

Resumen

La red recurrente lineal jerárquicamente controlada (HGRN, Qin et al. 2023) ha demostrado una velocidad de entrenamiento y un rendimiento competitivos en modelado de lenguaje, además de ofrecer una inferencia eficiente. Sin embargo, el tamaño del estado recurrente de HGRN sigue siendo relativamente pequeño, lo que limita su expresividad. Para abordar este problema, inspirados por la atención lineal, introducimos un mecanismo simple de expansión del estado basado en el producto externo, de modo que el tamaño del estado recurrente pueda ampliarse significativamente sin introducir parámetros adicionales. La forma de atención lineal también permite un entrenamiento eficiente en hardware. Nuestros extensos experimentos verifican la ventaja de HGRN2 sobre HGRN1 en modelado de lenguaje, clasificación de imágenes y Long Range Arena. Nuestro modelo HGRN2 más grande, de 3B, supera ligeramente a Mamba y al Transformer de arquitectura LLaMa en modelado de lenguaje en un entorno de experimentación controlada; y compite de manera competitiva con muchos modelos de código abierto de 3B en evaluaciones posteriores, utilizando muchos menos tokens de entrenamiento en total.

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.