RotRNN: Modelling Long Sequences with Rotations
进 Q 交流群：922230617 或加 VX：CV_EDPJ 进 V 交流群
目录
0. 摘要
1. 引言
2. 背景
3. RotRNN
3.1. 归一化
3.2. 旋转参数化
4. 实验
5. 结论与未来工作
0. 摘要
线性递归模型，例如状态空间模型（State Space Models，SSM）和线性递归单元（Linear Recurrent Units，LRU），在长序列建模基准测试中最近表现出最先进的性能。尽管它们取得了成功，但也有一些缺点，最显著的是它们复杂的初始化和归一化方案。在这项工作中，我们通过提出 RotRNN 来解决其中的一些问题——一种利用旋转矩阵便利特性的线性递归模型。我们展示了 RotRNN 提供了一个比先前工作具有更少理论假设的简单模型，并且其实际实现忠实于其理论推导，在多个长序列建模数据集上取得了与 LRU 和 SSM 相当的成绩。
1. 引言
状态空间模型（SSM）[4, 5, 10] 在长序列建模任务中表现出了最先进的性能。SSM 利用线性递归层，避免了 Transformers [12] 的二次扩展问题以及经典递归神经网络（RNN）[2, 6, 7, 9] 的梯度消失和梯度爆炸问题。它们既可以作为快速推理的线性 RNN，也可以作为高效训练的卷积神经网络（CNN）。然而，无论哪种实现方式，都涉及到矩阵幂的计算，如果没有仔细的初始化和归一化，可能会导致数值不稳定的状态值。线性递归单元（LRU）[8] 则引入了一种由递归状态矩阵的特征分解驱动的对角矩阵参数化方法。尽管这使得初始化比之前的工作更简单，但归一化 LRU 层依赖于限制性假设，并且算法的实际实现偏离了理论推导。在这项工作中，我们提出了一种算法来解决 LRU 的一些问题。通过将状态矩阵参数化为旋转矩阵，我们能够在归一化过程中使用更少的限制性假设，并允许实现忠实于理论动机。我们在附录 C 中对我们的工作与之前的算法进行了更深入的比较。
2. 背景
状态空间模型（SSM）[4, 5, 10] 源自连续时间线性常微分方程，给定
其中 B′ ∈ R^(N×K) 是输入矩阵，A′ ∈ R^(N×N) 是状态矩阵，C ∈ R^(O×N) 是输出矩阵，u(t) ∈ R^K 是连续时间输入。由于这是一个线性时不变系统，它可以很容易地离散化并表示为离散递归
其中，u_t、A 和 B 是 u(t)、A' 和 B' 的离散版本。方程 1 是 S4 [4] 及其变体 [3, 5, 10] 中 SSM 层的基础，前提是 A 经过仔细初始化以确保数值稳定性。
线性递归单元（LRU）。LRU 不是离散化潜在的连续时间常微分方程，而是通过仔细初始化和归一化线性 RNN 来实现类似的性能。为了推导 LRU 参数化方案，第一步是对矩阵 A 进行特征值分解。由于 A ∈ R^(N×N) 是一个实数方阵，可以写成（经过任意小的扰动 [1]）：
其中 Λ 是特征值的对角矩阵，P ∈ C^(N×N) 是复数值的可逆特征向量矩阵。将方程 1 中的递归展开为 x_t =
两边同时左乘 P^(-1)，并代入方程 2，我们得到：
LRU 旨在直接学习矩阵 ˜B 和 ˜ C，以及特征值
对应可学习参数 γj, θj ∈ R。
然而，在实践中，LRU 偏离了这一理论。
首先，它不强制要求实矩阵 A 的特征值成共轭对 [1]。
此外，没有约束确保 P 和 P^(-1) 一致，这些矩阵仅被吸收到可学习的 ˜C 和 ˜B 中。
最后，LRU 状态归一化（见附录 C）仅在无限序列长度的情况下成立。
3. RotRNN
在我们的方法中，我们称之为旋转递归神经网络（Rotational Recurrent Neural Network，RotRNN），我们旨在通过将状态矩阵 A 参数化为旋转矩阵来克服 LRU 的缺点。选择这个方法的原因有三点：
旋转矩阵的特征值大小总是 1——这使得确保递归状态不爆炸或消失变得容易；
旋转矩阵可以从任何实值矩阵平滑生成（见 §3.2），对初始化具有鲁棒性；
旋转矩阵的正交性使我们可以推导出一个简单的归一化方案，而不需要无限序列长度的假设
3.1. 归一化
我们定义 RotRNN 的递归关系为：
其中 α 是归一化常数，γ ∈ [0, 1] 是可学习的标量衰减因子。写成方程 4 的卷积形式我们得到：
定理 1：根据 LRU [8] 的白噪声输入假设，u_t ∼ N(0, I) 独立同分布，那么如果
则对于所有时间步 t，我们有
我们在附录 A 中提供了主要文本中所有定理和引理的证明。注意，通过适当地缩放 SSM 的第一个状态，可以轻松地强制初始条件 ||x1|| = 1。仔细观察方程 5，归一化常数 α 被提升到 t + 1 − k 次方。因此，如果 α < 1，我们观察到序列中早期输入的权重趋于零，模型只关注最近的过去。由于我们希望递归衰减仅由 γ 控制，我们强制 α = 1，将所有归一化转移到矩阵 B。这需要 Tr[B^⊤·B] = 1 − γ^2，可以通过简单地用系数 ξ :=
重新缩放 B 来实现。因此，我们的最终表达式变为：
多头衰减（Multi-Head Decay）：我们发现单一标量衰减 γ ∈ [0, 1] 对于需要在不同尺度上保留信息的问题来说，泛化性不好。为了解决这个问题，我们将 K 维向量 u_t 输入到 H 个独立的 RotRNN 头中，每个头的输入矩阵 B_h ∈ R^(N/H×K)。最终投影 y_t = C·x_t，其中 x_t 是每个头的隐藏状态的连接，可以看作是信息混合线性层。
3.2. 旋转参数化
所有 R^(N×N) 旋转矩阵组成的群称为特殊正交群（special orthogonal group），SO(N)。要将任何矩阵 M ∈ R^(N×N) 平滑映射到 SO(N)，可以使用以下引理：
引理 2：设 S = M − M^⊤，则 A = exp(S) ∈ SO(N)，其中 exp 是矩阵指数，定义为
然而，计算方程 5 涉及对 A 取矩阵幂，这对于高维密集矩阵（dense matrices）来说很慢。为了提高这种计算的效率，我们使用以下引理将密集旋转矩阵分解为块对角形式。
引理 3：设 P ∈ O(N) 为正交矩阵，D ∈ SO(N) 为块对角旋转矩阵，则 A = PDP^⊤ ∈ SO(N)。
如果 N 是偶数，矩阵 D 在对角线上有 N/2 个形式为
的块；如果 N 是奇数，则有 (N−1)/2个块，剩余的对角线上值为 1。当计算矩阵幂时，密集正交矩阵相互抵消，留下 A^k = P·D^k·P^⊤，D^k 的块形式为
因此，通过直接学习 P 和 θ = {θ1, . . . , θm}，我们可以轻松生成适合快速计算矩阵幂的旋转矩阵。需要注意的是，尽管很难参数化 P 以使 P 的空间覆盖整个正交矩阵群，我们可以使用引理 2 获得 P ∈ SO(N) ⊂ O(N)。虽然这意味着 A = PDP^⊤ 不是SO(N)上的满射映射，但我们发现它在实践中足够通用，可以取得良好的结果（见 §4）。
长程竞技场（Long Range Arena，LRA）。我们在 LRA [11]上评估了 RotRNN 的性能，该数据集包括 6 个序列建模任务，序列长度在 1K 到 16K 个 token 之间，数据模态各异。表 1 显示了 RotRNN 的结果，通过报告具有最佳验证准确性的模型的测试准确性。这些实验的大多数超参数是手动选择的，仅在 ListOps 和 Text 上进行了超参数搜索（见附录 B）。因此，我们的初步结果显示，我们的模型与现有文献中的结果具有可比性。我们还注意到 Path-X 对超参数非常敏感，我们没有设法重现 LRU [8] 的性能，也没有让 RotRNN 在 Path-X 上的表现优于随机猜测。
核图（Kernel Plot）。为了了解 RotRNN 如何随着时间回忆长程信息，在图 1 中，我们绘制了在初始化时一个 16 头之一的单位脉冲下隐藏状态的演变。我们看到不同的通道对应于不同的正弦分量，正如预期的那样，随着序列长度的增加而衰减。
5. 结论与未来工作
在本文中，我们提出了 RotRNN，这是一种利用旋转矩阵便利特性的创新线性递归模型。我们展示了，尽管没有进行广泛的超参数搜索，RotRNN 在 LRA 基准测试中提供了与最先进水平相当的结果。未来的工作应侧重于探索 RotRNN 性能的极限，以及进行超参数搜索和消融研究，以更好地理解我们模型的优缺点。
论文地址：https://arxiv.org/abs/2407.07239
