【LLM】从零开始学大语言模型 | SSM篇 | 5 | SSM 优化(中):DPLR 结构与高效离散化
SSM 优化(中):DPLR 结构与高效离散化 上一篇介绍了 HiPPO 理论如何推导出理论上最优的 $\mathbf{A}$ 矩阵,同时也指出了两个工程问题:HiPPO-LegS 的时变性和密集矩阵 $O(N^3)$ 的计算代价。 本文将详细介绍 S4 如何通过三个关键技术解决这些问题:用 DPLR 结构将密集矩阵转化为稀疏表示,用双线性变换替代 ZOH 离散化以保持 DPLR 结构,以及用 Z 变换将卷积核序列的计算转化为矩阵求逆问题。 一、DPLR 结构:将密集矩阵变为稀疏表示 1. 问题的起点 HiPPO 推导出的 $\mathbf{A}$ 是一个 $N \times N$ 的密集下三角矩阵。如果直接使用它,离散化后的 $\bar{\mathbf{A}} = e^{\mathbf{A}\Delta}$ 和卷积核 $\bar{K}_i = \mathbf{C}\bar{\mathbf{A}}^i\bar{\mathbf{B}}$ 的计算都需要密集矩阵运算,复杂度为 $O(N^3)$。 我们需要一种特殊的矩阵结构,使得幂运算和矩阵求逆都能在 $O(N)$ 内完成。 2. 什么矩阵的幂运算最快 在高等代数中,对角矩阵的幂运算是最简单的。如果 $\mathbf{\Lambda} = \text{diag}(\lambda_0, \lambda_1, \ldots, \lambda_{N-1})$,那么 $\mathbf{\Lambda}^i = \text{diag}(\lambda_0^i, \lambda_1^i, \ldots, \lambda_{N-1}^i)$,复杂度仅为 $O(N)$。 但纯对角矩阵意味着各个状态维度之间完全独立,没有信息交互,无法表达复杂的序列模式。 2.1 S4 内部的状态组织方式:每个特征维度独立维护一个 $N$ 维状态 S4 并非只使用“一个 $N$ 维状态”,这一点在与 Transformer 的对比中非常关键。 ...