torch_numopt.curvature package#
Submodules#
- torch_numopt.curvature.exact_block_hessian module
- torch_numopt.curvature.exact_hessian module
- torch_numopt.curvature.gauss_newton_approximation module
- torch_numopt.curvature.gauss_newton_block_approximation module
- torch_numopt.curvature.hutchinson_diagonal_approximation module
- torch_numopt.curvature.naive_identity module
Module contents#
Curvature estimators for second-order optimization.
This package provides classes that approximate or compute the Hessian matrix (and its products) in various ways:
Exact Hessian (full or block-diagonal)
Gauss-Newton approximation (full or block)
Hutchinson diagonal approximation (via random projections)
Identity (no curvature)
All estimators inherit from CurvatureEstimator and implement the
scaling_matrix, hvp, and quadratic_form methods.
- class NaiveIdentityCalculator[source]#
Bases:
CurvatureEstimatorCurvature estimator that always returns the identity matrix.
The scaling matrix is 1 (scalar), the Hessian-vector product is the vector itself, and the quadratic form is the squared norm.
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, step_dir)Compute the quadratic form vᵀ H v.
reset()Reset any internal state (intended to be used for quasi-Newton methods).
scaling_matrix(objective, params)Obtain the curvature matrix in its native representation.
update()Updates the parameters of the curvature estimator.
- scaling_matrix(objective, params)[source]#
Obtain the curvature matrix in its native representation.
- Return type:
float- Parameters:
- objectiveObjectiveFunction
Objective function.
- paramsParams
Parameter tensors.
- Returns:
- iterable or torch.Tensor
Representation of the matrix (scalar, vector, tuple of blocks, or full tensor).
- class ExactHessianCalculator(damping=None, mu=0.0001)[source]#
Bases:
CurvatureEstimatorCompute the exact Hessian matrix (full or block) of the objective.
The Hessian is obtained using
torch.func.hessian, which computes the full second-order derivatives. For large models, this can be memory- intensive; use the block version for parameter groups.- Parameters:
- dampingstr or None, default=None
Damping strategy:
"identity"addsmu * I,"fletcher"addsmu * diag(H). IfNone, no damping is applied.- 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 exact hessian of the Neural network given a dataset.
update()Updates the parameters of the curvature estimator.
- scaling_matrix(objective, params)[source]#
Calculation of the exact hessian of the Neural network given a dataset.
- 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.
- class ExactBlockHessianCalculator(damping=None, mu=0.0001)[source]#
Bases:
CurvatureEstimatorExact Hessian approximated as a block-diagonal matrix.
Each block corresponds to a single parameter tensor (e.g., a weight matrix). This is often sufficient for many optimization problems and is cheaper than the full Hessian.
- Parameters:
- dampingstr or None, default=None
Damping strategy (see
ExactHessianCalculator).- 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 exact hessian of the Neural network given a dataset.
update()Updates the parameters of the curvature estimator.
- scaling_matrix(objective, params)[source]#
Calculation of the exact hessian of the Neural network given a dataset.
- 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.
- class GaussNewtonApproximation(vectorize=True, damping=None, mu=0.0001)[source]#
Bases:
CurvatureEstimatorFull Gauss-Newton Hessian approximation.
The matrix is computed as Jᵀ J, where J is the Jacobian of the residual vector with respect to the parameters. This estimator forms a single dense matrix.
- Parameters:
- vectorizebool, default=True
If
True, use vectorized Jacobian computation (may be faster).- dampingstr or None, default=None
Damping strategy (identity or Fletcher).
- 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.
- 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.
- class HutchinsonDiagonalApproximation(n_samples=1, skip_iters=0)[source]#
Bases:
CurvatureEstimatorDiagonal Hessian estimator via Hutchinson’s method.
The diagonal is estimated as the average of z ⊙ (H z) over random Rademacher vectors z.
- Parameters:
- n_samplesint, default=1
Number of random samples to average.
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, step_dir)Compute the quadratic form vᵀ H v.
reset()Reset any internal state (intended to be used for quasi-Newton methods).
scaling_matrix(objective, params)Obtain the curvature matrix in its native representation.
update()Updates the parameters of the curvature estimator.
- scaling_matrix(objective, params)[source]#
Obtain the curvature matrix in its native representation.
- Return type:
Tuple[Tensor]- Parameters:
- objectiveObjectiveFunction
Objective function.
- paramsParams
Parameter tensors.
- Returns:
- iterable or torch.Tensor
Representation of the matrix (scalar, vector, tuple of blocks, or full tensor).
- hvp(objective, params, step_dir)[source]#
Compute the Hessian-vector product H * v.
- Return type:
Tuple[Tensor]- Parameters:
- objectiveObjectiveFunction
Objective function.
- paramsParams
Parameter tensors.
- step_dirParams
Vector v (same structure as params).
- Returns:
- Params
Result of H * v.