Cautious AdamW

Cautious AdamW: Consensus-Based Update Masking.

Cautious AdamW applies a gradient clipping pre-filter before the standard AdamW step. By clamping gradient magnitudes to a configurable threshold, it masks outlier updates that would otherwise destabilize training. The paper reports up to 1.47x training speedup on large-scale models by reducing the number of harmful large updates.

Reference:

Liang, K., Zhou, L., Liu, B., Zhao, L., Jiang, Y., Pan, S., Zhang, R., & Bengio, Y. (2024). Cautious Optimizers: Improving Training with One Line of Code. arXiv:2411.16085. https://arxiv.org/abs/2411.16085

class src.model.optimizer.cautious_adamw.CautiousAdamW(*args: Any, **kwargs: Any)[source]

Bases: AdamW

Cautious AdamW with gradient clipping pre-filter.

Extends AdamW by clamping each parameter’s gradient to [-cautious_clip, cautious_clip] before invoking the base optimizer step. This acts as a consensus-based mask that suppresses outlier gradient components.

State buffers (per parameter):

exp_avg: First-moment estimate \(m_t\) (1 buffer). exp_avg_sq: Second-moment estimate \(v_t\) (1 buffer). Total: 2 buffers, O(2n) memory.

Parameters:

  • params – Iterable of parameters or parameter groups.

  • *args – Positional arguments forwarded to torch.optim.AdamW.

  • cautious_clip – Maximum absolute gradient value allowed before the AdamW step. Set to 0 to disable clipping (default 1.0).

  • **kwargs – Keyword arguments forwarded to torch.optim.AdamW.

cautious_clip

The gradient clipping threshold.

Type:

float

Reference:

Liang, K., Zhou, L., Liu, B., Zhao, L., Jiang, Y., Pan, S., Zhang, R., & Bengio, Y. (2024). Cautious Optimizers: Improving Training with One Line of Code. arXiv:2411.16085.

step(closure=None)

Perform a single optimization step with gradient pre-clipping.

Clamps all gradients to [-cautious_clip, cautious_clip] in-place, then delegates to the base AdamW.step().

Parameters:

closure – Optional callable that reevaluates the model and returns the loss. Passed through to the base optimizer.

Returns:

The loss value returned by closure (if provided), or None.