本项目是一个专注于二阶优化动力学与信息几何的研究框架。它通过能量-几何解耦(Energy-Geometry Decoupling)原则,实现了在黎曼流形上高效滑行的优化范式。
ARS2-Neo 的核心设计基于对现代优化算法的深度重构,旨在克服一阶优化器在病态曲率地形下的局限性。
通过 Muon 的 Newton-Schulz 迭代,ARS2-Neo 强制更新矩阵保持正交性(Stiefel 流形约束)。在数学上,正交化更新等价于在参数空间执行去协关联,消除了内部协变量偏移,使梯度信息更加纯净。
对于任意正交矩阵 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 的类比也随之削弱。
虽然 NGD 提供了极速的收敛,但极易陷入“针尖极小值”(过拟合)。ARS2-Neo 引入了流形感知 SAM (Sharpness-Aware Minimization):
- 平坦度约束:通过在黎曼流形上寻找对抗方向,算法被引导至损失景观中更宽阔的盆地。
- MDL 对应:根据最小描述长度 (MDL) 原则,平坦的区域对应于更简单的模型解释,从而具备更强的泛化能力。
ARS2-Neo 将优化过程分解为两个独立的算子:
- 统计算子 (能量):利用 AdamW 的二阶矩修正动量范数确定更新步长,作为自由能下降速率的代理。
- 结构算子 (几何):通过预白化 (Pre-whitening) 与正交投影,确保更新方向严格遵循流形的测地线 (Geodesic)。
实验设置: 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 的全场最佳表现,证明了二阶几何信息在捕捉语义规律方面的代际优势。
实验设置: 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% | 纯几何优化,上限受限。 |
为了验证优化器在泛化相变(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 倍以上,有力证明了能量-几何解耦能避免模型在过拟合吸引盆中的无效游走,直接穿越高维峡谷抵达泛化解。
# 推荐使用 uv
uv sync# 运行 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- LRP/Main 实验:目录名为
outputs/lrp_*,用于主要对比结论。 - Verify/Smoke 实验:目录名为
outputs/verify_*,主要用于短程连通性验证(通常 1 epoch),不与长程 LRP 结果直接对比。
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}
}