torch_numopt.algorithms.adahessian module

AdaHessian optimizer (diagonal Hessian with momentum).

This module implements the AdaHessian algorithm, which combines the adaptive learning rate mechanism of Adam with a diagonal Hessian approximation computed via Hutchinson’s method. It maintains moving averages of both the gradient and the Hessian diagonal, and uses them to compute a preconditioned step direction.

class AdaHessianMixin(*args, beta1=0.9, beta2=0.999, k=1, eps=1e-08, **kwargs)

Bases: object

Mixin that implements the AdaHessian algorithm.

AdaHessian uses a diagonal Hessian approximation (via Hutchinson’s method) and maintains moving averages of the gradient and the squared diagonal Hessian.

Parameters:

beta1float, default=0.9

Exponential decay rate for the gradient moment.

beta2float, default=0.999

Exponential decay rate for the Hessian diagonal moment.

kfloat, default=1

Exponent used in the denominator; typically 0.5 for AdaHessian (root), but here set to 1 to allow flexibility.

epsfloat, default=1e-4

Small constant for numerical stability in the division.

Methods

get_step_direction

get_step_direction(objective, grad_params)

class AdaHessian(params, lr_init=0.01, lr_method=None, beta1=0.9, beta2=0.999, k=1, eps=1e-08, n_samples=10, skip_iters=0)

Bases: AdaHessianMixin, NumericalOptimizer

AdaHessian optimizer (diagonal Hessian with momentum).

Uses Hutchinson diagonal Hessian approximation and momentum for both gradient and Hessian diagonal, similar to Adam but using second-order information.

Parameters:

paramsParams

Parameter tensors.

lr_initfloat, default=1

Initial learning rate.

lr_methodstr or None, default=None

Learning rate initialization method.

beta1float, default=0.9

Exponential decay rate for the first moment estimate (gradient).

beta2float, default=0.999

Exponential decay rate for the second moment estimate (Hessian diagonal).

kfloat, default=1

Exponent for the Hessian diagonal in the step calculation (0.5 for AdaHessian).

epsfloat, default=1e-4

Small constant for numerical stability.

n_samplesint, default=5

Number of Hutchinson samples for diagonal estimation.

skip_itersint, default=0

Number of times in which we reuse the previously computed diagonal approximation (No skipping by default).

Methods

add_param_group(param_group)	Add a param group to the Optimizer s param_groups.
apply_gradients(objective, params, grad_params)	Update parameters using the current gradient and curvature.
load_state_dict(state_dict)	Load the optimizer state.
register_load_state_dict_post_hook(hook[, ...])	Register a load_state_dict post-hook which will be called after load_state_dict() is called. It should have the following signature::.
register_load_state_dict_pre_hook(hook[, ...])	Register a load_state_dict pre-hook which will be called before load_state_dict() is called. It should have the following signature::.
register_state_dict_post_hook(hook[, prepend])	Register a state dict post-hook which will be called after state_dict() is called.
register_state_dict_pre_hook(hook[, prepend])	Register a state dict pre-hook which will be called before state_dict() is called.
register_step_post_hook(hook)	Register an optimizer step post hook which will be called after optimizer step.
register_step_pre_hook(hook)	Register an optimizer step pre hook which will be called before optimizer step.
state_dict()	Return the state of the optimizer as a dict.
step(objective)	Perform one optimization step.
zero_grad([set_to_none])	Reset the gradients of all optimized torch.Tensor s.

OptimizerPostHook
OptimizerPreHook
get_step_direction
profile_hook_step

class AdaHessianLS(params, lr_init=1, lr_method=None, beta1=0.9, beta2=0.999, k=1, eps=1e-08, n_samples=10, skip_iters=0, c1=0.0001, c2=0.9, tau=0.1, max_iter=20, tol=1e-08, line_search_method='backtrack', line_search_cond='armijo')

Bases: AdaHessianMixin, LineSearchOptimizer

AdaHessian with line search.

Same as AdaHessian, but instead of a fixed learning rate it performs a line search to determine the step length.

Works well in practice, but theoretically it’s not well supported.

Parameters:

paramsParams

Parameter tensors.

lr_initfloat, default=1

Initial learning rate.

lr_methodstr or None, default=None

Learning-rate initialization method.

beta1, beta2, k, eps, n_samples, skip_iterssame as in AdaHessian.

c1, c2, tau, max_iter, tolline-search parameters.

line_search_methodstr, default=”backtrack”

Line-search method.

line_search_condstr, default=”armijo”

Line-search condition.

Methods

add_param_group(param_group)	Add a param group to the Optimizer s param_groups.
apply_gradients(objective, params, grad_params)	Update parameters using the current gradient and curvature.
load_state_dict(state_dict)	Load the optimizer state.
register_load_state_dict_post_hook(hook[, ...])	Register a load_state_dict post-hook which will be called after load_state_dict() is called. It should have the following signature::.
register_load_state_dict_pre_hook(hook[, ...])	Register a load_state_dict pre-hook which will be called before load_state_dict() is called. It should have the following signature::.
register_state_dict_post_hook(hook[, prepend])	Register a state dict post-hook which will be called after state_dict() is called.
register_state_dict_pre_hook(hook[, prepend])	Register a state dict pre-hook which will be called before state_dict() is called.
register_step_post_hook(hook)	Register an optimizer step post hook which will be called after optimizer step.
register_step_pre_hook(hook)	Register an optimizer step pre hook which will be called before optimizer step.
state_dict()	Return the state of the optimizer as a dict.
step(objective)	Perform one optimization step.
zero_grad([set_to_none])	Reset the gradients of all optimized torch.Tensor s.

OptimizerPostHook
OptimizerPreHook
get_step_direction
profile_hook_step