利用深度强化学习优化自动微分

自动微分用于计算雅可比矩阵在机器学习、计算流体动力学、机器人和金融等众多科学领域广泛应用。即使在雅可比计算中节省少量计算次数或内存使用，即可在能源消耗和运行时间方面产生巨大节省。虽然已有许多方法可以实现此类节省，但它们通常以计算效率为代价来近似精确雅可比。本文提出一种新方法，通过利用深度强化学习 (RL) 和称为“跨国消除”的概念，仍计算精确雅可比，同时优化必要的乘法数量。跨国消除是一种自动微分框架，将雅可比累积表述为对计算图上所有顶点的有序消除，而每一次消除都产生一定的计算成本。我们将寻找最小化必要乘法数量的最优消除顺序搜索表述为一个由 RL 智能体玩的单人游戏。我们展示了该方法在来自多个相关领域的任务上相对于最先进方法实现了最高可达33%的改进。此外，我们显示这些理论收益转化为实际的运行时改进，通过提供一个在 JAX 中高效执行所获得消除顺序的跨国消除解释器。