[Paper] $U(d)$ 的子群诱导自然的 RNN 与 Transformer 架构

Source: arXiv - 2602.18417v1

Overview

本文提出了一种统一的方法来构建循环神经网络（RNN）和 Transformer 模型，其隐藏状态位于数学上行为良好的群上——具体而言是酉群 (U(d)) 的闭子群。通过将子群的选择（例如正交群 (O(d))）视为可插拔组件，作者推导出简洁的、可“直接替换”的架构，这些架构继承了诸如范数保持等理想的几何属性。在经典语言建模基准上的实验表明，当参数数量保持不变时，这些受群约束的模型能够匹配或超越标准基线。

关键贡献

• 统一的群论框架，同时适用于 RNN 和 transformer，基于一个共同的骨架，其中隐藏状态空间、切空间投影和更新规则由子群 (U(d)) 参数化。
• 针对正交群 (O(d)) 的具体实现，产生 正交状态 RNN 和 transformer，在训练过程中保持严格的范数不变。
• 切空间中的线性混合扩展，一种轻量级修改，适用于任何子群，并在参数受限的情况下提升性能。
• 实证验证，在 Tiny Shakespeare 和 Penn Treebank 上展示正交状态模型在参数相匹配的基线下实现竞争性的困惑度。
• 开源实现（随论文发布），让实践者无需重写模型代码即可切换子群。

方法论

1. 骨架公式化 – 作者从一组最小公理出发描述序列模型：隐藏状态流形 (\mathcal{M})、切空间投影 (\Pi) 和更新映射 (\Phi)。
2. 群替换 – 通过挑选闭子群 (G \subseteq U(d))（例如 (O(d)) 或完整酉群），令 (\mathcal{M}=G)。任一点的切空间是李代数 (\mathfrak{g})，(\Pi) 成为到 (\mathfrak{g}) 的正交投影。
3. RNN 模板 – 循环更新被表达为群乘法：
   [ h_{t+1}=h_t \exp\bigl(\Pi(\mathbf{W}x_t + \mathbf{U}h_t + b)\bigr), ]
   其中 (\exp) 为矩阵指数，将投影向量映射回群上。
4. Transformer 模板 – 将自注意力重新表述，使查询、键、值向量都是李代数的元素，注意力权重通过群作用（矩阵乘法）而非加性残差连接来应用。
5. 线性混合微调 – 不再直接将原始投影向量送入指数，而是与一个可学习的标量线性混合，实际上对切空间中的步长进行缩放。当模型规模受限时，这一简单改动可提升收敛性。

所有这些步骤均使用标准深度学习原语实现（矩阵乘法、用于重新正交化的 QR 分解等），因此可以轻松地嵌入现有的 PyTorch 或 JAX 代码库。

结果与发现

• Orthogonal‑state 模型在参数预算固定的情况下始终优于其无约束对应模型，验证了几何正则化的益处。
• 线性混合扩展在普通正交版本上带来约 2–3 % 的相对提升，尤其在较小的 Tiny Shakespeare 数据集上表现更佳。
• 训练稳定性提升：梯度保持良好尺度，得益于群的保持范数特性，模型出现梯度爆炸/消失的情况更少。

实际意义

• 即插即用的稳定性 – 开发者可以在现有 RNN 或 transformer 代码中，用正交（或其他子群）版本替换隐藏状态表示，而无需重新设计整个架构。这可以减少对梯度裁剪或学习率技巧的需求。
• 内存高效模型 – 由于群约束消除了额外正则化项（例如正交惩罚）的需求，您可以在参数更少的情况下实现相当的性能，这对边缘设备或对延迟敏感的服务非常有价值。
• 更好的长程建模 – 范数保持有助于在许多时间步上维持信息，使正交状态 RNN 在语音合成、时间序列预测或需要稳定隐藏动力学的强化学习代理等任务中具有吸引力。
• 可扩展到其他群 – 该框架并不限于 (O(d))；对复数网络、辛动力学或其他李群感兴趣的开发者可以通过替换为不同的子群进行实验，从而打开面向特定领域归纳偏置的大门（例如物理感知模型）。

限制与未来工作

• 计算开销 – 计算矩阵指数（或其近似）并在每一步后重新正交化，相比于普通 RNN 会增加一个适度的常数因子。
• 对超大模型的可扩展性 – 本文聚焦于 ≤ 2 M 参数的模型；对于生产环境中常见的数亿参数的 Transformer，尚不清楚该方法的表现如何。
• 子群探索受限 – 仅对正交群进行了实证评估；其他子群（例如酉群、特殊正交群）可能提供不同的权衡，但被留作后续研究。
• 切空间线性混合 – 虽然有效，但线性混合的启发式方法缺乏正式的理论依据；更深入的分析可能揭示最佳步长调度或自适应方案。

作者建议将该框架扩展到 结构化 群（例如块对角正交矩阵），并将其与现代训练技巧（如混合精度和梯度检查点）结合，以降低开销。

作者

• Joshua Nunley

论文信息

• arXiv ID: 2602.18417v1
• 分类: cs.LG, cs.CL
• 出版日期: 2026年2月20日
• PDF: Download PDF