清华团队提出RL专用神经网络优化器，性能位居榜首

原创已于 2024-12-20 21:34:52 修改 · 908 阅读

11 ·

CC 4.0 BY-SA版权

文章标签：

#神经网络 #人工智能 #深度学习 #机器学习

于 2024-12-20 21:18:38 首次发布

在人工智能领域，强化学习（RL）一直被视为解决复杂序列决策问题的有力工具。这项技术不仅在电子游戏、棋类智能、机器人控制、自动驾驶等前沿领域大放异彩，还在大语言模型（LLM）的微调、对齐、推理等关键阶段发挥着重要作用。

然而，在深度神经网络作为函数近似器的背景下，RL训练常常表现出显著的不稳定性，面临诸如过拟合、值高估和策略发散等一系列挑战。这些问题不仅影响了RL算法的最终性能，也限制了其在真实物理场景中大规模应用的潜力。

面对这一挑战，清华大学的研究团队提出了RL专用的神经网络优化方法——RAD优化器（Relativistic Adaptive gradient Descent）。相较于9个主流神经网络优化器（包括SGD-M，Adam和AdamW等），RAD优化器在12个测试环境（包括1个CartPole任务、6个MuJoCo任务、4个Atari任务和1个自动驾驶任务）及5种主流RL算法（包括DQN、DDPG、TD3、SAC和ADP）下，综合性能均排名第一。特别在图像类标准测试环境Seaquest任务中，RAD性能达到Adam优化器的2.5倍，得分提升了155.1%。

在这里插入图片描述

论文标题：Conformal Symplectic Optimization for Stable Reinforcement Learning
论文地址：https://ieeexplore.ieee.org/document/10792938
代码仓库：https://github.com/TobiasLv/RAD

神经网络作为RL价值函数和策略函数的核心载体，其训练过程通常依赖于神经网络优化器以实现参数更新。然而，目前主流的神经网络优化器（如SGD-M、Adam和AdamW等）虽然在缓解局部最优和加速收敛方面有所帮助，但其算法设计和参数选择均依赖于人工经验和实用技巧，缺乏对优化动态特性的解释与分析，难以从理论上保障RL训练的稳定性。

研究者从动力学视角出发，将神经网络参数的优化过程建模为多粒子相对论系统状态的演化过程，通过引入狭义相对论的光速最大原理，从理论上抑制了网络参数的异常更新速率，同时提供了各网络参数的独立自适应能力，有力保障了RL训练的长时域稳定性和收敛速率。

神经网络参数优化=动力学系统状态演变

本质上，RL的神经网络优化要求迭代求解以下标准形式的非凸随机优化问题：

在这里插入图片描述

其中， $J$ 为目标函数， $\mathcal{L}$ 为损失函数， $\mathcal{P}$ 为样本 $x$ 的分布， $\theta$ 为可训练网络参数。然而，该优化过程的稳定性通常难以直接分析，那么是否存在可行的替代方案呢？

研究者提出了一种全新视角，将神经网络参数的离散优化过程对偶为共形哈密顿系统状态的离散演变过程，通过分析和保留动力学系统的关键特征，向神经网络优化过程引入稳定性和收敛性等重要性能。

在这里插入图片描述

共形哈密顿系统（Conformal Hamiltonian system）是一类特殊的能量耗散系统，其哈密顿量（即总能量）和辛形式（即相空间面积）随时间单调收敛，从而使得系统趋向于稳定点。这类系统的正则方程具有特殊的反对称结构，保证了系统的长期稳定性：

在这里插入图片描述

其中 $q$ 为系统状态， $p$ 为系统动量， $r$ 为耗散因子，哈密顿量 $H$ 由系统动能 $T (p)$ 和系统势能 $U (q)$ 组成，即 $H (q, p) = T (p) + U (q)$ 。此外，大量研究表明，使用保辛离散方法得到的离散系统能够在长时域下以高精度近似原始连续系统，从而具备后者的稳定性和收敛速度。

这一工作的研究者发现，神经网络的收敛条件与共形哈密顿系统的收敛条件呈现出高度相似性：

1）神经网络收敛条件：目标函数极小且参数更新速度为0，即 $\nabla J(\theta) = 0$ 且 $\dot{\theta} = 0;$
2）共形哈密顿系统收敛条件：系统势能极小且系统动量为0，即 $\nabla U(q) = 0$ 且 $p = 0$ 。

因此，通过将网络参数 $\theta$ 类比为系统状态 $q$ ，目标函数 $J(\theta)$ 类比为系统势能 $U (q)$ ，可直接建立二者间的对偶关系：

