Skip to content

Latest commit

 

History

History
118 lines (79 loc) · 6.61 KB

File metadata and controls

118 lines (79 loc) · 6.61 KB

ARS2-Neo: 沿着损失景观的测地线直接滑向全局最优解

License: AGPL v3 Python PyTorch Ask DeepWiki

本项目是一个专注于二阶优化动力学与信息几何的研究框架。它通过能量-几何解耦(Energy-Geometry Decoupling)原则,实现了在黎曼流形上高效滑行的优化范式。

1. 理论基础:从对角 Fisher 到满秩 NGD

ARS2-Neo 的核心设计基于对现代优化算法的深度重构,旨在克服一阶优化器在病态曲率地形下的局限性。

1.1 参数去协关联 (De-correlation)

通过 Muon 的 Newton-Schulz 迭代,ARS2-Neo 强制更新矩阵保持正交性(Stiefel 流形约束)。在数学上,正交化更新等价于在参数空间执行去协关联,消除了内部协变量偏移,使梯度信息更加纯净。

1.2 满秩 Fisher 近似与 NGD

对于任意正交矩阵 R 与对角矩阵 D,提升后的乘积 RDRᵀ 是一个满秩矩阵,其谱分布仍在旋转后的坐标系中——这种恒等式独立于曲率的来源。Adam 通常被解读为由梯度二阶矩构建的对角预条件,而 ARS2-Neo 将这一对角缩放与矩阵级正交化步骤(类似极分解的混合)复合。如果混合基 R 缓慢漂移并保持与曲率特征基的相关性,同时对角 D 跟踪对应的谱,那么 RDRᵀ 在原始坐标系下就可以被视为一个结构化的自然梯度预条件器,这与 Amari(1998)以及 K-FAC、Shampoo 等实践近似保持一致。

这种表述更适合作为一个经验可验证的假设而非数学恒等:当 ARS2-Neo 保持必要的对齐关系时,复合算子可以近似自然梯度下降(NGD),我们在 Wikitext-2 的训练(20 epochs 训练损失约 0.9)验证了这一强预条件化下降的效能。若正交化仅仅重塑奇异值而与曲率统计脱钩,那么提升后的 RDRᵀ 就会失去与真实 Fisher/Hessian 的联系,NGD 的类比也随之削弱。

1.3 全局最优与 MDL 原则

虽然 NGD 提供了极速的收敛,但极易陷入“针尖极小值”(过拟合)。ARS2-Neo 引入了流形感知 SAM (Sharpness-Aware Minimization)

  • 平坦度约束:通过在黎曼流形上寻找对抗方向,算法被引导至损失景观中更宽阔的盆地。
  • MDL 对应:根据最小描述长度 (MDL) 原则,平坦的区域对应于更简单的模型解释,从而具备更强的泛化能力。

2. 核心机制:能量-几何解耦

ARS2-Neo 将优化过程分解为两个独立的算子:

  1. 统计算子 (能量):利用 AdamW 的二阶矩修正动量范数确定更新步长,作为自由能下降速率的代理。
  2. 结构算子 (几何):通过预白化 (Pre-whitening) 与正交投影,确保更新方向严格遵循流形的测地线 (Geodesic)。

3. 关键实验结果 (LRP 验证)

3.1 Wikitext-2 语言建模

实验设置: Qwen3 (RoPE, 3-layer), Context 255. 旨在探测病态曲率流形上的优化稳定性。

优化器 Best PPL Last PPL 动力学特征 说明
AdamW 116.46 213.52 标准欧氏空间基准 缓慢收敛,后期过拟合
Muon 111.35 475.65 谱约束收敛 缺乏自适应能量,后期崩溃
ARS2-Neo (Base) 96.10 3055.47 过拟合 极速坠入针尖极小值,泛化崩溃
ARS2-Neo (Sync) 90.69 330.85 最优泛化上限 ρ=0.3, 成功抑制过拟合
ARS2-Neo (AGA) 93.23 414.83 效率与稳定性的折衷 λ=0.5, 自适应几何感知

核心洞察:ARS2-Neo (AGA) 仅需 3 个 Epoch 即可达到 93.23 PPL,远超 AdamW 的全场最佳表现,证明了二阶几何信息在捕捉语义规律方面的代际优势。

3.2 CIFAR-10 视觉分类

实验设置: ResNet-18, Batch Size 256.

优化器 Best Acc Final Acc 备注
ARS2-Neo (Sync) 95.87% 95.73% SOTA。60 Epoch 极速收敛。
AdamW 94.60% 94.47% 标准基准。
Muon 93.76% 93.69% 纯几何优化,上限受限。

3.3 Grokking 现象加速

为了验证优化器在泛化相变(Phase Transition)中的动力学特征,我们在模加法任务 (p=113, train_frac=0.3) 上对比了各优化器的表现。

优化器 拟合 (Epoch) 顿悟 (Epoch) 收敛 (Epoch) 状态
AdamW ~140 >600 N/A 严重泛化延迟,600 Epoch 未能实现顿悟。
Muon ~150 >400 N/A 纯几何优化在缺乏能量自适应时收敛极慢。
ARS2-Neo (Base) 20 180 250 极速 Grokking。能量-几何解耦显著加速相变。
ARS2-Neo (AGA) 20 150 200 最优动力学。自适应几何感知进一步缩短了泛化延迟。

核心洞察:ARS2-Neo 将 Grokking 发生时间提前了 4 倍以上,有力证明了能量-几何解耦能避免模型在过拟合吸引盆中的无效游走,直接穿越高维峡谷抵达泛化解。

4. 快速开始

4.1 安装

# 推荐使用 uv
uv sync

4.2 运行实验

# 运行 WikiText-2 同步模式 (最优泛化)
# 注意:实验目录为 `exp/wikitext-2`,因此使用脚本路径启动。
python exp/wikitext-2/train.py --config config/lrp_wikitext2_ars2_neo_sync_10e.toml

# 运行 CIFAR-10 AGA 模式 (高效收敛)
python -m exp.cifar.train --config config/lrp_cifar10_ars2_neo_aga_20e.toml

4.3 结果分层与解释口径

  • LRP/Main 实验:目录名为 outputs/lrp_*,用于主要对比结论。
  • Verify/Smoke 实验:目录名为 outputs/verify_*,主要用于短程连通性验证(通常 1 epoch),不与长程 LRP 结果直接对比。

5. 框架结构

  • optimizer/: 核心优化器实现,包括 ars2_neo.py
  • exp/: 原子化实验脚本,解耦数据流与模型逻辑。
  • model/: 包含 Qwen3 (RoPE) 与 ResNet 等标准研究模型。
  • config/: 基于 TOML 的实验配置管理。

引用

@software{ARS2_Neo_2025,
  author = {Rui, L.},
  title = {ARS2-Neo: Gliding Directly Towards Global Optima Along Geodesics of the Loss Landscape},
  year = {2026},
  publisher = {GitHub},
  url = {https://github.com/dmf-archive/ARS}
}