augmented_lagrangian

Augmented-Lagrangian controller for a single scalar inequality constraint.

Classes:

Name	Description
`AugmentedLagrangianInequalityConstraint`	Powell-Hestenes-Rockafellar augmented-Lagrangian controller for a single scalar inequality constraint `g <= tolerance`.

AugmentedLagrangianInequalityConstraint ¶

Bases: Module

Powell-Hestenes-Rockafellar augmented-Lagrangian controller for a single scalar inequality constraint g <= tolerance.

Maintains the dual state used to enforce one inequality constraint as a penalty added to a training objective. The constraint violation is v = g - tolerance; the penalty is

``penalty = multiplier * relu(v) + (penalty_coefficient / 2) * relu(v) ** 2``

where multiplier (the dual variable, often written lambda) and penalty_coefficient (the quadratic coefficient, often written rho) are registered buffers. Only the active part relu(v) enters the penalty, so a satisfied constraint contributes nothing and never pushes the objective negative. The multiplier is advanced by dual ascent from the mean accumulated violation:

``multiplier <- clamp(multiplier + penalty_coefficient * mean_violation, min=0)``.

The two public methods split across the training step because the penalty enters the backward pass while the multiplier update mutates it:

penalty runs in the forward pass; its output is added to the loss and backpropagated.
observe runs after backward (the multiplier is a saved-for-backward multiplicand, so mutating it earlier would trip autograd's in-place version check). It records the violation and applies the dual-ascent update on the configured cadence, or immediately with force=True (e.g. at an epoch boundary).

All state — multiplier, penalty_coefficient, and the running violation/example accumulators — is held in buffers, so a state_dict checkpoint captures a partially accumulated window and resume is well defined. The module performs no cross-replica reduction: it updates its own multiplier from the values it is given, so under data-parallel training each replica maintains its own multiplier and accumulators unless the caller reduces constraint_value across replicas before calling.

relax selects the dual-update regime. With relax=True (default) the raw violation is accumulated, so a feasible window (mean violation below zero) drives the multiplier back down toward zero — the textbook PHR inequality update — letting the constrained quantity settle at the tolerance boundary rather than staying over-constrained after a transient violation. With relax=False only active (positive) violations are accumulated, so the multiplier is a one-way ratchet that never decreases — classic dual (sub)gradient ascent. Use relax=False when you want to be strict about the constraint: once the penalty ramps up to enforce the bound it never backs off, keeping the constrained quantity firmly within tolerance even after it becomes feasible.

The penalty coefficient is held constant (no escalation schedule).

Parameters:

Name	Type	Description	Default
`tolerance` ¶	`float`	The constraint bound `tolerance` in `g <= tolerance`.	required
`multiplier_init` ¶	`float`	Initial value of the dual multiplier.	`0.0`
`penalty_coefficient` ¶	`float`	The (constant) quadratic penalty coefficient. Must be positive.	`0.5`
`relax` ¶	`bool`	If `True` (default), accumulate raw violations so the multiplier may decrease during feasible windows (the PHR inequality update). If `False`, accumulate only active violations so the multiplier never decreases — a strict, one-way ratchet for firmly enforcing the constraint.	`True`

Raises:

Type	Description
`ValueError`	If `penalty_coefficient` is not positive.

Example

import torch from stainedglass_core.loss import augmented_lagrangian constraint = augmented_lagrangian.AugmentedLagrangianInequalityConstraint( ... tolerance=0.2, penalty_coefficient=0.5 ... )

A value within tolerance incurs no penalty.¶

penalty, violation = constraint.penalty(torch.tensor(0.1)) float(penalty) 0.0

A value above tolerance incurs a quadratic penalty; add it to the loss, then, after¶

backward, record the violation and apply the dual-ascent update.¶

penalty, violation = constraint.penalty(torch.tensor(0.6)) round(float(penalty), 4) 0.04 constraint.observe(violation, force=True) True round( ... float(constraint.multiplier), 4 ... ) # multiplier rose to enforce the constraint 0.2

Added in version v3.53.0.

Methods:

Name	Description
`observe`	Record a violation and apply the dual-ascent update on cadence.
`penalty`	Compute the augmented-Lagrangian penalty and the signed violation.

observe ¶

observe(
    violation: Tensor | None = None,
    *,
    num_examples: int = 1,
    update_examples: int = 0,
    world_size: int = 1,
    force: bool = False
) -> bool

Record a violation and apply the dual-ascent update on cadence.

Call after backward so the in-place multiplier update does not corrupt the penalty's autograd graph. Typical use records each step's violation and updates every update_examples examples; pass force=True (with or without a new violation) to flush the accumulator immediately, e.g. at an epoch boundary.

With relax=False only the active (relu'd) part of violation is accumulated, so satisfied-constraint windows cannot lower the multiplier; with relax=True the raw violation is accumulated. The dual-ascent update averages the accumulated violations over the number of recorded observations (not weighted by num_examples); num_examples only advances the example counter that governs the update_examples cadence.

Parameters:

Name	Type	Description	Default
`violation` ¶	`Tensor \| None`	The signed violation to record, as returned by `penalty`. `None` records nothing (use with `force=True` to flush a pending window without a new observation).	`None`
`num_examples` ¶	`int`	Number of examples represented by `violation`, added to the example counter that governs the `update_examples` cadence. Must be positive. Ignored when `violation` is `None`.	`1`
`update_examples` ¶	`int`	Global number of examples between cadence updates. `0` or less disables the example-driven update (rely on `force` for an epoch-boundary flush instead).	`0`
`world_size` ¶	`int`	Number of data-parallel replicas sharing the global example budget. Must be positive.	`1`
`force` ¶	`bool`	If `True`, apply the dual-ascent update immediately regardless of the cadence.	`False`

Returns:

Type	Description
`bool`	`True` if a dual-ascent update was applied on this call; `False` otherwise.

Raises:

Type	Description
`ValueError`	If `violation` is given and `num_examples` is not positive, or if a cadence update is attempted and `world_size` is not positive.

penalty ¶

penalty(
    constraint_value: Tensor,
) -> tuple[torch.Tensor, torch.Tensor]

Compute the augmented-Lagrangian penalty and the signed violation.

Call in the forward pass and add the returned penalty to the loss. The penalty is differentiable with respect to constraint_value for active violations, providing the gradient that drives the constrained quantity back toward the tolerance.

Parameters:

Name	Type	Description	Default
`constraint_value` ¶	`Tensor`	The current value of the constrained quantity `g`.	required

Returns:

Type	Description
`tuple[torch.Tensor, torch.Tensor]`	A tuple `(penalty, violation)` where `penalty` is the scalar penalty term to add to the
`tuple[torch.Tensor, torch.Tensor]`	objective and `violation` is the signed violation `g - tolerance` (pass it to
`tuple[torch.Tensor, torch.Tensor]`	`observe` after `backward` for the dual update).

augmented_lagrangian