Compute the product of the Hessian of a loss with a vector v by applying forward-mode differentiation to the gradient, without ever forming the Hessian matrix.
The loss depends on a weight vector w of length D:
f(w)=N1n=1∑N(log(1+exn⋅w)−ynxn⋅w)+4αi=1∑Dwi4Here x has shape (N, D) with rows xn, y has shape (N,) with labels in {0, 1}, and α is the argument alpha. The Hessian H(w) is the D by D matrix of second derivatives ∂2f/∂wi∂wj, and the Hessian-vector product is H(w)v.
Return H(w)v as a float32 JAX array of shape (D,).
Input: w = [0.5,-0.5], v = [1.0,0.0], x = [[1.0,2.0],[0.0,1.0],[-1.0,1.0]], y = [1.0,0.0,1.0], alpha = 0.0
Output: [0.143872,0.091132]
Explanation: With alpha = 0 only the logistic term has curvature. The returned vector equals the first column of the 2 by 2 Hessian, because v = [1, 0] selects that column.
Input: w = [1.0,-2.0,0.5], v = [0.5,1.0,-1.0], x = [[0.5,-1.0,2.0],[1.0,0.0,-0.5]], y = [0.0,1.0], alpha = 0.5
Output: [0.839386,6.039123,-0.50772]
Input: w = [0.0,1.0,0.0,-1.0], v = [0.0,0.0,1.0,0.0], x = [[1.0,-1.0,2.0,0.5]], y = [1.0], alpha = 0.2
Output: [0.298293,-0.298293,0.596586,0.149146]
jax.jvp(f, (x,), (v,)) returns (f(x), directional derivative of f along v).
Pass jax.grad(f) as the f of jax.jvp to differentiate the gradient.
jax.nn.softplus keeps log(1 + exp(z)) finite for large z.
Accepts: array
Accepts: array
Accepts: array
Accepts: array
Accepts: number
Compute the product of the Hessian of a loss with a vector v by applying forward-mode differentiation to the gradient, without ever forming the Hessian matrix.
The loss depends on a weight vector w of length D:
f(w)=N1n=1∑N(log(1+exn⋅w)−ynxn⋅w)+4αi=1∑Dwi4Here x has shape (N, D) with rows xn, y has shape (N,) with labels in {0, 1}, and α is the argument alpha. The Hessian H(w) is the D by D matrix of second derivatives ∂2f/∂wi∂wj, and the Hessian-vector product is H(w)v.
Return H(w)v as a float32 JAX array of shape (D,).
Input: w = [0.5,-0.5], v = [1.0,0.0], x = [[1.0,2.0],[0.0,1.0],[-1.0,1.0]], y = [1.0,0.0,1.0], alpha = 0.0
Output: [0.143872,0.091132]
Explanation: With alpha = 0 only the logistic term has curvature. The returned vector equals the first column of the 2 by 2 Hessian, because v = [1, 0] selects that column.
Input: w = [1.0,-2.0,0.5], v = [0.5,1.0,-1.0], x = [[0.5,-1.0,2.0],[1.0,0.0,-0.5]], y = [0.0,1.0], alpha = 0.5
Output: [0.839386,6.039123,-0.50772]
Input: w = [0.0,1.0,0.0,-1.0], v = [0.0,0.0,1.0,0.0], x = [[1.0,-1.0,2.0,0.5]], y = [1.0], alpha = 0.2
Output: [0.298293,-0.298293,0.596586,0.149146]
jax.jvp(f, (x,), (v,)) returns (f(x), directional derivative of f along v).
Pass jax.grad(f) as the f of jax.jvp to differentiate the gradient.
jax.nn.softplus keeps log(1 + exp(z)) finite for large z.
Accepts: array
Accepts: array
Accepts: array
Accepts: array
Accepts: number
