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, Цин и др., 2023) продемонстрировала конкурентоспособную скорость обучения и производительность в языковом моделировании, обеспечивая при этом эффективный вывод. Однако размер рекуррентного состояния HGRN остается относительно небольшим, что ограничивает его выразительность. Для решения этой проблемы, вдохновленные линейным вниманием, мы представляем простой механизм расширения состояния на основе внешнего произведения, благодаря которому размер рекуррентного состояния может значительно увеличиться без введения дополнительных параметров. Форма линейного внимания также позволяет эффективно использовать аппаратное обеспечение для обучения. Наши обширные эксперименты подтверждают преимущество HGRN2 перед HGRN1 в языковом моделировании, классификации изображений и в долгосрочной арене. Наша крупнейшая модель HGRN2 объемом 3 млрд нейронов незначительно превосходит Mamba и трансформерную архитектуру LLaMa в языковом моделировании в контролируемом экспериментальном окружении; и конкурентоспособна с многими открытыми моделями объемом 3 млрд в оценке на следующем уровне, используя при этом гораздо меньше общего количества обучающих токенов.

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.