如何加速Zygote.jl自动微分计算Black-Scholes期权希腊值的速度

你观测到的数十秒耗时绝大多数是Zygote第一次生成梯度代码的编译耗时，实际函数计算和梯度计算的运行时耗时只有微秒级。你可以尝试在main函数中重复调用一次bs_call'(theta)，第二次调用的耗时会降到几乎可以忽略的程度。

下面是可以进一步大幅降低耗时的优化方案：

• 替换为前向模式自动微分
  你当前的输入参数维度仅为5，属于典型的低维度求导场景，反向模式自动微分（Zygote）的优势在高输入维度、低输出维度的场景，低维度下ForwardDiff.jl的性能远优于Zygote，编译耗时也低很多。仅需要将梯度计算的代码替换为ForwardDiff.gradient(bs_call, theta)，其余代码无需修改，即可获得数倍到数十倍的性能提升。
• 为函数自定义反向微分规则（rrule）
  Black-Scholes模型的希腊值有成熟的解析表达式，不需要Zygote对整个函数做源到源的自动微分，你可以直接为bs_call函数定义自定义反向规则，完全消除自动微分的编译和计算开销，示例代码如下：

using Zygote, SpecialFunctions, ChainRulesCore

function cdf(x)
    0.5 * (1 + erf(x / sqrt(2)))
end
# 正态分布PDF，用于计算梯度
pdf(x) = exp(-x^2/2)/sqrt(2π)

function bs_call(theta)
    s, k, r, t, sigma = theta
    vol = sigma * sqrt(t)
    d1 = (log(s / k) + (r + 0.5 * sigma ^ 2) * t) / vol
    d2 = d1 - vol
    price = cdf(d1) * s - cdf(d2) * k * exp(-r * t)
    return price
end

# 自定义rrule，直接返回解析梯度
function ChainRulesCore.rrule(::typeof(bs_call), theta)
    s, k, r, t, sigma = theta
    vol = sigma * sqrt(t)
    d1 = (log(s / k) + (r + 0.5 * sigma ^ 2) * t) / vol
    d2 = d1 - vol
    price = cdf(d1) * s - cdf(d2) * k * exp(-r * t)
    N_d1 = cdf(d1)
    n_d1 = pdf(d1)
    discount = exp(-r*t)
    # 解析希腊值
    delta = N_d1
    theta_greek = - (s * n_d1 * sigma)/(2*sqrt(t)) - r*k*discount*cdf(d2)
    rho = k * t * discount * cdf(d2)
    vega = s * n_d1 * sqrt(t)
    kappa = - discount * cdf(d2)
    function bs_call_pullback(Δ)
        return (NoTangent(), Δ * [delta, kappa, rho, theta_greek, vega])
    end
    return price, bs_call_pullback
end

定义完该规则后，Zygote求导时会直接调用你写的解析梯度，完全不需要自动遍历函数代码求导，速度可以达到微秒级。

• 使用静态数组作为输入
  你的输入参数theta是固定长度为5的小数组，使用StaticArrays.jl的SVector替代普通Vector作为输入，编译器可以完全推断类型、消除堆分配，同时大幅降低Zygote的微分编译开销，示例如下：

using StaticArrays
theta = @SVector [100.0, 110.0, 0.2, 1.0, 0.5]

• 批量计算时向量化代码
  如果你需要批量计算大量期权的希腊值，不要循环逐个求导，将参数向量化后一次性计算，或者使用广播机制，能进一步降低平均耗时。

内容的提问来源于stack exchange，提问作者MikeRand