Optimizers¶

The batched optimizers of the candidate fit, selected by train.optimizer.name and configured by train.optimizer.optimizer_params. See Optimizers.

BaseOptimizer ¶

BaseOptimizer(shared_param_names: Iterable[str], lr: LRLike, clip_value: float | None = None, clip_norm: float | None = None, shared_param_lr_multiplier: float = 1.0)

Bases: ABC

Base class of the ensemble-batched optimizers.

Every parameter tensor carries the candidate ensemble on its leading dimension, one row per candidate neuron, and a subclass's step updates all rows at once. Parameters named in shared_param_names have no ensemble dimension; the subclasses update them once per step from the gradient averaged over the active members, at the mean of those members' learning rates times shared_param_lr_multiplier.

Parameters:

Name	Type	Description	Default
`shared_param_names`	`iterable of str`	Names of the parameters shared by every ensemble member.	required
`lr`	`float or (ensemble_size,) tensor`	Learning rate, one per member when a tensor.	required
`clip_value`	`float`	Clamp every gradient element to [-clip_value, clip_value].	`None`
`clip_norm`	`float`	Rescale each member's gradient to at most this L2 norm.	`None`
`shared_param_lr_multiplier`	`float`	Factor on the learning rate of the shared parameters. Default 1.0.	`1.0`

state_dict ¶

state_dict() -> dict[str, Any]

Return the optimizer state as a dict; implemented by subclasses.

load_state_dict ¶

load_state_dict(state_dict: dict[str, Any]) -> None

Restore the state written by state_dict; implemented by subclasses.

gradient_clipping ¶

gradient_clipping(g: Tensor) -> torch.Tensor

Clip a batched gradient by value, then by per-member norm.

Parameters:

Name	Type	Description	Default
`g`	`(B, ...) tensor`	Gradient with the ensemble on the leading dimension.	required

Returns:

Name	Type	Description
`g`	`(B, ...) tensor`	`g` clamped to +-`clip_value` (when set), then each member's slice scaled down to L2 norm `clip_norm` (when set). Unchanged when both are None.

BatchedLBFGS ¶

BatchedLBFGS(params: dict[str, Tensor], shared_param_names: Iterable[str], lr: LRLike, history_size: int = 10, clip_value: float | None = None, clip_norm: float | None = None, shared_param_lr_multiplier: float = 1.0, max_step: float | None = 1.0, curvature_eps: float | None = 1e-08)

Bases: BaseOptimizer

Limited-memory BFGS, batched over the candidate ensemble.

Each member keeps its own history of the last history_size correction pairs s = delta theta, y = delta g per parameter tensor. The two-loop recursion turns the clipped gradient into a quasi-Newton direction d ~ H^-1 g, and the step is

theta = theta - lr*d

with no line search. A pair is stored only if it passes the curvature condition (curvature_eps), and each member's update of a parameter tensor is capped at norm max_step. Members with active_mask False in step keep their parameters and store no pair. Shared parameters keep one history, driven by the gradient averaged over the active members. stats counts steps, pairs, rejected pairs, updates and capped updates for the trainer's log.

Parameters:

Name	Type	Description	Default
`params`	`dict of str -> (B, ...) tensor`	Initial parameters.	required
`shared_param_names`	`iterable of str`	See `BaseOptimizer`.	required
`lr`	`float or (B,) tensor`	See `BaseOptimizer`.	required
`clip_value`	`float or None`	See `BaseOptimizer`.	`None`
`clip_norm`	`float or None`	See `BaseOptimizer`.	`None`
`shared_param_lr_multiplier`	`float`	See `BaseOptimizer`.	`1.0`
`history_size`	`int`	Correction pairs kept per member. Default 10.	`10`
`max_step`	`float or None`	Cap on the norm of one member's update of one parameter tensor, in parameter units (a crude trust region); None for no cap. Default 1.0.	`1.0`
`curvature_eps`	`float or None`	Store a pair only if y.s > curvature_eps\|s\|\|y\| (the BFGS curvature condition, with the standard skipping rule); None stores every pair with s != 0. Default 1e-8.	`1e-08`

Notes

Why the two guards. The two-loop recursion scales its direction by s.y / y.y and takes a fixed step. On a direction whose gradient and curvature vanish together (an RBF bump losing its mass) that scale is a ratio of two vanishing numbers and the step does not shrink with the gradient: one step can throw centres tens of standard units off the data. A pair with y.s <= 0 makes the inverse-Hessian estimate indefinite and the direction can point uphill. The cap bounds the damage of any single step; the curvature test keeps the estimate positive definite. Neither changes a well-conditioned convex fit (the polynomial families) beyond its first steps.

Storage. For the batched (per-member) parameters the correction history is a pair of tensors s_hist[k], y_hist[k] of shape (B, history_size, P) plus a per-member fill count hist_count[k]; the newest pair sits in the last slot and the valid slots are the last hist_count ones. The two-loop recursion runs as history_size batched operations over the whole ensemble instead of a Python loop over its members. Shared parameters keep a single deque history and the single-vector recursion.

step ¶

step(params: dict[str, Tensor], grads: dict[str, Tensor], active_mask: Tensor | None = None) -> dict[str, torch.Tensor]

Take one optimizer step for every active member.

Parameters:

Name	Type	Description	Default
`params`	`dict of str -> (B, ...) tensor`	Current parameters, with the members on the leading dimension.	required
`grads`	`dict of str -> (B, ...) tensor`	Their gradients, with the same keys and shapes.	required
`active_mask`	`(B,) bool tensor`	True for the members that step; None steps every member.	`None`

Returns:

