iALM

class humancompatible.train.dual_optim.iALM(m: int = None, beta: float = 1.0, sigma: float = 1.0, gamma: float = 1.0, init_duals: float | Tensor = None, penalty: float = 1.0, *, dual_range: Tuple[float, float] = (-100.0, 100.0), momentum: float = 0.0, dampening: float = 0.0, is_ineq: bool = False, device=None)

A Dual Optimizer that works on the dual maximization tasks according to the Augmented Lagrangian rule, with adaptive stepsize based on https://doi.org/10.1007/s10589-023-00521-z, Algorithm 1. Creates and updates dual variables.

\[ \begin{align}\begin{aligned}\pmb{\lambda}_{t+1} & \leftarrow \pmb{\lambda}_t + \min\left\{ \beta_k, \frac{\gamma_k}{\|\mathbf{c}_t(\theta_t)\|} \right\} \mathbf{c}_t(\theta_{t})\\\mathcal{L}_{t+1} & \leftarrow f_t(\theta_{t}) + \pmb{\lambda}_{t+1}^T \mathbf{c}_t(\theta_{t}) + \frac{\rho}{2} \| \mathbf{c}_t(\theta_{t}) \|^2_2\end{aligned}\end{align} \]

Parameters:

m (int) – Number of constraints (determines the number of dual variables to create)
beta (float) – Dual variable update rate.
sigma (float) – Multiplier for increasing`beta`.
gamma (float) – Penalty update parameter.
init_duals (float | Tensor) – Initial values for the new dual variables. Defaults to 0 for all.
penalty (float) – Augmented Lagrangian penalty parameter. Defaults to`1.`
dual_range (Tuple[float, float]) – Safeguarding range for dual variables; they will be`clamp`-ed to this range.
momentum (float) – Momentum/Smoothing factor for dual variables. Equivalent to SGD momentum. Set to 0 to disable.
dampening (float) – Dampening for momentum. Equivalent to SGD dampening. Set to 0 to disable.
is_ineq (bool) – Whether to treat the constraints as equality or inequality. If`True`, dual variables will be decreased on strict satisfaction and lower-bounded by max(dual_range[0], 0).
ctol (float) – Constraint tolerance; allows tiny violations of constraints to account for noise.

add_constraint_group(m: int = None, beta: float = 1.0, sigma: float = 1.0, gamma: float = 1.0, momentum: float = None, dampening: float = None, init_duals: Tensor = None, dual_range: tuple[float, float] = None, is_ineq: bool = False, device=None) → None

Allows to add a group of dual variables with separate initial values and learning rates.

Parameters:

m (int) – Size of group (number of dual variables to add)
beta (float) – Dual variable update rate
sigma (float) – Multiplier for increasing beta
gamma (float) – Penalty update parameter
momentum (float) – Momentum for dual variable updates
dampening (float) – Dampening for momentum
init_duals (Tensor) – Initial values for the new dual variables
dual_range (Tuple[float, float]) – After each dual update, the dual variables will be clamped to this range.
is_ineq (bool) – Whether to treat the constraints as equality or inequality. If`True`, dual variables will be relaxed on strict satisfaction and lower-bounded by max(dual_range[0], 0).

property duals: Tensor

Returns:: Dual variables, concatenated into a single tensor.
Return type:: Tensor

forward(loss: Tensor, constraints: Tensor) → Tensor

Calculates and returns the Augmented Lagrangian.

Parameters:

loss (Tensor) – Loss (objective function) value
constraints (Tensor) – Tensor of constraint values

Returns:

Lagrangian

Return type:

Tensor

forward_update(loss: Tensor, constraints: Tensor) → Tensor

Combines forward and update; slightly faster.

Parameters:

loss (Tensor) – Loss (objective function) value
constraints (Tensor) – Tensor of constraint values

Returns:

Lagrangian

Return type:

Tensor

load_state_dict(state_dict: dict[str, Any]) → None

Load the optimizer state.

Args:

state_dict (dict): optimizer state. Should be an object returned: from a call to state_dict().

Warning

Make sure this method is called after initializing torch.optim.lr_scheduler.LRScheduler, as calling it beforehand will overwrite the loaded learning rates.

Note

The names of the parameters (if they exist under the “param_names” key of each param group in state_dict()) will not affect the loading process. To use the parameters’ names for custom cases (such as when the parameters in the loaded state dict differ from those initialized in the optimizer), a custom register_load_state_dict_pre_hook should be implemented to adapt the loaded dict accordingly. If param_names exist in loaded state dict param_groups they will be saved and override the current names, if present, in the optimizer state. If they do not exist in loaded state dict, the optimizer param_names will remain unchanged.

Example:

>>> # xdoctest: +SKIP
>>> model = torch.nn.Linear(10, 10)
>>> optim = torch.optim.SGD(model.parameters(), lr=3e-4)
>>> scheduler1 = torch.optim.lr_scheduler.LinearLR(
...     optim,
...     start_factor=0.1,
...     end_factor=1,
...     total_iters=20,
... )
>>> scheduler2 = torch.optim.lr_scheduler.CosineAnnealingLR(
...     optim,
...     T_max=80,
...     eta_min=3e-5,
... )
>>> lr = torch.optim.lr_scheduler.SequentialLR(
...     optim,
...     schedulers=[scheduler1, scheduler2],
...     milestones=[20],
... )
>>> lr.load_state_dict(torch.load("./save_seq.pt"))
>>> # now load the optimizer checkpoint after loading the LRScheduler
>>> optim.load_state_dict(torch.load("./save_optim.pt"))

state_dict() → dict[str, Any]

Return the state of the optimizer as a dict.

It contains two entries:

state: a Dict holding current optimization state. Its content
differs between optimizer classes, but some common characteristics hold. For example, state is saved per parameter, and the parameter itself is NOT saved. state is a Dictionary mapping parameter ids to a Dict with state corresponding to each parameter.
param_groups: a List containing all parameter groups where each
parameter group is a Dict. Each parameter group contains metadata specific to the optimizer, such as learning rate and weight decay, as well as a List of parameter IDs of the parameters in the group. If a param group was initialized with named_parameters() the names content will also be saved in the state dict.

NOTE: The parameter IDs may look like indices but they are just IDs associating state with param_group. When loading from a state_dict, the optimizer will zip the param_group params (int IDs) and the optimizer param_groups (actual nn.Parameter s) in order to match state WITHOUT additional verification.

A returned state dict might look something like:

{
    'state': {
        0: {'momentum_buffer': tensor(...), ...},
        1: {'momentum_buffer': tensor(...), ...},
        2: {'momentum_buffer': tensor(...), ...},
        3: {'momentum_buffer': tensor(...), ...}
    },
    'param_groups': [
        {
            'lr': 0.01,
            'weight_decay': 0,
            ...
            'params': [0]
            'param_names' ['param0']  (optional)
        },
        {
            'lr': 0.001,
            'weight_decay': 0.5,
            ...
            'params': [1, 2, 3]
            'param_names': ['param1', 'layer.weight', 'layer.bias'] (optional)
        }
    ]
}

step(constraints: Tensor) → None

Updates the dual variables

Parameters:: constraints (Tensor) – Tensor of constraint values

update(constraints: Tensor) → None

Updates the dual variables

Parameters:: constraints (Tensor) – Tensor of constraint values