在Adam优化器近乎统治深度学习的当下，二阶优化器正凭借其理论上的收敛优势重回视野。谷歌早在2022年就成功将二阶优化器Shampoo应用于千亿级广告模型，其工程实践方案极具参考价值，为解决大规模模型训练效率瓶颈提供了新思路。
智能速览
二阶优化器利用曲率信息，理论上比Adam等一阶方法收敛更快。
谷歌通过异构计算，将Shampoo的重负载从TPU转移到CPU，避免阻塞主训练。
提出“学习率嫁接”技巧，结合Shampoo的更新方向和AdaGrad的步长大小。
该方案在千亿参数模型上，实现了AUC提升0.44%，而训练时间仅增加10%。
大模型对效率的极致渴求，正重新点燃对二阶优化的研究热情。
精华内容
将二阶优化器落地工业界并非易事，谷歌通过一系列精密的工程设计，巧妙地平衡了计算开销与模型性能，其核心创新点值得我们深入拆解。
为何选择二阶
在超大规模模型训练中，传统的SGD、Adam等一阶优化方法仅利用梯度信息，往往存在收敛速度慢的问题。二阶优化方法通过引入Hessian矩阵（即二阶导数或曲率信息），理论上能够实现更快的收敛速度和更精准的参数更新方向。
然而，理想的丰满无法掩盖现实的骨感。直接计算并存储Hessian矩阵的逆矩阵，其计算复杂度高达O(d^3)，其中d为模型参数维度。对于千亿参数级别的模型，这样的计算成本是完全无法接受的，这也是二阶优化器长期以来“叫好不叫座”的根本原因。
异构计算之妙
谷歌的创新核心在于一套精巧的异构计算架构，它将训练任务进行了智能化拆分。模型的梯度计算与反向传播等常规操作，仍在高性能的TPU上进行，以发挥其矩阵运算优势。
而Shampoo算法中最繁重、最耗时的部分——预条件的逆P次根计算，则被异步地流水线式转移到CPU集群上处理。这种设计使得TPU的主训练流程几乎不会因为预条件的计算而阻塞，从而巧妙地规避了二阶计算带来的巨大延迟，为整个系统的流畅运行提供了保障。
学习率嫁接术
在在线学习场景中，数据流持续不断，不存在固定的Epoch，因此无法使用Cosine Decay等预设的学习率衰减策略。为了解决这个问题，谷歌论文中提出了一个名为“Learning Rate Grafting”的精妙技巧。
该方法将优化过程一分为二：使用Shampoo来计算参数更新的精准方向（一个向量），但步长的大小（一个标量）则交由适应性更强的AdaGrad来决定。这种“移花接木”的方式，既保留了二阶方法在方向上的优越性，又继承了AdaGrad对稀疏特征和在线数据环境的良好适应性，极具工程智慧。
精度的代价与回报
经过一系列硬核的工程改造，分布式Shampoo的效果在工业级应用中得到了验证。实验数据显示，相较于AdaGrad、Adam、Yogi等主流一阶优化器，Shampoo为CTR模型带来了0.44%的PerQueryAUC提升。在广告推荐这类业务中，微小的AUC提升意味着巨大的商业价值。
更令人惊喜的是成本。尽管引入了复杂的二阶计算，但由于异步异构设计，端到端的训练时间仅仅增加了10%。同时，在相同的训练步数下，Shampoo的Loss下降速度显著快于其他优化器，证明了其高效的收敛能力。
谷歌的实践证明，二阶优化器并非空中楼阁，而是能在工业环境中创造真实价值的技术。随着大模型对训练效率的要求越来越高，这些曾被认为成本高昂的“旧技术”正焕发新生，未来的优化器发展值得持续关注。
