用于矩阵极分解近似的解耦多项式状态空间迭代及其在 Muon 中的应用
Decoupled Polynomial State-Space Iterations for Matrix Polar Approximation, with an Application to Muon
基于摘要分析。本文研究受限矩阵计算预算下的矩阵极分解近似,提出 Gram-matrix-based iterations(GMBI),通过解耦多项式状态空间参数化,分别控制方形辅助状态的传播与矩形输出的累积,并将其应用于 Muon 优化器。 方法建立在 Polar Express 的有限深度多项式框架上,支持紧凑的短计算图和稳健的完整重启计算图。对于重启计算图,作者冻结带保护机制的两层前缀,将 FP16 校准运行中的实际输出显式生成并用于重建真实 Gram 状态,再依据实际重启谱进行标量展开以优化后缀。该流程旨在适应前缀引起的分布变化与有限精度效应;摘要未说明保护机制及具体优化目标。 作者报告,在受控 FP16 实验中,无论采用仅面向稳健性的拟合协议还是 Muon-aware 拟合协议,解耦 GMBI 均优于学习得到的绑定多项式对照;Muon-aware 拟合还进一步改善了解耦版本的表现。作者报告,对未见谱、矩阵形状及留出的 GPT-2 Medium 更新的测试表明,收益不局限于校准集;重启对照实验还区分了有限精度 Gram 状态修复的贡献。摘要未提供具体误差指标或改善幅度。 作者报告,内核基准中,紧凑 GMBI 计算图在编译后端与专用后端均可降低大型矩形矩阵的极分解调用延迟。在采用 Muon 的多随机种子 GPT-2 Small 训练中,其最终验证损失与 Direct PE 相当,平均训练步耗时略有降低。这表明矩阵层面的收益可以延续至该端到端训练设置,但不能推广为任意模型上的加速结论。 这里的“状态空间”指矩阵迭代参数化,并非摘要已验证的神经信号序列建模方法。实验协议、消融及统计信息尚未验证;复现时需核对计算预算、矩阵形状、FP16 校准方式、拟合目标、后端配置与计时口径。代码和完整复现条件尚未验证,已有 EEG 相关应用尚未检索确认。
值得核对其解耦参数化与有限精度 Gram 状态修复各自的贡献,以及矩阵极分解近似的加速如何转化为 Muon 端到端训练收益;摘要中的验证对象是矩阵计算与 GPT-2 训练,尚不能据此推断 EEG/BCI 效果。
来源:OpenReview · State space models · openreview.net