在这里插入图片描述

研究者据此开发了一个用于设计神经网络离散优化方法的全新框架，包括两个核心步骤：

1）确定动能形式：选取合适的系统动能形式 $T (p)$ 构建动力学系统，向优化方法引入恰当的动态特性；
2）采用保辛离散化：选取合适的保辛离散方法实现特定精度的动力学近似，确保优化方法的稳定性和收敛性。

在这一框架下，神经网络离散优化方法将自然继承共形哈密顿系统的长时域迭代稳定性。然而，研究者特别强调，动能 $T (p)$ 对于优化动力学具有决定性作用，选择合适的动能形式对于实现高效优化至关重要。

具有相对论约束的自适应梯度下降机制

经典动能 $T(p) = \| p \|^2 / (2m)$ 是一个直观的选择，但其对应于牛顿力学系统，该类系统速度（即参数更新速度 $\dot{\theta}$ ）没有上限，从而可能在异常梯度影响下产生不受控制的训练行为。

为有效防止训练中的参数发散现象，受狭义相对论光速最大原理启发，研究者指出通过对相对论系统进行保辛离散化，可向神经网络优化过程中自然引入参数更新速度的限制。

研究者进一步指出，通过构建由大量独立1维粒子组成的相对论共形哈密顿系统模型，并将每个网络参数 $\theta_i$ 对应于一个1维粒子 $q_i$ ，可以在优化过程中实现参数的个体自适应性。该性质使得优化方法能够更灵活地应对参数空间的复杂性，进而加速收敛。

具体而言，由 $n$ 个具有相同质量 $m$ 的1维粒子组成的相对论系统具有以下哈密顿量：

在这里插入图片描述

对该系统进行保辛离散化，并引入变量替换以适应深度学习领域常用的符号体系，可以得到如下网络参数离散优化方法：

在这里插入图片描述

其中 $v$ 为一阶动量， $\beta_1$ 为一阶动量系数， $\alpha$ 为学习率， $\delta$ 为速度系数。

由于建模为多粒子相对论系统，每个网络参数的有效学习率 $\alpha_{k, i} = \alpha / (\delta^2 v_{k+1, i}^2 + 1)^{1/2}$ 可以根据其二阶动量 $v_{k+1, i}^2$ 自适应调节，这有助于在非凸随机优化设定下实现快速收敛。

同时， $|\theta_{k+1, i} - \theta_{k, i}| \leq \alpha / \delta$ 总成立，参数更新速度受到了预设限制。当 $\delta \to 0$ 时（即系统趋向于牛顿力学系统），参数更新将不再受到限制，优化过程容易受到异常梯度的影响。

研究者进一步引入了二阶动量 $y = v^2$ 的指数移动平均估计和动量估计偏差修正技术，从而得到了适用于深度神经网络非凸随机优化场景的离散迭代方法：

在这里插入图片描述

其中， $\beta_2$ 为二阶动量系数，保辛因子 $\zeta_k = 1 - \beta_2^{k+1}$ ，所有运算均逐元素进行。

研究者观察到，遵循原始退火机制的保辛因子 $\zeta_k$ 会导致有效学习率 $\alpha_k$ 始终低于初始学习率 $\alpha$ 。鉴于训练初期较高的有效学习率有利于网络参数迅速接近最优解空间，他们提出了一种改进方案：对 $\zeta_k$ 的退火过程进行人为调整，使其在训练初期尽量小，然后随训练进行逐渐恢复到 $\beta_2^{k+1}$ 。这种调整策略旨在平衡初期的收敛速度和后期的训练稳定性。

最终，研究者提出了既具备稳定动力学特性又适用于深度神经网络非凸随机优化的RAD优化器（伪代码见算法3）。

在这里插入图片描述

值得注意的是，当速度系数 $\delta$ 设为1且保辛因子 $\zeta_k$ 固定为小常数 $\epsilon$ 时，RAD优化器将退化为深度学习中广泛采用的Adam优化器。研究者据此阐明了RAD优化器能够无缝切换的两阶段优化过程：

1）优化初期：RAD类似于Adam，迅速引导网络参数接近最优解空间；
2）优化后期：随着保辛性质的恢复，RAD保持原系统动态特性，确保长时域稳定训练。

研究者进一步指出，在Adam优化器中， $\epsilon$ 是一个为避免数值运算中分母为零而人为加入的小常数。与之相比，RAD中的保辛因子 $\zeta$ 与系统的“质量 $\times$ 质能”（即 $m^2 c^2$ ）相关，具有明确的物理根源。该发现一定程度上阐释了为何在以往研究中，适度增加 $\epsilon$ （使得优化过程更贴近原始系统动力学）总能提升Adam优化器的性能。

