Resurrecting Recurrent Neural Networks for Long Sequences
Antonio Orvieto ETH Zurich Work done at DeepMind. Samuel L Smith DeepMind Albert Gu DeepMind Anushan Fernando DeepMind Caglar Gulcehre DeepMind Razvan Pascanu DeepMind Soham De DeepMind
Abstract
Recurrent Neural Networks (RNNs) offer fast inference on long sequences but are hard to optimize and slow to train. Deep state-space models (SSMs) have recently been shown to perform remarkably well on long sequence modeling tasks, and have the added benefits of fast parallelizable training and RNN-like fast inference. However, while SSMs are superficially similar to RNNs, there are important differences that make it unclear where their performance boost over RNNs comes from. In this paper, we show that careful design of deep RNNs using standard signal propagation arguments can recover the impressive performance of deep SSMs on long-range reasoning tasks, while also matching their training speed. To achieve this, we analyze and ablate a series of changes to standard RNNs including linearizing and diagonalizing the recurrence, using better parameterizations and initializations, and ensuring proper normalization of the forward pass. Our results provide new insights on the origins of the impressive performance of deep SSMs, while also introducing an RNN block called the Linear Recurrent Unit that matches both their performance on the Long Range Arena benchmark and their computational
中文速览
循环神经网络(RNN)在长序列任务上推理快但训练难,而新兴的深度状态空间模型(SSM,如S4)既能高效并行训练、又能像RNN一样快速推理,在长程推理基准上表现亮眼——但究竟是哪些设计让SSM如此出众,一直不清楚。本文系统地对普通深度RNN做了一系列改造:去掉递归中的非线性激活、将递归矩阵对角复数化、用指数参数化来稳定特征值并调整初始化分布,以及在前向传播中对隐状态做归一化,最终得到一个名为"线性递归单元"(Linear Recurrent Unit,LRU)的新RNN模块。在Long Range Arena长程推理基准上,LRU的准确率和训练速度均与S4族模型持平,而每一步改造的消融实验也揭示出SSM性能提升的真正来源——关键不是HiPPO理论给出的特定确定性初始化,而是递归层特征值在初始化时的分布以及前向传播的归一化。这一发现不仅为理解SSM的工作机制提供了新视角,也为未来设计更简洁、更高效的序列模型提供了实用指导。
原文 arXiv:2303.06349;中英对照 + 大白话阅读 https://aha.fim.ai/paper/2303.06349v1