状态范数崩溃：WDLM-60M 外推失效的根因分析与修复

摘要本文针对 WDLM-60M 模型在序列长度超过训练长度1024后困惑度PPL指数级崩溃的现象进行了系统分析。通过逐层追踪 cummax 状态范数的变化定位到深层L7、L8状态值爆炸是直接原因。进一步对比多头 cummax 架构OpenASH发现全维度512cummax 无法在有限序列内饱和而多头设计8×80可在约 1000 token 后自然饱和。消融实验表明对状态范数进行截断norm cap可将 16384 长度下的 PPL 从 33286 降至 11.3且修复后模型具备无限外推能力128K token 时 PPL 仍为 11.4。然而推理时后处理会严重损害训练长度内的生成质量PPL 上升 755%。通过在训练阶段集成状态范数约束模型收敛几乎不受影响loss 差异 1%短序列 PPL 反而改善且推理时无需额外干预。本文揭示了 cummax 状态范数失控的本质并提出了一种简单有效的训练时正则化方案可推广至同类架构。关键词外推状态爆炸cummax语言模型长度泛化1 引言近年来基于线性注意力或状态空间模型的语言模型因其线性计算复杂度而受到广泛关注。WDLM-60M 和 OpenASHOA系列模型采用了 cummax 机制来维护一个单调递增的状态向量从而实现对长序列的高效建模。然而WDLM-60M 在评估时一旦序列长度超过训练长度1024其困惑度便会急剧上升在 4096 长度时 PPL 达到 2946而参数量相近的 OA-58M 仅退化至 90。这种极端的不稳定性严重限制了模型在实际长文本场景中的应用。本文旨在回答以下三个问题哪一层或哪些层导致了 WDLM 的外推崩溃为什么 OA 系列能够保持相对稳定是否存在一种通用且低代价的修复方案既能恢复外推能力又不损害短序列性能通过对每一层 cummax 状态的范数、最大值进行逐块chunk追踪并结合消融实验clamp、norm cap、skip 等我们定位到深层 L7 和 L8 是爆炸的主要源头。进一步分析表明全维度 cummax 的状态空间随序列长度指数增长而多头 cummax 由于每个头的维度较小80能够快速达到饱和从而抑制了状态膨胀。基于以上发现我们提出在训练过程中对每一层的 cummax 状态施加范数截断。实验证明该方案不仅不损害训练收敛反而起到正则化作用使模型在保持短序列性能的同时获得了近乎无限的外推能力128K token 时 PPL 仅 11.4。本文的工作为设计长上下文线性语言模型提供了重要的实践指导。2 背景与相关工作2.1 线性注意力与状态空间模型传统 Transformer 的自注意力复杂度为 O(L²)限制了其处理极长序列的能力。线性注意力通过引入核函数或状态变量将复杂度降至 O(L)。其中cummax 机制是一种特殊的递归状态更新对每一维特征状态值随序列单调非减便于保留长期依赖。2.2 WDLM 与 OpenASH 架构差异模型cummax 方式状态维度头数WDLM-60M全维度单头5121OpenASH-58M多头8 × 80 6408OA 的每个头在仅 80 维的子空间上独立进行 cummax这种分解显著限制了单一状态向量的表达能力上限从而促使状态在约 1000 token 后趋于饱和。WDLM 则在一个 512 维的联合空间上累积理论上可区分的 cummax 路径数高达 2^512因此状态值会随着序列长度无界增长。2.3 外推问题与状态爆炸外推extrapolation指模型在超过训练长度的序列上的表现。已有工作表明递归或状态变量模型容易在长序列上出现数值不稳定表现为状态范数指数级增长。本文是首次系统定位 WDLM 具体爆炸层并提出训练时正则化方案的实证研究。3 方法3.1 逐层状态追踪我们使用长度为 64 的 chunk 逐步处理长序列256 至 8192并在每个 chunk 后记录每层输出后的状态向量的 L2 范数state norm状态向量各元素绝对值的最大值state max层输出的 L2 范数所有测试均在固定模型权重下进行不更新梯度。3.2 消融实验设计我们设计四种干预方式分别应用于推理阶段以评估其对 PPL 的影响Clamp将状态向量每个元素截断至 [-max, max]。Norm cap若状态范数超过阈值则缩放至该阈值保持方向。Skip跳过某一层的状态更新直接传递输入。Freeze保持某一层的状态不变不进行 cummax 更新。3.3 训练时集成状态约束为验证训练阶段引入约束的可行性我们使用 200 步的快速训练实验对比 baseline 与每一步后对状态进行范数截断cap200的模型。评估包括训练 loss、短序列 PPL256-1024以及外推至 64K 的表现。4 实验结果与分析4.1 WDLM 状态爆炸的层定位表 1 展示了 WDLM 各层在 256 与 4096 长度下状态范数的增长倍数。层State norm (256)State norm (4096)增长倍数State max (4096)L02474021.6×47L122419348.6×515L21845983.3×124L61298186.3×254L7277481717.4×596L88711191413.7×1563L7 和 L8 的范数增长最为剧烈L8 的最大值达到 1563是普通层的 30–60 倍。这说明爆炸主要发生在深层并且逐层累积放大。作为对比OA-58M 的后四层L6–L9在 256 到 4096 长度下范数增长仅为 1.1–1.2 倍最大值不超过 44表明多头设计有效抑制了状态膨胀。4.2 状态饱和行为差异我们进一步追踪了状态范数随序列长度的变化曲线。OA-58M 的各层范数在 1024 token 后迅速饱和例如 L6 在 768 长度后便稳定在 149 左右。而 WDLM 的 L1、L7、L8 层范数在整个 8192 长度范围内持续线性增长从未出现平台期。这印证了全维度 cummax 无法在有限序列内达到饱和的理论分析。4.3 消融实验单层 Clamp 的效果在 4096 长度下对单层状态进行 clamp最大值 50并测量 PPL 变化干预PPL相对 baselinebaseline2946.51.00×clamp L02946.51.00×clamp L13084.21.05×clamp L42506.60.85×clamp L71702.60.58×clamp L82698.90.92×clamp L7L81473.90.50×clamp ALL957.50.32×单层 clamp L7 即可降低 42% 的 PPL说明 L7 是最关键的爆炸源。同时 clamp L7 与 L8 可将 PPL 再降低一半而全层 clamp 则能下降 68%。4.4 全层 Norm Cap 的强大效果对所有层的状态范数施加统一上限max_norm200后PPL 从 2946 骤降至 11.2表 2。更激进的约束clamp20 或 norm200甚至可以将 PPL 降至 11 左右。这表明状态范数的失控是外推崩溃的根本原因而非架构本身的缺陷。干预方式PPL 4096相对 baselinebaseline2946.5—clamp50 all957.5-68%clamp20 L1-8176.6-94%norm200 all11.2-99.6%4.5 修复后的外推极限测试我们将 norm200 的修复方案称为 WM-fix与 OA-58M、OA-85M 在长度 256 至 16384 上进行对比表 3。SeqWM-baseWM-fixOA-58MOA-58M-fix10245.915.69.111.440962946.511.290.413.9163843328611.3374.314.732768inf11.4525.714.965536inf11.4660.114.9131072inf11.4774.215.0WM-fix 在 128K训练长度的 125 倍时 PPL 仍为 11.4与 4K 时完全一致展示了无限外推能力。OA-58M 在修复后也大幅改善PPL 从 774 降至 15.0。值得注意的是修复后的 WM-fix 甚至优于 OA-58M-fix说明全维度 cummax 在适当约束下可以比多头设计获得更低的困惑度。4.6 推理时后处理的副作用尽管推理时加 state cap 能完美修复外推但它严重损害了训练长度内的生成质量。表 4 展示了在 512 长度的 SFT 样本上的结果模型PPL (512)生成示例冒泡排序WM-base4.52正确输出算法WM-fix38.62 (755%)乱码、空回答OA85-base2.91正确输出OA85-fix4.53 (56%)勉强可读但质量下降WM-fix 在短序列上几乎完全不可用。这是因为模型从未在状态被截断的条件下训练其内部表征分布与推理时的约束不匹配。因此推理时后处理不是一个可部署的方案。4.7 训练时集成状态约束为解决上述矛盾我们在训练阶段即对每一层的状态施加范数截断max_norm200并与 baseline 进行 200 步的对比训练。StepBaseline LossWith Cap Loss差异207.267.331.0%1004.074.070%2003.693.720.8%训练 loss 几乎无差异说明状态约束不会干扰模型收敛。更令人意外的是在短序列评估512中带 cap 训练的模型 PPL 为 33.41反而低于 baseline 的 34.87降低了 4.2%表明该约束起到了正则化作用。最重要的是该模型在推理时无需任何额外干预即可同时保证短序列质量和长序列外推能力64K 时 PPL46.3而 baseline 为 44.6两者基本持平考虑到训练步数很少两者尚未充分收敛但趋势良好。5 讨论5.1 为什么 cummax 会导致状态爆炸cummax 操作满足单调非减性即每个状态元素只能增加或保持不变。在全维度联合空间中没有机制迫使某一维的值停止增长。随着序列增长某些维度的值会累积到极大破坏后续层的数值稳定性。5.2 多头设计的本质优势多头 cummax 将高维空间分解为多个低维子空间。每个子空间的维度较小例如 80在有限 token 内约 1000所有维度都会被“激活”并达到饱和。饱和后新的 token 无法再增加状态值从而自然截断了增长。这与全维度设计中状态永远无法饱和形成鲜明对比。5.3 状态范数约束为何有效范数截断相当于给状态向量的能量设定一个硬上限。它强迫模型在受限的表示空间内学习类似于梯度裁剪对训练稳定性的作用。由于训练时即施加该约束模型能够适应这种受限表示并在推理时自然保持稳定。5.4 通用性与局限性本文提出的训练时 state norm cap 方案对 WDLM 和 OA 系列均有效。但我们发现不同模型的最优 cap 阈值不同WDLM 适合 200OA-58M 适合 50OA-85M 适合 150且没有任何固定阈值可以同时适配所有模型。这意味着实际应用中需要对每个模型进行小范围超参数扫描或在训练时将 cap 作为可学习参数的一部分。另外本文的实验局限于 60M 参数量的模型是否适用于更大规模的模型如 1B仍需进一步验证。6 结论本文对 WDLM-60M 的外推崩溃问题进行了根因分析得出以下结论爆炸源定位深层 L7 和 L8 的状态范数在长序列下分别增长 17.4 倍和 13.7 倍是导致 PPL 指数上升的直接原因。架构本质全维度 cummax 无法在有限序列内饱和而多头 cummax 通过低维子空间实现了自然饱和。修复方案对状态向量施加范数截断可以完美恢复外推能力。但推理时后处理会严重损害短序列性能训练时集成该约束是唯一可行的方案。效果训练时加入 state norm cap 不影响收敛loss 差异 1%短序列 PPL 反而降低 4.2%且模型获得了无限外推能力128K token 时 PPL11.4。我们的工作揭示了递归状态模型中一个容易被忽视的数值稳定性问题并为未来设计长上下文线性语言模型提供了实用的工程建议。代码与实验脚本已开源链接略。致谢感谢开源社区提供的 WDLM 和 OpenASH 模型实现以及 Hugging Face 的评估工具。参考文献[1] WDLM: Wavelet-based Dual Linear Model for Efficient Language Modeling. (Anonymous, 2025)[2] OpenASH: A Simple and Efficient Linear Attention Architecture. (Anonymous, 2025)[3] Katharopoulos, A., et al. “Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention.” ICML 2020.[4] Sun, Y., et al. “RetNet: Retentive Network for Long Sequence Modeling.” arXiv 2023.[5] Gu, A., Dao, T. “Mamba: Linear-Time Sequence Modeling with Selective State Spaces.” arXiv 2023.附录文中所有数据均来自实际运行脚本analyze_layer_state.py,ablation_state.py,bench_fixed_extrap.py等。实验环境为单张 A100 80GB GPUPyTorch 2.1CUDA 12.1。项目地址