如何对批量输入下神经网络的单个输出求导（JAX实现）

错误原因 #

你的代码中，p和q函数将网络输出reshape为(-1,1)，导致单个样本输入时输出形状为(1,1)。此时jacfwd计算的雅可比矩阵形状为(1,1,1)，经过vmap批量处理后，p_x(inputs)的形状变成(10,1,1)——后续和kappa(inputs).reshape(-1,1)（形状(10,1)）相乘时，维度不匹配触发索引错误。此外，如果输入inputs的形状是(10,)而非(10,1)，也会加剧维度计算混乱。

修复方案 #

方案一：简化输出维度（推荐） #

直接让p和q返回标量输出（批量时为一维数组），这样jacfwd计算的梯度维度更简洁：

# 去掉reshape，单个样本输出为标量，批量输出为(10,)
p = lambda inputs: f(params, inputs)[:, 0]
q = lambda inputs: f(params, inputs)[:, 1]

# vmap批量应用jacfwd，得到形状为(10,1)的梯度数组
p_x = vmap(jacfwd(p, argnums=0))
q_x = vmap(jacfwd(q, argnums=0))

# 后续计算无需调整维度
k_p_x = lambda inputs: kappa(inputs).reshape(-1,1) * p_x(inputs)

方案二：保留reshape并压缩冗余维度 #

如果需要维持p和q的输出为二维数组，可通过squeeze去掉梯度结果中的冗余维度：

p = lambda inputs: f(params, inputs)[:,0].reshape(-1,1)
q = lambda inputs: f(params, inputs)[:,1].reshape(-1,1)

# 用squeeze去掉vmap后梯度的中间维度，得到(10,1)的结果
p_x = lambda inputs: jnp.squeeze(vmap(jacfwd(p, argnums=0))(inputs), axis=1)
q_x = lambda inputs: jnp.squeeze(vmap(jacfwd(q, argnums=0))(inputs), axis=1)

k_p_x = lambda inputs: kappa(inputs).reshape(-1,1) * p_x(inputs)

注意事项 #

确保传入的inputs形状为(批量大小, 输入维度)，即示例中的(10,1)，而非(10,)——若输入是一维数组，可先通过inputs.reshape(-1,1)转换形状。

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