论文

Zeta:通过坐标自适应预处理进行矩阵优化的双重美白

Zeta: Dual Whitening for Matrix Optimization via Coordinate-Adaptive Preconditioning

模型训练预训练

摘要

大规模神经网络训练越来越依赖于矩阵感知优化器,这些优化器利用权重参数的结构,超越了元素自适应。然而,现有的矩阵感知方法(例如 Muon)有一个未被充分认识的漏洞:它们的核心操作 Newton-Schulz 迭代严重依赖于输入条件,但原始动量矩阵表现出严重的坐标尺度异质性。在本文中,我们首先通过卡方均匀性检验验证了这种尺度异质性,表明矩阵内尺度不平衡在 Transformer 层中普遍存在,并且坐标白化有效地纠正了它。受这一发现的启发,我们提出了 Zeta,一种双重白化优化器,它在严格排序的管道中应用坐标白化和光谱白化。排序不是可调的选择,而是遵循数学依赖性:坐标白化建立了光谱白化可靠运行所需的统计各向同性。我们进一步证明,通过改进输入的条件数,这种双管道相对于纯谱方法严格减少了正交化误差。根据经验,Zeta 在语言建模(0.6B 到 8B 参数)、混合专家架构和视觉任务方面匹配或超越了强大的基线,证明在正交化之前解决尺度不平衡可以带来更快的收敛和更好的泛化。代码可在 https://github.com/AIGCodeOS/aigcode_zeta_optimizer 获取。