Type	Description
`dict of str -> (B, ...) tensor`	The new parameters, with the same keys and shapes.

state_dict ¶

state_dict() -> dict[str, Any]

Snapshot of curvature history + prev params/grads + lr so a resumed run keeps the L-BFGS approximation rather than starting from identity.

load_state_dict ¶

load_state_dict(state_dict: dict[str, Any]) -> None

Restore the state written by state_dict; every field must be present.

Shared parameters get their single deque history back; batched ones their (B, history_size, P) correction tensors and per-member fill counts.

BatchedAdam ¶

BatchedAdam(params: dict[str, Tensor], shared_param_names: Iterable[str], lr: LRLike, betas: tuple[float, float] = (0.9, 0.999), eps: float = 1e-08, clip_value: float | None = None, clip_norm: float | None = None, shared_param_lr_multiplier: float = 1.0)

Bases: BaseOptimizer

Adam, batched over the candidate ensemble.

Per member and parameter tensor, with g the clipped gradient and t the step count (shared by all members):

m = b1*m + (1 - b1)*g,      v = b2*v + (1 - b2)*g^2
theta = theta - lr * m_hat / (sqrt(v_hat) + eps)
m_hat = m / (1 - b1^t),     v_hat = v / (1 - b2^t)

Members with active_mask False in step keep their parameters; their moment buffers still advance.

Parameters:

Name	Type	Description	Default
`params`	`dict of str -> (B, ...) tensor`	Initial parameters; they set the shapes of the moment buffers.	required
`shared_param_names`	`iterable of str`	See `BaseOptimizer`.	required
`lr`	`float or (B,) tensor`	See `BaseOptimizer`.	required
`clip_value`	`float or None`	See `BaseOptimizer`.	`None`
`clip_norm`	`float or None`	See `BaseOptimizer`.	`None`
`shared_param_lr_multiplier`	`float`	See `BaseOptimizer`.	`1.0`
`betas`	`(float, float)`	Decay rates b1, b2 of the moment estimates. Default (0.9, 0.999).	`(0.9, 0.999)`
`eps`	`float`	Added to sqrt(v_hat) in the denominator. Default 1e-8.	`1e-08`

step ¶

step(params: dict[str, Tensor], grads: dict[str, Tensor], active_mask: Tensor | None = None, b: Any = None) -> dict[str, torch.Tensor]

Take one optimizer step for every active member.

Parameters:

Name	Type	Description	Default
`params`	`dict of str -> (B, ...) tensor`	Current parameters, with the members on the leading dimension.	required
`grads`	`dict of str -> (B, ...) tensor`	Their gradients, with the same keys and shapes.	required
`active_mask`	`(B,) bool tensor`	True for the members that step; None steps every member.	`None`
`b`	`any`	Not used.	`None`

Returns:

Type	Description
`dict of str -> (B, ...) tensor`	The new parameters, with the same keys and shapes.

state_dict ¶

state_dict() -> dict[str, Any]

Return lr, betas, eps, the moment buffers, the step count and the shared-parameter settings.

load_state_dict ¶

load_state_dict(state_dict: dict[str, Any]) -> None

Restore the state written by state_dict; every field must be present.

BatchedSGD ¶

BatchedSGD(params: dict[str, Tensor], shared_param_names: Iterable[str], lr: LRLike, momentum: float = 0.9, weight_decay: float = 0.0, nesterov: bool = True, clip_value: float | None = None, clip_norm: float | None = None, shared_param_lr_multiplier: float = 1.0)

Bases: BaseOptimizer

SGD with momentum, batched over the candidate ensemble.

Per member and parameter tensor, with g the gradient plus weight_decay*theta, then clipped:

v = momentum*v + g
theta = theta - lr*(momentum*v + g)     (nesterov=True)
theta = theta - lr*v                    (nesterov=False)

Members with active_mask False in step keep their parameters; their momentum buffer still advances.

Parameters:

Name	Type	Description	Default
`params`	`dict of str -> (B, ...) tensor`	Initial parameters; they set the shapes of the momentum buffers.	required
`shared_param_names`	`iterable of str`	See `BaseOptimizer`.	required
`lr`	`float or (B,) tensor`	See `BaseOptimizer`.	required
`clip_value`	`float or None`	See `BaseOptimizer`.	`None`
`clip_norm`	`float or None`	See `BaseOptimizer`.	`None`
`shared_param_lr_multiplier`	`float`	See `BaseOptimizer`.	`1.0`
`momentum`	`float`	Momentum coefficient. Default 0.9.	`0.9`
`weight_decay`	`float`	L2 penalty added to the gradient (coupled, not decoupled). Default 0.0.	`0.0`
`nesterov`	`bool`	Use the Nesterov form of the update. Default True.	`True`

step ¶

step(params: dict[str, Tensor], grads: dict[str, Tensor], active_mask: Tensor | None = None) -> dict[str, torch.Tensor]

Take one optimizer step for every active member.

Parameters:

Name	Type	Description	Default
`params`	`dict of str -> (B, ...) tensor`	Current parameters, with the members on the leading dimension.	required
`grads`	`dict of str -> (B, ...) tensor`	Their gradients, with the same keys and shapes.	required
`active_mask`	`(B,) bool tensor`	True for the members that step; None steps every member.	`None`

Returns:

Type	Description
`dict of str -> (B, ...) tensor`	The new parameters, with the same keys and shapes.

state_dict ¶

state_dict() -> dict[str, Any]

Return lr, weight decay, the momentum settings and buffers, and the shared-parameter settings.

load_state_dict ¶

load_state_dict(state_dict: dict[str, Any]) -> None

Restore the state written by state_dict; every field must be present.