Skip to content

← deep_learning package

RMSNorm

Apply RMS normalization to the given data.

This standardizes data across the spatial and/or feature dimensions, separately per instance (example), but does not remove the mean (unlike LayerNorm, which does remove the mean). RMS normalization can be the best choice when dealing with data that has a small batch size, or when used inside an RNN (both cases rendering the batch norm potentially inapplicable) and/or when no additional large spatial axes are present (rendering the instance norm inapplicable). Like most normalizations, RMS normalization typically includes a learned scale parameter, whose shape (and thus dimensionality) can be configured, and which in case of the RMS norm varies between different NN suites. This can also be optionally overridden with an externally generated value. If packet data is given, this node ensures that the instance axes come first and the feature axes come last. More Info... Version 0.2.0

Ports/Properties

data

Data to process.

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

scale_init

Initializer for the trainable scale.

  • verbose name: Scale Init
  • default value: None
  • port type: DataPort
  • value type: BaseNode (can be None)
  • data direction: IN

scale_prior

Optional prior distribution for the scale.

  • verbose name: Scale Prior
  • default value: None
  • port type: DataPort
  • value type: Distribution (can be None)
  • data direction: IN

axes

"Optional comma-separated list of axis names or indices over which to accumulate the normalization statistics. If unspecified, the statistics will be accumulated over all except the instance axes. This parameter is not limited to the predefined choices.

  • verbose name: Accumulate Across Axes
  • default value: (non-instance)
  • port type: ComboPort
  • value type: str (can be None)

param_axes

List of axis names/indices across which to learn separate per-element scale and bias parameters. Like with the layer norm, different NN libraries use different conventions for this parameter. Haiku and Sonnet use the feature axis (or last axis) by default, meaning that each feature is post-scaled independently as in the batch norm, but some other ML may set this to the same as the axes parameter, which causes a separate scale/bias to be learned also across all entries of the spatial axes. Like axis, this parameter is not limited to the predefined choices.

  • verbose name: Learn Scale/bias Across Axes
  • default value: feature
  • port type: ComboPort
  • value type: str (can be None)

epsilon

Small value to add to the variance to avoid division by zero.

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

learnable_scale

Whether to learn a trainable scale parameter. Normalizations typically include such a parameter in order to drive the subsequent activation function in a regime that is desirable for downstream computations (e.g., saturating or linear). Note the shape (and thus dimensionality) of the learned parameter is governed by the param_axes (learn scale/bias across axes) parameter.

  • verbose name: Learnable Scale
  • default value: True
  • port type: BoolPort
  • value type: bool (can be None)

scale_initializer

Choice of scale initializer. This can either be one of the provided initializers, or the value "custom", in which case one of the Initializer nodes must be wired into the respective input port. For beginners it is recommended to stick to the defaults, since initialization of deep net layers is nuanced and can be tricky, otherwise be prepared to experiment with different choices. In general, the variance-scaling (lecun, glorot/xavier, he/kaiming) initializers are recommended, except for very simple/small layers where you may have a good default assumption as to the distribution of the weights (e.g., truncated_normal or uniform). Bias layers are typically zero-initialized. For initializers that take arguments, you can also type out the arguments positionally as in "truncated_normal(1.0,0.0)" (note reversed order of stddev, mean). The following initializers have arguments (here listed with their defaults): constant(value), those ending in normal(stddev=1, mean=0), those ending in uniform(min=0,max=1), orthogonal(scale=1,axis=-1), identity(gain=1), and variance_scaling(scale=1, "fan_in" (default)/"fan_avg"/"fan_out", "truncated_normal"(default)/"normal"/"uniform",optional-axis-indices=auto). Note that glorot and xavier are aliases for each other, and likewise he and kaiming are aliases for each other.

  • verbose name: Scale Initializer
  • default value: ones
  • port type: ComboPort
  • value type: str (can be None)

layername

Name of the layer. Used for naming of the trainable parameters.

  • verbose name: Layer Name
  • default value: rmsnorm
  • port type: StringPort
  • 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)