这些发现不仅深化了对Adam优化器的理解，同时也为从动力学视角探究其他主流自适应梯度优化方法（如AdaGrad、NAdam、AdamW和Lion等）开辟了全新路径。

标准测试任务的实验分析

研究者在5种主流RL算法（包括DQN、DDPG、TD3、SAC和ADP）和12个测试环境（包括1个CartPole任务、6个MuJoCo任务、4个Atari任务和1个自动驾驶任务）中对RAD优化器进行了广泛测试，并与9种主流神经网络优化器（包括SGD、SGD-M、DLPF、RGD、NAG、Adam、NAdam、SWATS和AdamW）进行了比较。实验主要聚焦三方面验证：

1）长期训练稳定性：RAD具备长时域训练稳定性，有助于RL策略收敛后达到高性能；
2）速度系数 $\delta$ 的影响：通过控制网络参数更新速度，有助于减轻异常梯度不利影响；
3）保辛因子的作用：通过调节有效学习率自适应能力，有助于减弱训练振荡情况。

性能对比

研究者首先在经典控制任务CartPole上展开了实验。结果表明，RAD呈现出最平滑的哈密顿量下降趋势（见图1右侧），这一结果证实了RAD有效保持了原动力学系统的耗散特性。在策略性能（Total Average Return，TAR）方面，RAD曲线迅速平稳上升，而Adam曲线波动较大。

在这里插入图片描述

研究者接着在MuJoCo多关节机器人控制任务和Atari游戏环境中进行了广泛测试。结果表明，RAD在所有基准测试中均展现出卓越的收敛速度和性能。与RL领域广泛使用的Adam优化器相比，RAD在Walker2d任务中的性能提升了5.4％（见表I），而在Seaquest任务中的性能提升更是达到了显著的155.1%（见表II）。

在这里插入图片描述

与其他SOTA自适应梯度优化器相比，RAD在保持计算效率和内存占用相当的情况下（均维护一阶动量 $v$ 和二阶动量 $y$ ），同样表现出优越性（见图4）。研究者还指出，诸如NAdam和AdamW等Adam变体所采用技术与RAD互补，未来工作将探索将这些技术整合到RAD中，以期进一步提高性能。

在这里插入图片描述

为评估RAD在受扰动环境中的鲁棒性，研究者进一步在Walker2d任务中针对不同观测噪声水平进行了对比实验。结果表明（见图5），RAD因其出色的保辛性质（保留原动力学系统的稳定演化和收敛特性），展现出对噪声的显著抵抗力。即使在最高噪声水平下，RAD性能下降也相对较小，仅为3.4%；相比之下，NAdam的性能下降达到了10.9%。

在这里插入图片描述

最后，为了验证RAD在处理复杂现实世界任务中的有效性，研究者在模拟城市路口的自动驾驶任务中进行了测试。结果显示，使用RAD训练出的自动驾驶汽车通过率达93.2%，显著优于使用Adam优化器的89.5%。

以上结果证明了RAD优化器在RL基准测试中达到SOTA性能，且适用于复杂现实世界应用，RAD的这一特性使其成为解决各类RL训练任务（特别是那些对长期稳定性和收敛性能有高要求的复杂任务）的一个强有力的优化工具。

速度系数消融实验

为探究速度系数 $\delta$ 对RAD训练稳定性的影响，研究者选择在Humanoid任务中进行消融实验，该任务在训练过程中易出现异常大的梯度值。如图6所示，当 $\delta$ 设为1时，RAD相较于Adam展现出较小波动。当 $\delta$ 增加至5时，RAD在训练初期的振荡抑制作用更加显著。

研究者指出， $\delta$ 控制了参数更新的最大速度，即 $\| \theta_{k+1} - \theta_k \|_{\infty} \leq \alpha / \delta$ 。因此，选择较大的 $\delta$ 可有效防止参数发生剧烈变化，从而增强RAD在面对异常梯度时的鲁棒性和恢复能力。

在这里插入图片描述

保辛因子消融实验

为研究保辛因子 $\zeta$ 对RAD自适应能力的影响，研究者在Hopper任务上开展了消融实验，该任务在训练过程中易产生方差异常大的梯度序列。如图7所示，当RAD采取固定的保辛因子 $\zeta = 1 \times 10^{-16}$ （即退化为Adam优化器）时，训练曲线出现了严重振荡。然而，将 $\zeta$ 增加到 $\times 10^{-6}$ 后，这种不稳定现象得到显著改善。