让二阶优化像一阶一样好用，还支持灵活组合。
Second-Order, First-Class: A Composable Stack for Curvature-Aware Training
- 将曲率感知训练封装为可组合的单步操作，模块可替换。
- 实验表明组合方式显著影响训练速度和收敛效率。
- 提前规划执行路径，减少重复计算，提升运行效率。
二阶优化方法虽能提升稳定性和收敛速度，但因实现复杂、调参困难且缺乏可组合接口而未被广泛使用。我们提出 Somax，一个原生兼容 Optax 的可组合优化栈，将曲率感知训练简化为单一可即时编译的步骤，由静态计划驱动。Somax 提供一等公民模块：曲率算子、估计器、线性求解器、预条件器与阻尼策略，通过统一接口实现组合，并可与 Optax 的标准梯度变换（如动量、权重衰减、学习率调度）无缝衔接。该设计使原本隐含的选择显式化且可互换。Somax 实现了规划与执行分离：根据模块需求生成静态计划（含执行节奏），再通过专用执行路径复用中间结果。系统级消融实验表明：(i) 组合选择显著影响缩放行为与达到精度所需时间；(ii) 规划机制相比无规划组合减少了冗余重算，降低每步开销。
原文摘要 · Abstract (English)
Second-order methods promise improved stability and faster convergence, yet they remain underused due to implementation overhead, tuning brittleness, and the lack of composable APIs. We introduce Somax, a composable Optax-native stack that treats curvature-aware training as a single JIT-compiled step governed by a static plan. Somax exposes first-class modules -- curvature operators, estimators, linear solvers, preconditioners, and damping policies -- behind a single step interface and composes with Optax by applying standard gradient transformations (e.g., momentum, weight decay, schedules) to the computed direction. This design makes typically hidden choices explicit and swappable. Somax separates planning from execution: it derives a static plan (including cadences) from module requirements, then runs the step through a specialized execution path that reuses intermediate results across modules. We report system-oriented ablations showing that (i) composition choices materially affect scaling behavior and time-to-accuracy, and (ii) planning reduces per-step overhead relative to unplanned composition with redundant recomputation.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。