在 JAX 中高效计算 Hessian 矩阵答案

【问题标题】：Compute efficiently Hessian matrices in JAX在 JAX 中高效计算 Hessian 矩阵
【发布时间】：2022-01-03 23:01:55
【问题描述】：

在 JAX 的快速入门教程中，我发现可以使用以下代码行有效地计算可微函数 fun 的 Hessian 矩阵：

from jax import jacfwd, jacrev

def hessian(fun):
  return jit(jacfwd(jacrev(fun)))

但是，也可以通过以下计算来计算 Hessian：

def hessian(fun):
  return jit(jacrev(jacfwd(fun)))

def hessian(fun):
  return jit(jacfwd(jacfwd(fun)))

def hessian(fun):
  return jit(jacrev(jacrev(fun)))

这是一个最小的工作示例：

import jax.numpy as jnp
from jax import jit
from jax import jacfwd, jacrev

def comp_hessian():

    x = jnp.arange(1.0, 4.0)

    def sum_logistics(x):
        return jnp.sum(1.0 / (1.0 + jnp.exp(-x)))

    def hessian_1(fun):
        return jit(jacfwd(jacrev(fun)))

    def hessian_2(fun):
        return jit(jacrev(jacfwd(fun)))

    def hessian_3(fun):
        return jit(jacrev(jacrev(fun)))

    def hessian_4(fun):
        return jit(jacfwd(jacfwd(fun)))

    hessian_fn = hessian_1(sum_logistics)
    print(hessian_fn(x))

    hessian_fn = hessian_2(sum_logistics)
    print(hessian_fn(x))

    hessian_fn = hessian_3(sum_logistics)
    print(hessian_fn(x))

    hessian_fn = hessian_4(sum_logistics)
    print(hessian_fn(x))


def main():
    comp_hessian()


if __name__ == "__main__":
    main()

我想知道什么时候最好使用哪种方法？我也想知道是否可以使用grad() 来计算Hessian？ grad() 与 jacfwd 和 jacrev 有何不同？

【问题讨论】：

标签： jax

【解决方案1】：

您的问题的答案在 JAX 文档中；例如看这个部分：https://jax.readthedocs.io/en/latest/notebooks/autodiff_cookbook.html#jacobians-and-hessians-using-jacfwd-and-jacrev

引用其对jacrev 和jacfwd 的讨论：

这两个函数计算相同的值（直到机器数字），但它们的实现不同：jacfwd 使用正向模式自动微分，这对于“高”雅可比矩阵更有效，而 jacrev 使用反向-mode，对于“宽”雅可比矩阵更有效。对于近似正方形的矩阵，jacfwd 可能比jacrev 有优势。

再往下，

要实现 hessian，我们可以使用 jacfwd(jacrev(f)) 或 jacrev(jacfwd(f)) 或两者的任何其他组合。但是正向反向通常是最有效的。这是因为在内部雅可比计算中，我们经常对函数范围的雅可比进行微分（可能像损失函数 ?:ℝⁿ→ℝ），而在外部雅可比计算中，我们用平方雅可比行列式（因为 ∇?:ℝⁿ→ℝⁿ），这是正向模式胜出的地方。

由于您的函数看起来像 ?:ℝⁿ→ℝ，因此 jit(jacfwd(jacrev(fun))) 可能是最有效的方法。

至于为什么不能用grad 实现hessian，这是因为grad 仅设计用于具有标量输出的函数的导数。根据定义，hessian 是向量值雅可比的组合，而不是标量梯度的组合。

【讨论】：