Module Kaun.Rms_norm

Root mean square normalization over the feature axis.

Root mean square normalization (Zhang and Sennrich, 2019).

RMS norm divides each feature vector — the last axis of the input — by its root mean square, then applies a learned per-feature scale (gamma). It is layer normalization without the centering and without a shift: one reduction instead of two, and the normalization Llama-class models use. Construct parameters with init or make and normalize with apply.

Types

#type ('a, _) t = {

#gamma : 'a;

}

The type for RMS-norm parameters over payload 'a. At tensor payloads gamma (the scale) has shape [| dim |], one entry per feature.

Constructors

val make : 
  dim:int ->
  (float, 'b) Nx.dtype ->
  'd Nx.Placement.t ->
  ((float, 'b, 'd) Nx.t, 'd) t

make ~dim dtype at is a fresh normalization over dim features at at, with gamma all ones.

Raises Invalid_argument if dim is not positive.

#val init : dim:int -> 'd Nx.Placement.t -> ('d Nx.float32_t, 'd) t

init ~dim at is make ~dim Nx.float32 at.

Applying

val apply : 
  ?eps:float ->
  ((float, 'b, 'd) Nx.t, 'd) t ->
  (float, 'b, 'd) Nx.t ->
  (float, 'b, 'd) Nx.t

apply p x scales each vector along x's last axis to unit root mean square and rescales it:

x / sqrt (mean (x * x) + eps) * gamma

where mean is taken along the last axis; every other axis is a batch axis. eps keeps the division finite for a zero vector and defaults to 1e-6. The result has x's shape. Differentiable through Rune.

For narrow inputs (float16, bfloat16, float8, float4) the mean square and the division run in a float32 island and the normalized values are cast back to x's dtype before the gamma scale, as in Layer_norm.apply.

Raises Invalid_argument if x is a scalar, if x's last axis does not have size dim, or if eps is negative.

Structure

#val walk : ('a, 'b, 'd, 'e) Nx.Ptree.Walk.cursor -> ('a, 'd) t -> ('b, 'e) t

walk c p walks p's parameters, gamma, at that path: the layer's Nx.Ptree.S instance. Nx.Ptree.instantiate (module Rms_norm) is the layer at one dtype, and Nx.Ptree.cast (module Rms_norm) dtype p converts its precision.