xai-org/x-algorithm · error · ValueError

RMSNorm: when weight_decay_mask > 0 the layer is re-paramete

Error message

RMSNorm: when weight_decay_mask > 0 the layer is re-parameterized as (1 + scale) * x and `scale_init` must be jnp.zeros (got {scale_init!r}). Pass scale_init=jnp.zeros or leave it as None to use the default.

What it means

When weight_decay_mask > 0, RMSNorm re-parameterizes the scale as (1 + scale) * x, which requires the scale to start at zero; hence scale_init must be exactly jnp.zeros (or None to get that default). Passing any other initializer (e.g. jnp.ones) breaks the re-parameterization invariant and is rejected.

Source

Thrown at phoenix/xrex/models/normalization.py:87

        weight_decay_mask: float = 0.0,
    ):
        super().__init__(name=name)
        if isinstance(axis, slice):
            self.axis = axis
        elif isinstance(axis, int):
            self.axis = (axis,)
        elif isinstance(axis, abc.Iterable) and all(isinstance(ax, int) for ax in axis):
            self.axis = tuple(axis)
        else:
            raise ValueError("`axis` should be an int, slice or iterable of ints.")

        self.eps = eps
        self.create_scale = create_scale
        self.reparameterize = weight_decay_mask > 0
        if scale_init is None:
            scale_init = jnp.zeros if self.reparameterize else jnp.ones
        if self.reparameterize and scale_init is not jnp.zeros:
            raise ValueError(
                "RMSNorm: when weight_decay_mask > 0 the layer is "
                "re-parameterized as (1 + scale) * x and `scale_init` must be "
                f"jnp.zeros (got {scale_init!r}). Pass scale_init=jnp.zeros or "
                "leave it as None to use the default."
            )
        self.scale_init = scale_init
        self.pspec = pspec
        self.lr_multiplier = lr_multiplier
        self.weight_decay_mask = weight_decay_mask

    def __call__(self, inputs: jax.Array):
        fprop_dtype = inputs.dtype
        param_shape = (inputs.shape[-1],)
        if self.create_scale:
            scale = get_parameter(
                "scale",
                param_shape,
                dtype=jnp.float32,

View on GitHub (pinned to 24c60942c5)

Solutions

Drop the scale_init argument (let it default to jnp.zeros when reparameterizing).
Or explicitly pass scale_init=jnp.zeros.
If you need nonzero init scale, set weight_decay_mask=0 for that layer.

Example fix

# before
RMSNorm(..., weight_decay_mask=1, scale_init=jnp.ones)

# after
RMSNorm(..., weight_decay_mask=1)  # scale_init defaults to jnp.zeros

Defensive patterns

Strategy: validation

Validate before calling

if weight_decay_mask > 0:
    scale_init = None  # force default jnp.zeros reparameterization

Prevention

Never pass scale_init when reparameterization is on; let the default apply.
Encapsulate norm construction in one builder that enforces the constraint.

When it happens

Trigger: Constructing the norm with weight_decay_mask > 0 while explicitly passing scale_init=jnp.ones or a custom initializer; copying constructor args from a non-reparameterized layer.

Common situations: Enabling weight decay on norms mid-project and reusing old constructor kwargs; defaults changed in shared builder code.

Related errors

AI-assisted analysis of xai-org/x-algorithm@24c60942c5 (2026-08-28). Data as JSON: /api/errors/1c304903ab464bed. Report an issue: GitHub.