torchref.refinement.optimizers.exploratory_lbfgs module

ExploratoryLBFGS: LBFGS optimizer with automatic landscape exploration.

After convergence, uses Lanczos eigenanalysis to find flat/negative-curvature directions in the Hessian, scans along them for alternative basins, and hops to better ones. Useful for escaping local minima in crystallographic refinement where parameter degeneracies are common.

State machine: OPTIMIZING -> CONVERGED -> EXPLORING -> HOPPING -> OPTIMIZING

class torchref.refinement.optimizers.exploratory_lbfgs.OptimizerPhase(value)

Bases: Enum

State machine phases for ExploratoryLBFGS.

OPTIMIZING = 1

CONVERGED = 2

EXPLORING = 3

HOPPING = 4

class torchref.refinement.optimizers.exploratory_lbfgs.Mode(eigenvalue, direction, index, is_negative)

Bases: object

A Hessian eigenmode from Lanczos analysis.

eigenvalue: float

direction: Tensor

index: int

is_negative: bool

__init__(eigenvalue, direction, index, is_negative)

class torchref.refinement.optimizers.exploratory_lbfgs.ScanPoint(t, loss)

Bases: object

A single evaluation point along a scan direction.

t: float

loss: float

__init__(t, loss)

class torchref.refinement.optimizers.exploratory_lbfgs.Basin(t, loss, direction, loss_improvement)

Bases: object

A detected basin along a scan direction.

t: float

loss: float

direction: Tensor

loss_improvement: float

__init__(t, loss, direction, loss_improvement)

class torchref.refinement.optimizers.exploratory_lbfgs.ParameterGroup(indices, direction, eigenvalue, mode_index)

Bases: object

Group of parameters with significant participation in a mode.

indices: List[int]

direction: Tensor

eigenvalue: float

mode_index: int

__init__(indices, direction, eigenvalue, mode_index)

class torchref.refinement.optimizers.exploratory_lbfgs.ConvergenceTracker(grad_threshold=1e-05, loss_threshold=1e-07, param_threshold=1e-06, n_stable=3, grad_norms=<factory>, loss_changes=<factory>, param_changes=<factory>, prev_loss=None, prev_params=None)

Bases: object

Tracks convergence criteria over multiple steps.

grad_threshold: float = 1e-05

loss_threshold: float = 1e-07

param_threshold: float = 1e-06

n_stable: int = 3

grad_norms: List[float]

loss_changes: List[float]

param_changes: List[float]

prev_loss: float | None = None

prev_params: Tensor | None = None

update(grad_norm, loss, params)

Record a new step’s metrics.

property is_converged: bool

Check if all criteria are met for n_stable consecutive steps.

reset()

Reset all tracking state.

__init__(grad_threshold=1e-05, loss_threshold=1e-07, param_threshold=1e-06, n_stable=3, grad_norms=<factory>, loss_changes=<factory>, param_changes=<factory>, prev_loss=None, prev_params=None)

class torchref.refinement.optimizers.exploratory_lbfgs.ExplorationResult(cycle, n_modes_found, n_negative_modes, n_degenerate_modes, n_basins_found, best_basin, hopped, eigenvalues=None)

Bases: object

Diagnostics from one exploration cycle.

cycle: int

n_modes_found: int

n_negative_modes: int

n_degenerate_modes: int

n_basins_found: int

best_basin: Basin | None

hopped: bool

eigenvalues: List[float] | None = None

__init__(cycle, n_modes_found, n_negative_modes, n_degenerate_modes, n_basins_found, best_basin, hopped, eigenvalues=None)

exception torchref.refinement.optimizers.exploratory_lbfgs.LanczosError

Bases: Exception

Raised when Lanczos iteration fails.

class torchref.refinement.optimizers.exploratory_lbfgs.ExploratoryLBFGS(params, lr=1.0, max_iter=20, history_size=100, m_modes=10, m_lanczos_iter=None, eigenvalue_threshold=0.01, participation_threshold=0.05, scan_points=20, scan_step_size=0.1, max_exploration_cycles=5, hvp_epsilon=0.0001, convergence_grad_threshold=1e-05, convergence_loss_threshold=1e-07, convergence_param_threshold=1e-06, n_stable=3, verbose=1)

Bases: Optimizer

LBFGS optimizer with automatic landscape exploration via Lanczos analysis.

Composes with (rather than subclasses) torch.optim.LBFGS. After the internal LBFGS converges, performs eigenanalysis of the Hessian to find degenerate/flat directions, scans along them, and hops to better basins if found.

Parameters:

• params (iterable) – Parameters to optimize.

• lr (float) – LBFGS learning rate. Default: 1.0.

• max_iter (int) – LBFGS max line search iterations per step. Default: 20.

• history_size (int) – LBFGS Hessian approximation memory. Default: 100.

• m_modes (int) – Number of lowest eigenmodes to compute. Default: 10.

• m_lanczos_iter (int, optional) – Lanczos iterations. Default: 2*m_modes + 10.

• eigenvalue_threshold (float) – Mode is degenerate if eigenvalue < threshold * median(positive). Default: 0.01.

• participation_threshold (float) – Parameter participates if |component| > threshold * ||mode||. Default: 0.05.

• scan_points (int) – Evaluation points per scan direction. Default: 20.

• scan_step_size (float) – Step size in parameter space units. Default: 0.1.

• max_exploration_cycles (int) – Cap on explore-hop cycles. Default: 5.

• hvp_epsilon (float) – Finite-difference epsilon for Hessian-vector products. Default: 1e-4.

• convergence_grad_threshold (float) – Gradient norm convergence threshold. Default: 1e-5.

• convergence_loss_threshold (float) – Loss change convergence threshold. Default: 1e-7.

• convergence_param_threshold (float) – Parameter change convergence threshold. Default: 1e-6.

• n_stable (int) – Consecutive converged steps required. Default: 3.

• verbose (int) – Verbosity level: 0=silent, 1=summary, 2=detailed. Default: 1.

__init__(params, lr=1.0, max_iter=20, history_size=100, m_modes=10, m_lanczos_iter=None, eigenvalue_threshold=0.01, participation_threshold=0.05, scan_points=20, scan_step_size=0.1, max_exploration_cycles=5, hvp_epsilon=0.0001, convergence_grad_threshold=1e-05, convergence_loss_threshold=1e-07, convergence_param_threshold=1e-06, n_stable=3, verbose=1)

exploration_history: List[ExplorationResult]

property phase: OptimizerPhase

Current optimizer phase.

step(closure)

Perform one step of the state machine.

Parameters:

closure (callable) – A closure that re-evaluates the model loss. Should call optimizer.zero_grad(), compute the loss, call loss.backward(), and return the loss.

Returns:

The loss value.

Return type:

float or None