Optax项目Polyak SGD优化器新增正部子最优间隙选项

在深度学习优化器库Optax中，Polyak SGD（随机梯度下降）是一种基于子最优间隙（suboptimality gap）的自适应学习率优化算法。近期社区讨论并通过了一项重要改进——为Polyak SGD增加仅使用子最优间隙正部的选项。

技术背景

Polyak SGD的核心思想是利用当前目标函数值f(x)与最优值f之间的差距f(x)-f来自适应调整学习率。传统实现中直接使用这个差值，无论其正负。但在实际应用中，当f(x)低于理论最优值f时（可能由于f估计不准确或噪声导致），会产生负的学习率调整，这可能引发数值不稳定问题。

改进方案

新增的plus选项将应用ramp函数处理子最优间隙，即只保留正部：(f(x)-f*)_+ = max{f(x)-f*, 0}。这种改进有以下优势：

1. 数值稳定性：避免负学习率调整
2. 理论保证：与Garrigos等人2023年提出的SPS+方法一致
3. 实用性：处理最优值估计不准确的情况

实现细节

在Optax的实现中，这个功能通过一个布尔参数plus控制。当设置为True时，优化器会在计算学习率前对子最优间隙应用max操作。这种实现保持了API的简洁性，同时提供了更强的鲁棒性。

应用场景

这项改进特别适用于：

• 最优值f*估计可能存在误差的情况
• 噪声较大的优化问题
• 需要更稳定训练过程的场景

总结

Optax项目对Polyak SGD的这项改进增强了算法的鲁棒性和实用性，使其能够更好地处理实际优化问题中的各种情况。这体现了开源社区持续优化算法实现、提升用户体验的努力。