最近在调研长序列建模方案时，我一直被一个问题困扰：Transformer 的注意力机制虽然效果好，但序列一长，计算量就是二次方往上涨。后来接触到 Mamba 这类状态空间模型（State Space Model，SSM），训练速度和显存占用确实香，但在实际跑实验时又发现一个隐蔽的坑——参数矩阵的谱结构（Spectral Structure）对训练稳定性和长期记忆能力影响极大，而常见的 AdamW 在这种场景下并不总是最优解。
这篇文章围绕“Muon 优化器 + Mamba 模型 + 频谱优化”展开，先讲清状态空间模型和 Mamba 的核心原理，再拆解 Muon 优化器为什么适合这类模型，最后用纯 PyTorch 实现一个简化版 Mamba，并在合成任务上对比 Muon 与 AdamW 的训练效果。适合对序列模型有基础了解、想自己动手跑实验的读者。
1. 从状态空间模型到 Mamba：背景与核心概念
1.1 状态空间模型解决什么问题
状态空间模型并不是新概念，它在控制论、信号处理领域已经发展了很多年。经典连续状态空间模型可以写成：
dx/dt = A x + B u y = C x + D u其中 x 是隐状态，u 是输入，y 是输出，A、B、C、D 是参数矩阵。把它用在前馈网络里，可以理解为：模型维护一个隐状态 h_t，每个时间步根据当前输入更新状态，再映射成输出。整个过程是循环的，类似 RNN，但数学形式和参数化方式更优雅。
相比 Transformer，SSM 的核心优势是线性复杂度。处理长度为 L 的序列时，Transformer 的 attention 复杂度是 O(L^2)，而 SSM 可以做到 O(L)。这意味着当序列长度达到几万甚至几十万时，SSM 仍然能跑得动，而 Transformer 早就爆显存了。
1.2 Mamba：选择性扫描与硬件感知
Mamba 是 2023 年底提出的状态空间模型架构，全称是Mamba: Linear-Time Sequence Modeling with Selective State Spaces。它在 S4 等早期 SSM 工作的基础上做了两个关键改进：
- 参数不再固定不变，而是依赖输入。A、B、C 这些矩阵会针对每个 token 动态调整，这就是“选择性”（Selective）的含义。模型可以根据当前输入决定“记住什么、忽略什么”，有点像 attention 的动态权重，但计算复杂度仍然是线性的。
- 硬件感知扫描。SSM 的递归计算可以看成一种“线性递归”，Mamba 用并行扫描（parallel scan）算法在 GPU 上高效实现，比逐时间步循环快得多。
Mamba 最早用于语言建模，后来很快扩展到视觉（Vision Mamba）、医学影像重建（例如 MRI 重建中的 LMO，Linear Mamba Operator）等方向。可以说，状态空间模型已经成了除 Transformer 之外最受关注的序列建模基座之一。
1.3 Muon 优化器与频谱优化：为什么要组合
Muon 是一个相对较新的优化器，核心思路是：对二维及以上的参数矩阵，先用动量累积梯度，再对梯度矩阵做正交化（Orthogonalization），最后按 SGD 方式更新。正交化这一步用的是 Newton-Schulz 迭代，本质上是计算梯度矩阵极分解中的正交因子。
为什么要做正交化？因为梯度矩阵可以分解为旋转部分和缩放部分。SGD 直接把缩放信息带给参数，如果梯度矩阵的奇异值分布很病态，优化路径就容易震荡。Muon 把缩放信息去掉、只保留旋转方向，相当于对梯度做了“白化”，在深层网络和循环结构上往往更稳。
2. 环境准备与版本说明
2.1 运行环境与依赖
本文的实战部分全部用纯 PyTorch 实现，不依赖官方 Mamba 的 CUDA 扩展，CPU 也能跑。推荐环境如下：
- Python 3.9 或 3.10
- PyTorch 2.0 以上（示例代码用到了
addcmul_、torch.linalg.svdvals等 API，旧版本可能不兼容） - 无需 GPU，CPU 版本即可跑通全部示例
版本需要根据你的项目实际情况调整，本文示例以常见环境为例，重点演示配置思路。
创建环境：
conda create -n ssm-muon python=3.10 -y conda activate ssm-muon pip install torch如果你有 NVIDIA GPU，并且想跑官方 Mamba 实现，可以额外安装：
pip install causal-conv1d pip install mamba-ssm需要特别提醒，mamba-ssm是 Linux + CUDA 环境下的编译扩展，Windows 支持很有限。如果你用的是 Windows 11，更推荐先用 WSL2 或 Linux 环境，或者在本文的纯 PyTorch 实现上跑通思路后再迁移。
另外，注意区分两个名字相似的东西：mamba-ssm是这里的模型库，而 conda 生态里的mamba是一个包管理器。搜索资料时经常看到“mamba 安装”的教程，先确认是哪一种，避免装错。
2.2 项目目录结构
实战代码按下面结构组织：
ssm-muon/ ├── model.py # 简化 Mamba 模型 ├── optimizer.py # Muon 优化器实现 └── train.py # 数据生成与训练对比3. 核心原理拆解：频谱视角下的 SSM 优化
3.1 SSM 的离散化与矩阵指数
连续 SSM 要放进神经网络，必须离散化。常用方法是零阶保持（Zero-Order Hold，ZOH），给定采样步长 Δ，离散化后的状态递推为：
h_t = A_bar h_{t-1} + B_bar x_t A_bar = exp(Δ A) B_bar = Δ B其中exp(Δ A)是矩阵指数。在 Mamba 的实现中，A 通常初始化为接近 1 的负对角线值，这样A_bar的谱半径会略小于 1，系统在长期递推时既不会发散也不会瞬间遗忘。
关键点来了：A_bar的谱性质直接决定模型的记忆长度。谱半径越接近 1，历史信息的“衰减”越慢，模型能记住的时间步越长；谱半径大于 1，递推过程数值上容易爆炸；远小于 1，模型退化为只看最近几步的短程模型。这就是频谱优化对 SSM 如此重要的原因。
3.2 谱半径、谱范数与长期记忆
需要区分两个概念：
- 谱半径（Spectral Radius）：矩阵特征值的最大模，只对方阵有意义，决定了递推动力系统的稳定性。
- 谱范数（Spectral Norm）：矩阵的最大奇异值，即
σ_max(M)，对任意形状矩阵都有定义，衡量矩阵作为线性变换的“最大放大倍数”。
在 Mamba 中，A 参数通常组织成(d_inner, d_state)，每个通道对应一个对角矩阵，因此谱半径可以直接取 A 中元素绝对值的最大值。离散化后，A_bar的元素都在 (0, 1) 区间，越接近 1 表示该通道的记忆保持能力越强。
训练过程中如果不加约束，A 的谱结构可能漂移，导致梯度消失或爆炸。常见的处理方式包括：
- 限制 A 的初始化范围；
- 对 A 做谱裁剪；
- 在损失函数里加谱正则项；
- 选择对病态梯度更
