romtools.workflows.inverse.vi_optimization_methods#

Optimization methods and configs used by inverse VI workflows.

Classes

AdamSolver([learning_rate, ...])

Stateful Adam ascent direction for stochastic VI gradients.

NewtonSolver(regularization[, hessian_type, ...])

Generic Newton solver supporting diagonal or dense Hessians.

SteepestDescentSolver()

Generic steepest-descent solver.

VIAdamOptimizerConfig([gradient_method, ...])

Configuration for Adam-based stochastic ELBO maximization.

VIGradientOptimizerConfig([gradient_method, ...])

VILegacyLineSearchConfig([...])

VINewtonOptimizerConfig([...])

VIStochasticNonmonotoneLineSearchConfig([...])

class romtools.workflows.inverse.vi_optimization_methods.AdamSolver(learning_rate=None, learning_rate_scale=0.1, beta1=0.9, beta2=0.999, epsilon=1e-08, fisher_damping_initial=0.01, fisher_damping_decay_start=50, fisher_damping_min=1e-06, gradient_clip_norm=1000000.0, parameter_dimension=None)[source]#

Bases: object

Stateful Adam ascent direction for stochastic VI gradients.

Parameters:

  • learning_rate (float)

  • learning_rate_scale (float)

  • beta1 (float)

  • beta2 (float)

  • epsilon (float)

  • fisher_damping_initial (float)

  • fisher_damping_decay_start (int)

  • fisher_damping_min (float)

  • gradient_clip_norm (float)

  • parameter_dimension (int)

load_restart_state_dict(restart_data)[source]#

Restore Adam state and reject incompatible optimizer configuration.

Return type:

None

prepare_gradient(gradient, fisher_diagonal=None)[source]#

Apply Fisher damping and clipping without updating Adam moments.

Parameters:

gradient (ndarray)

Return type:

ndarray

resolved_learning_rate(gradient)[source]#

Return the learning rate used for the supplied VI gradient.

Parameters:

gradient (ndarray)

Return type:

float

restart_state_dict()[source]#

Return restart-safe Adam state and configuration metadata.

Return type:

dict

class romtools.workflows.inverse.vi_optimization_methods.NewtonSolver(regularization, hessian_type='diagonal', regularization_strategy='absolute', additive_regularization=0.0, regularization_epsilon=1e-12, fallback_learning_rate=0.01)[source]#

Bases: object

Generic Newton solver supporting diagonal or dense Hessians.

Parameters:

  • regularization (float)

  • hessian_type (str)

  • regularization_strategy (str)

  • additive_regularization (float)

  • regularization_epsilon (float)

  • fallback_learning_rate (float)

class romtools.workflows.inverse.vi_optimization_methods.SteepestDescentSolver[source]#

Bases: object

Generic steepest-descent solver.

class romtools.workflows.inverse.vi_optimization_methods.VIAdamOptimizerConfig(gradient_method='natural', learning_rate=None, learning_rate_scale=0.1, beta1=0.9, beta2=0.999, epsilon=1e-08, fisher_damping_initial=0.01, fisher_damping_decay_start=50, fisher_damping_min=1e-06, gradient_clip_norm=1000000.0, gradient_norm_tolerance=1e-05, max_iterations=1000, max_log_std_update=0.5, min_variational_std=1e-06, max_variational_std=1000000.0)[source]#

Bases: object

Configuration for Adam-based stochastic ELBO maximization.

By default Adam is applied to the mean-field Gaussian natural gradient. Following the ABRIS setup, the Fisher matrix is damped early in the optimization and the preconditioned gradient is clipped by its L2 norm. When learning_rate is None, the ABRIS convention 0.1 / parameter_dimension is used.

Parameters:

  • gradient_method (str)

  • learning_rate (float)

  • learning_rate_scale (float)

  • beta1 (float)

  • beta2 (float)

  • epsilon (float)

  • fisher_damping_initial (float)

  • fisher_damping_decay_start (int)

  • fisher_damping_min (float)

  • gradient_clip_norm (float)

  • gradient_norm_tolerance (float)

  • max_iterations (int)

  • max_log_std_update (float)

  • min_variational_std (float)

  • max_variational_std (float)

