RotRNN: Modelling Long Sequences with Rotations

技术摘要：RotRNN —— 利用旋转建模长序列

问题陈述
由于捕捉长程依赖关系的难度，长序列建模仍然是机器学习领域的一个挑战。Transformer 在计算复杂度上呈二次方增长且缺乏归纳偏置，而非线性 RNN 则受困于梯度消失或梯度爆炸问题。近期的状态空间模型（SSM）和线性递归单元（LRU）已取得了最先进的性能，但这些模型仍存在显著缺陷：

1. 复杂的初始化： 像 S4 这样的 SSM 依赖于复杂的理论推导（例如基于最优多项式投影的 HiPPO 初始化），这些在训练过程中并不易于理解。
2. 实现与理论的不匹配： LRU 虽然在概念上更简单，但在实践中并未完全体现其理论动机。具体而言，它无法强制实矩阵的特征值成共轭对，缺乏确保实值输出的约束，并且依赖于一个学习到的归一化因子，该因子无法保证训练过程中隐藏状态量级的连贯性。
3. 稳定性： 现有方法通常需要复杂的归一化方案或初始化程序，而这些程序在不同深度和随机种子下的鲁棒性较差。

方法论
作者提出了 RotRNN（旋转循环神经网络），这是一种将循环状态矩阵 $A$ 参数化为旋转矩阵的线性递归模型。该方法建立在三个核心支柱之上：

1. 通过旋转矩阵进行参数化：
   与其直接学习一个受限矩阵，RotRNN 通过反对称矩阵的矩阵指数来学习一个通用的权重矩阵 $M$，并将其映射到特殊正交群 $SO(N)$（即$A = \exp(M - M^\top)$）。这确保了$A$ 在整个训练过程中始终是一个有效的旋转矩阵。
   为了实现高效计算矩阵幂（这是线性递归卷积形式所必需的），作者将旋转矩阵 $A$ 分解为块对角形式：$A = P \Theta P^\top$。其中 $P$ 是一个正交矩阵，$\Theta$ 由定义为角度 $\theta_i$ 的 $2\times2$ 旋转块组成。这种结构允许在不进行昂贵的稠密矩阵运算的情况下，快速计算 $A^k$。

2. 原则性的归一化：
   利用旋转矩阵的正交性（其特征值位于单位圆上），作者推导出了一个归一化程序，该程序能在整个训练过程中保持恒定的期望隐藏状态范数（$E[\|x_t\|^2] = c$）。
   递归定义为 $x_t = \gamma A x_{t-1} + \xi B u_t$，其中 $\gamma$ 是可学习的标量衰减因子，$\xi$ 是为了满足恒定范数条件而推导出的重缩放系数。与 LRU 基于无限序列极限进行归一化不同，RotRNN 的归一化对所有时间步都成立，从而防止了当归一化常数被提升至 $t$ 次方时常见的早期输入指数级衰减现象。

3. 多头架构：
   为了解决单个标量衰减因子 $\gamma$ 的局限性（可能无法很好地泛化到需要不同记忆时长的任务），RotRNN 采用了 $H$ 个独立的低维头。每个头都拥有各自的旋转矩阵 $A^{(h)}$、输入矩阵 $B^{(h)}$ 和衰减因子 $\gamma^{(h)}$。输出经过拼接并进行线性混合，从而使模型能够捕捉不同尺度的依赖关系。

核心贡献

• 新型线性 RNN 层： 提出了一种将状态矩阵参数化为旋转的循环层，确保算法在整个训练过程中始终忠实于其理论动机。
• 高效计算： 通过将参数化旋转分解为块对角形式的方法，实现了高效的矩阵幂计算，从而实现快速线性递归。
• 鲁棒归一化： 推导出了一个原则性的归一化程序，该程序能保持恒定的期望隐藏状态量级，相比于依赖渐近收敛的先前工作，提供了更强的稳定性保证。
• 理论洞察： 推导出了 LRU 的一个特例（其特征值形成共轭对）与多头 RotRNN 之间的数学等价性，这可能有助于阐明 LRU 的内部机制。

实验结果
作者在多个长序列建模基准测试上评估了 RotRNN：

• 长程竞技场 (LRA)： RotRNN 在与最先进的线性递归模型（S4, S5, LRU）的对比中表现出了竞争力。在离散输入任务（如 ListOps, Text (IMDB) 和 Retrieval）中，它取得了基准模型中的最高分。然而，在像素级图像任务（Path-X, CIFAR）上，它的表现略低于某些基准模型。
• 语音命令 (Speech Commands)： 在 35 类语音命令分类任务中，RotRNN 以相似的参数量实现了与复现的 LRU 相同的性能（95.2% 准确率），并与更复杂的深度状态空间模型保持了竞争力。
• 稳定性分析： 通过对 ListOps 数据集上隐藏状态范数的经验分析显示，RotRNN 在训练、层深度和随机种子方面都能保持几乎恒定的量级。相比之下，LRU 的隐藏状态范数表现出剧烈波动，且需要更长时间才能收敛到稳定量级。

意义与主张
本文声称 RotRNN 为现有的线性递归模型提供了一个概念简单、高效且具有数学原则性的替代方案。其主要意义在于弥合了理论动机与实际实现之间的差距。通过使用旋转矩阵，RotRNN 避免了 SSM 复杂的初始化方案以及 LRU 的实现不一致性。作者断言，其模型提供了一种鲁棒的归一化程序，可以在无需依赖“黑盒”初始化的情况下确保稳定性，同时在标准基准测试中取得竞争力的表现。此外，这项工作表明，通过旋转块结构的视角来看待 LRU，有助于澄清先前模型中线性递归的机制。