基于JAX的看涨期权蒙特卡洛模拟对标的资产求一阶导数问题

问题根因 #

• 原代码的Simulation函数逻辑和grad、vmap的规则冲突：grad要求被求导函数输出必须是单个标量，而你原函数内部写了批量处理xi的逻辑，当vmap把单个xi值传入函数时，len(xi)会直接报错（标量没有长度属性），同时返回值也不是标量，不符合grad的输入要求。
• 原逻辑里随机数生成的位置没问题，能保证前向和反向传播用同一组布朗运动路径，这点不需要调整。

修复后可运行代码 #

from jax import random
from jax import jit, grad, vmap
import jax.numpy as jnp

# 固定全局参数，预先生成布朗运动，避免重复计算同时保证前后向路径一致
K,T,number_sim,sigma,r,q = 1.,1.,100,0.4,0,0
r_numb = random.PRNGKey(10)
mean = -.5 * sigma * sigma * T
volatility = sigma*jnp.sqrt(T)
BM = mean + volatility * random.normal(r_numb, shape=(number_sim,))
exp_BM = jnp.exp(BM)

# 函数接收单个标的价格S0，输出单个期权价格标量，符合grad的输入要求
def Simulation(S0):
    product = S0 * exp_BM
    payoff = jnp.maximum(product - K, 0)
    result = jnp.average(payoff) * jnp.exp(-q*T)
    return result

xi = jnp.linspace(1,1.2,5)
# 先对单个点求导，再批量映射到所有xi点
first_derivative = vmap(grad(Simulation))(xi)
print(first_derivative)

优化建议 #

1. 上述代码已经把布朗运动的生成移到了函数外部，避免每次求导重复生成随机数，大幅提升运行速度，导数计算的稳定性也更高。
2. 如果要提升蒙特卡洛估计的精度，可以适当提高number_sim的数值，AD计算的开销和前向模拟的开销基本是线性关系，不会额外增加太多成本。
3. 可以给外层加jit编译进一步提速：first_derivative = jit(vmap(grad(Simulation)))(xi)，运行效率会有数量级的提升。

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