Skip to content

← deep_learning package

AdamaxWStep

The AdamaxW optimizer step (adamax with weight decay).

Based on Loshchilov et al, 2019 (but see the note about the weight decay parameter), AdamaxW is a variant of the Adamax optimizer that additionally regularizes weights to have small l2 norm, which helps with generalization, especially when the amount of training data is low relative to the number of parameters. This is analogous to l2 regularization terms on the weights, but works better with adaptive gradient algorithms like Adam. Like all step nodes, this node only processes gradients, and the resulting updates must be applied manually to the weights (this can be accomplished using the Add node). However, you can also pass it to the StepSolver node which implements the full optimization loop. The learning rate can instead be given as a schedule, by wiring one of the Schedule nodes into the learning_rate_schedule port. The weight decay can be used in conjunction with a mask data structure that has the same nested structure as the weights being optimized, but which contains booleans indicating which weights should be decayed. More Info... Version 0.2.0

Ports/Properties

gradients

Gradients to be transformed.

  • verbose name: Gradients
  • default value: None
  • port type: DataPort
  • value type: object (can be None)
  • data direction: INOUT

weights

Optional current weights.

  • verbose name: Weights
  • default value: None
  • port type: DataPort
  • value type: object (can be None)
  • data direction: IN

state

Explicit state of the node.

  • verbose name: State
  • default value: None
  • port type: DataPort
  • value type: object (can be None)
  • data direction: INOUT

learning_rate_schedule

Optional learning rate schedule.

  • verbose name: Learning Rate Schedule
  • default value: None
  • port type: DataPort
  • value type: BaseNode (can be None)
  • data direction: IN

weight_decay_mask

Mask structure for the weight decay.

  • verbose name: Weight Decay Mask
  • default value: None
  • port type: DataPort
  • value type: object (can be None)
  • data direction: IN

learning_rate

Learning rate. A typical choice may be 0.001 here, but this is problem dependent. If a learning rate schedule is provided, this value should be left unspecified.

  • verbose name: Learning Rate
  • default value: None
  • port type: FloatPort
  • value type: float (can be None)

beta1

Exponential decay rate for the first moment estimates.

  • verbose name: Beta1
  • default value: 0.9
  • port type: FloatPort
  • value type: float (can be None)

beta2

Exponential decay rate for the second moment estimates.

  • verbose name: Beta2
  • default value: 0.999
  • port type: FloatPort
  • value type: float (can be None)

epsilon

Small value applied to the denominator outside the square root to avoid dividing by zero when rescaling. Note that larger epsilon values have been explored in the literature.

  • verbose name: Epsilon
  • default value: 1e-08
  • port type: FloatPort
  • value type: float (can be None)

epsilon_inroot

Small value applied to the denominator inside the square root to avoid dividing by zero when rescaling. A case where this is needed is when differentiating the optimizer itself, eg for bilevel optimization.

  • verbose name: Epsilon (Inside Root)
  • default value: 0
  • port type: FloatPort
  • value type: float (can be None)

weight_decay

Strength of the weight decay. This is multiplied by the learning rate as in e.g., PyTorch and Optax, but differs from the paper, where it is only multiplied by the schedule multiplier but not the base learning rate.

  • verbose name: Weight Decay
  • default value: 0.0001
  • port type: FloatPort
  • value type: float (can be None)

mu_precision

Numeric precision for the first-order accumulator. Keep resolves to the precision of the inputs.

  • verbose name: Mu Precision
  • default value: keep
  • port type: EnumPort
  • value type: str (can be None)

set_breakpoint

Set a breakpoint on this node. If this is enabled, your debugger (if one is attached) will trigger a breakpoint.

  • verbose name: Set Breakpoint (Debug Only)
  • default value: False
  • port type: BoolPort
  • value type: bool (can be None)

metadata

User-definable meta-data associated with the node. Usually reserved for technical purposes.

  • verbose name: Metadata
  • default value: {}
  • port type: DictPort
  • value type: dict (can be None)