class romtools.workflows.inverse.vi_optimization_methods.VIGradientOptimizerConfig(gradient_method: str = 'standard', gradient_norm_tolerance: float = 1e-05, max_iterations: int = 1000, max_log_std_update: float = 0.5, min_variational_std: float = 1e-06, max_variational_std: float = 1000000.0)[source]#

Bases: object

Parameters:

  • gradient_method (str)

  • gradient_norm_tolerance (float)

  • max_iterations (int)

  • max_log_std_update (float)

  • min_variational_std (float)

  • max_variational_std (float)

class romtools.workflows.inverse.vi_optimization_methods.VILegacyLineSearchConfig(initial_step_size: float = 0.01, max_step_size: float = inf, step_size_growth_factor: float = 1.01, step_size_decay_factor: float = 3.0, max_step_size_decrease_trys: int = 5, relaxation_parameter: float = 2.05, line_search_objective: str = 'elbo', line_search_sample_growth_factor: float = 1.0, log_std_learning_rate_factor: float = 0.1)[source]#

Bases: object

Parameters:

  • initial_step_size (float)

  • max_step_size (float)

  • step_size_growth_factor (float)

  • step_size_decay_factor (float)

  • max_step_size_decrease_trys (int)

  • relaxation_parameter (float)

  • line_search_objective (str)

  • line_search_sample_growth_factor (float)

  • log_std_learning_rate_factor (float)

class romtools.workflows.inverse.vi_optimization_methods.VINewtonOptimizerConfig(gradient_norm_tolerance: float = 5e-05, max_iterations: int = 1000, max_log_std_update: float = 0.5, max_mean_update_std: float = None, min_variational_std: float = 1e-08, max_variational_std: float = 1000000.0, newton_metric: str = 'standard', newton_fallback_learning_rate: float = 0.02, newton_regularization: float = 0.01, newton_hessian_type: str = 'diagonal', newton_curvature_strategy: str = 'same_sample', newton_hessian_num_samples: int = None, newton_hessian_averaging_factor: float = 0.9, newton_regularization_strategy: str = 'absolute', newton_additive_regularization: float = 0.0, newton_regularization_epsilon: float = 1e-12, newton_adaptive_regularization: bool = False, newton_regularization_increase_factor: float = 5.0, newton_regularization_decrease_factor: float = 1.25, newton_regularization_max_multiplier: float = 10000.0)[source]#

Bases: object

Parameters:

  • gradient_norm_tolerance (float)

  • max_iterations (int)

  • max_log_std_update (float)

  • max_mean_update_std (float)

  • min_variational_std (float)

  • max_variational_std (float)

  • newton_metric (str)

  • newton_fallback_learning_rate (float)

  • newton_regularization (float)

  • newton_hessian_type (str)

  • newton_curvature_strategy (str)

  • newton_hessian_num_samples (int)

  • newton_hessian_averaging_factor (float)

  • newton_regularization_strategy (str)

  • newton_additive_regularization (float)

  • newton_regularization_epsilon (float)

  • newton_adaptive_regularization (bool)

  • newton_regularization_increase_factor (float)

  • newton_regularization_decrease_factor (float)

  • newton_regularization_max_multiplier (float)

class romtools.workflows.inverse.vi_optimization_methods.VIStochasticNonmonotoneLineSearchConfig(initial_step_size: float = 0.01, max_step_size: float = inf, step_size_growth_factor: float = 1.05, step_size_decay_factor: float = 2.0, max_step_size_decrease_trys: int = 15, relaxation_parameter: float = 3.05, line_search_objective: str = 'elbo', line_search_nonmonotone_window: int = 5, line_search_armijo_coefficient: float = 1e-06, line_search_uncertainty_sigma: float = 3.0, line_search_sample_growth_factor: float = 1.0, log_std_learning_rate_factor: float = 1.0)[source]#

Bases: object

Parameters:

  • initial_step_size (float)

  • max_step_size (float)

  • step_size_growth_factor (float)

  • step_size_decay_factor (float)

  • max_step_size_decrease_trys (int)

  • relaxation_parameter (float)

  • line_search_objective (str)

  • line_search_nonmonotone_window (int)

  • line_search_armijo_coefficient (float)

  • line_search_uncertainty_sigma (float)

  • line_search_sample_growth_factor (float)

  • log_std_learning_rate_factor (float)