torch_numopt.numerical_optimizer module#

Base classes for numerical optimizers that use curvature information.

This module provides the core NumericalOptimizer abstract class and its subclasses for line-search and trust-region strategies. It handles parameter updates, learning-rate initialization, and direction computation using a curvature estimator.

class NumericalOptimizer(params, curvature_estimator, lr_method=None, lr_init=1, solver='solve', fix_ascent=True, min_lr=0, max_lr=100)[source]#

Bases: Optimizer, ABC

Base optimizer that uses a curvature estimator to compute a step direction.

Subclasses must implement the strategy for determining the step (line-search or trust-region). This class handles storage of previous iterates and learning-rate initialization methods.

Parameters:
paramsParams

Iterable of parameter tensors to optimize.

curvature_estimatorCurvatureEstimator

Object that provides the second-order (or approximate) curvature.

lr_initfloat, default=1

Initial guess for the learning rate (or radius).

lr_methodstr or None, default=None

Method for initializing the learning rate. Options: - None: use the supplied lr_init directly. - "keep": reuse the previous learning rate. - "scaled": scale based on gradient and step direction. - "quadratic": use the curvature quadratic form. - "interpolate": interpolate from previous loss change. - "lipschitz": estimate Lipschitz constant. - "BB1", "BB2": Barzilai-Borwein formulas.

lr_tolfloat, default=1e-18

Tolerance for the learning rate; if the estimated rate falls below this, it is replaced by the initial guess.

solverstr, default=”solve”

Solver used to invert the curvature system (see solve_system).

fix_ascentbool, default=True

If True, detect and correct ascent directions (fall back to steepest descent).

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.

get_step_direction(objective, grad_params)

Compute the un-scaled step direction by solving the system H * p = -grad (or an approximation).

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

profile_hook_step

get_step_direction(objective, grad_params)[source]#

Compute the un-scaled step direction by solving the system H * p = -grad (or an approximation).

Parameters:
objectiveObjectiveFunction

Objective function.

grad_paramsParams

Current gradient.

Returns:
Params

Step direction (not yet multiplied by learning rate).

apply_gradients(objective, params, grad_params)[source]#

Update parameters using the current gradient and curvature.

This method computes a step direction, determines a step length (via learning-rate initialization or line-search/trust-region), and applies the update. It also updates stored previous iterates.

Parameters:
objectiveObjectiveFunction

Objective function.

paramsParams

Current parameters (in-place updated).

grad_paramsParams

Current gradient.

step(objective)[source]#

Perform one optimization step.

This is the main entry point called by the training loop. It calls objective.closure() to compute the loss and gradients, then calls apply_gradients for each parameter group.

Parameters:
objectiveObjectiveFunction

Objective function.

class LineSearchOptimizer(params, curvature_estimator, line_search, lr_method=None, lr_init=1, solver='solve', min_lr=0, max_lr=100)[source]#

Bases: NumericalOptimizer, ABC

Numerical optimizer that uses a line-search algorithm to determine the step length.

Parameters:
paramsParams

Parameter tensors.

curvature_estimatorCurvatureEstimator

Curvature estimator.

line_searchLineSearchSolver

Line-search solver (backtracking, interpolation, etc.).

lr_initfloat, default=1

Initial learning-rate guess.

lr_methodstr or None, default=None

Learning-rate initialization method.

solverstr, default=”solve”

Linear solver for the step direction.

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.

get_step_direction(objective, grad_params)

Compute the un-scaled step direction by solving the system H * p = -grad (or an approximation).

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

profile_hook_step

apply_gradients(objective, params, grad_params)[source]#

Update parameters using the current gradient and curvature.

This method computes a step direction, determines a step length (via learning-rate initialization or line-search/trust-region), and applies the update. It also updates stored previous iterates.

Parameters:
objectiveObjectiveFunction

Objective function.

paramsParams

Current parameters (in-place updated).

grad_paramsParams

Current gradient.

class TrustRegionOptimizer(params, trust_region, lr_init=1.0, curvature_estimator=None, *, accept_tol=0.1, contract_tol=0.25, expand_tol=0.75, growth_factor=2, shrink_factor=0.25, radius_max=1000.0)[source]#

Bases: NumericalOptimizer, ABC

Numerical optimizer that uses a trust-region strategy.

Parameters:
paramsParams

Parameter tensors.

trust_regionTrustRegionSolver

Trust-region solver that computes the step within a region.

lr_initfloat, default=1.0

Initial trust-region radius.

curvature_estimatorCurvatureEstimator, optional

Curvature estimator; not directly used by the method, kept for compatibility. The system will be solved internally by the trust region solver.

accept_tolfloat, default=0.1

Threshold for the ratio rho; if rho > accept_tol the step is accepted.

contract_tolfloat, default=0.25

Threshold for the ratio rho; if rho < contract_tol the model is considered poor and the trust-region radius is shrunk by shrink_factor.

expand_tolfloat, default=0.75

Threshold for the ratio rho; if rho > expand_tol and the step is at the trust-region boundary, the radius is expanded by growth_factor.

growth_factorfloat, default=2

Factor by which the trust-region radius is multiplied when expanded.

shrink_factorfloat, default=0.25

Factor by which the trust-region radius is multiplied when contracted.

radius_maxfloat, default=1e3

Maximum allowed trust-region radius.

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.

get_step_direction(objective, grad_params)

Compute the un-scaled step direction by solving the system H * p = -grad (or an approximation).

load_state_dict(state_dict)

Load the optimizer state.

new_model_radius(objective, radius, loss, ...)

Update the trust-region radius based on the ratio rho.

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

profile_hook_step

new_model_radius(objective, radius, loss, params, grad_params, new_loss, step_dir)[source]#

Update the trust-region radius based on the ratio rho.

The ratio rho measures the agreement between the actual reduction and the predicted reduction. If rho is small, the model is poor, so shrink the radius; if rho is large and the step is at the boundary, expand it.

Parameters:
objectiveObjectiveFunction

Objective function.

radiusfloat

Current trust-region radius.

radius_initfloat

Initial radius (used as upper bound for expansion).

losstorch.Tensor

Loss value at current parameters.

paramsParams

Current parameters.

grad_paramsParams

Gradient at current parameters.

new_losstorch.Tensor

Loss value at proposed new parameters.

step_dirParams

Proposed step direction.

Returns:
tuple (rho, new_radius)
rhofloat

Ratio of actual to predicted reduction.

new_radiusfloat

Updated trust-region radius.

apply_gradients(objective, params, grad_params)[source]#

Update parameters using the current gradient and curvature.

This method computes a step direction, determines a step length (via learning-rate initialization or line-search/trust-region), and applies the update. It also updates stored previous iterates.

Parameters:
objectiveObjectiveFunction

Objective function.

paramsParams

Current parameters (in-place updated).

grad_paramsParams

Current gradient.