torch_numopt.curvature.gauss_newton_block_approximation module#
Block-diagonal Gauss-Newton approximation.
Similar to the full Gauss-Newton, but only the diagonal blocks of Jᵀ J are computed, one per parameter group. This saves memory and is often sufficient.
- class GaussNewtonBlockApproximation(vectorize=True, damping=None, mu=0.0001)[source]#
Bases:
CurvatureEstimatorBlock-diagonal Gauss-Newton Hessian approximation.
Each block is formed as J_iᵀ J_i, where J_i is the Jacobian of the residual with respect to the i-th parameter group. Cross-group derivatives are ignored.
- Parameters:
- vectorizebool, default=True
Use vectorized Jacobian computation.
- dampingstr or None, default=None
Damping strategy.
- mufloat, default=1e-4
Damping coefficient.
Methods
full_scaling_matrix(objective, params)Return the curvature matrix as a single dense tensor.
hvp(objective, params, step_dir)Compute the Hessian-vector product H * v.
quadratic_form(objective, params, grad_params)Compute the quadratic form vᵀ H v.
reset()Reset any internal state (intended to be used for quasi-Newton methods).
scaling_matrix(objective, params)Calculation of the an approximate hessian of the Neural network given a dataset as in the Gauss-Newton algorithm.
update()Updates the parameters of the curvature estimator.
jvp
- scaling_matrix(objective, params)[source]#
Calculation of the an approximate hessian of the Neural network given a dataset as in the Gauss-Newton algorithm. The approximate Hessian is calculated as the square of the Jacobian of the residual of every data point with respect to the parameters.
Let the loss function be, for example the MSE:
\(\mathcal{L}(x,y;\theta) = \sum^{N}_{i=1} (f(x_i; \theta) - y_i)^2 = \sum^{N}_{i=1} r_i\)
Then the Jacobian of the residuals will be the matrix:
\((J_{\theta}[\mathcal{L}])_{i,j} = \dfrac{\partial r_i}{\partial \theta_j}\)
Then, we will approximate the hessian as the product of the Jacobian with it’s transpose, noting that the result will be a square matrix with size \(p\\times p\) with \(p\) being the number of parameters of the model:
\(H_{\theta}[\mathcal{L}] \approx J_{\theta}[\mathcal{L}]^{\intercal} \cdot J_{\theta}[\mathcal{L}]\)
- Return type:
Tuple[Tensor]- Parameters:
- x: torch.Tensor
Input dataset for calculating the loss.
- y: torch.Tensor
Target dataset for calculating the loss.
- loss_fn: torch.Module
Loss function for which to calculate the hessian.
- vectorize: boolean
Use vectorization in pytorch’s implementation of the hessian calculation.