[论文解读] On Training Implicit Models

Zhengyang Geng, Xinyu Zhang+3 more|arXiv (Cornell University)|Nov 9, 2021

Domain Adaptation and Few-Shot Learning参考文献 59被引用 13

一句话总结

本文提出了幽灵梯度（phantom gradient），这是一种隐式模型中精确梯度的轻量化近似方法，在避免昂贵的反向求解器的同时，保持或提升了性能。通过利用阻尼展开（damped unrolling）和诺伊曼级数（Neumann series），该方法在 ImageNet 上将训练速度提升 1.7 倍，并超越了精确梯度方法的性能，为训练深层隐式网络提供了一种稳定且高效的常数内存替代方案。

ABSTRACT

This paper focuses on training implicit models of infinite layers. Specifically, previous works employ implicit differentiation and solve the exact gradient for the backward propagation. However, is it necessary to compute such an exact but expensive gradient for training? In this work, we propose a novel gradient estimate for implicit models, named phantom gradient, that 1) forgoes the costly computation of the exact gradient; and 2) provides an update direction empirically preferable to the implicit model training. We theoretically analyze the condition under which an ascent direction of the loss landscape could be found, and provide two specific instantiations of the phantom gradient based on the damped unrolling and Neumann series. Experiments on large-scale tasks demonstrate that these lightweight phantom gradients significantly accelerate the backward passes in training implicit models by roughly 1.7 times, and even boost the performance over approaches based on the exact gradient on ImageNet.

研究动机与目标

• 通过隐式微分解决隐式模型中精确梯度计算的高计算成本问题。
• 开发一种计算高效且实证有效的梯度近似方法，用于训练隐式模型。
• 提供理论条件，确保近似梯度在损失曲面中保持上升方向。
• 在不改变隐式模型前向传播或推理行为的前提下实现更快的训练速度。
• 在视觉、语言和图学习等多种任务上验证该方法的有效性。

提出的方法

• 提出幽灵梯度作为一阶预言器（first-order oracle），在无需求解昂贵的反向线性系统的情况下近似精确梯度。
• 推导出确保幽灵梯度在损失曲面中保持上升方向的理论条件。
• 提出两种实现方式：阻尼固定点展开（damped fixed-point unrolling）和基于诺伊曼级数的 DEQ 模型近似。
• 保持相同的前向传播和推理流程，确保与现有求解器和架构的兼容性。
• 采用带幽灵梯度的随机梯度下降（SGD），在适当超参数选择下实现收敛。
• 引入显式雅可比正则化（Jacobian regularization, JR），以在不规则损失曲面（尤其是基于 Transformer 的模型）上提升训练稳定性。

实验结果

研究问题

• RQ1 一阶梯度近似能否在训练隐式模型时超越精确梯度计算？
• RQ2 在何种条件下，近似梯度仍能提供损失曲面中的有效上升方向？
• RQ3 与精确梯度方法相比，幽灵梯度在训练速度和模型精度方面表现如何？
• RQ4 幽灵梯度能否在如基于 Transformer 的模型等不规则损失曲面上稳定训练？
• RQ5 该方法在视觉、语言和图学习等多种任务上是否依然有效？

主要发现

• 在 ImageNet 上，幽灵梯度将反向传播速度提升至最多 12 倍，总训练时间减少 1.4–1.7 倍。
• 在 ImageNet 上，幽灵梯度实现的 top-1 准确率优于精确梯度方法，展现出更优的泛化能力。
• 即使在精确梯度求解器发散的病态条件下，该方法仍能保持稳定训练。
• 幽灵梯度的诺伊曼级数与阻尼展开实现方式满足理论上的上升方向条件，并在适当阻尼下实现收敛。
• 显式雅可比正则化显著提升了稳定性和性能，尤其在基于 Transformer 的模型中，可将显式正则化权重降低 20 倍。
• 在 Wikitext-103 上，带 JR 的 UPG 变体实现测试困惑度 24.4（训练速度提升 2.2 倍），以及困惑度 24.0（训练速度提升 1.7 倍）。