ScaledWSLinear

ScaledWSLinear#

class braintrace.nn.ScaledWSLinear#

Linear layer with weight standardization.

Applies weight standardization [1] to normalize weights before the linear transformation, which can improve training stability and performance.

Parameters:

  • in_size (int or tuple of int) – The input feature size.

  • out_size (int or tuple of int) – The output feature size.

  • w_init (Callable, optional) – Weight initializer. Default is KaimingNormal().

  • b_init (Callable, optional) – Bias initializer. Default is ZeroInit().

  • w_mask (ArrayLike, Callable, or None, optional) – Optional mask for the weights.

  • ws_gain (bool, optional) – Whether to use a learnable gain parameter for weight standardization. Default is True.

  • eps (float, optional) – Small constant for numerical stability in standardization. Default is 1e-4.

  • name (str, optional) – Name of the module.

  • param_type (type, optional) – Type of parameter state. Default is ParamState.

Variables:

  • in_size (tuple) – Input feature size.

  • out_size (tuple) – Output feature size.

  • w_mask (ArrayLike or None) – Weight mask if provided.

  • eps (float) – Epsilon for numerical stability.

  • weight (ParamState) – Parameter state containing ‘weight’, optionally ‘bias’ and ‘gain’.

References

Examples

>>> import braintrace as braintrace
>>> import jax.numpy as jnp
>>>
>>> # Create a weight-standardized linear layer
>>> layer = braintrace.nn.ScaledWSLinear((10,), (5,))
>>> x = jnp.ones((32, 10))
>>> y = layer(x)
>>> y.shape
(32, 5)
>>>
>>> # Without learnable gain
>>> layer = braintrace.nn.ScaledWSLinear((10,), (5,), ws_gain=False)
>>> y = layer(x)
>>> y.shape
(32, 5)
update(x)#

Apply the weight-standardized linear transform through ETP matmul.

Weight standardization (and the optional mask) are applied inside weight_fn, which closes over the current eps value only. Routing the transform through braintrace.matmul() with weight_fn= causes the ETP compiler to track the gradient w.r.t. the raw weight leaf exactly (the standardization Jacobian is recovered via jax.vjp).

Note on post-ops: gain and bias are applied OUTSIDE the matmul primitive as post-operations, so the eligibility trace tracks only the standardized weight leaf. gain and bias are therefore non-temporal for the eligibility trace — in genuine online training their trace-based gradient is partial. (The multi-step VJP oracle path used in tests, which autodiffs through the full rollout, still recovers them exactly; only the online eligibility-trace gradient is non-temporal.)

Parameters:

x (ArrayLike) – Input array, of shape (..., in_size).

Returns:

ArrayLike – The transformed output, of shape (..., out_size).

ScaledWSLinear.__init__(in_size, out_size, w_init=KaimingNormal(   scale=2.0, mode='fan_in', in_axis=-2, out_axis=-1, distribution='truncated_normal', rng=RandomState(Array((), dtype=key<fry>) overlaying:   [2233333421 2029709265]), unit=Unit("1") ), b_init=ZeroInit(   unit=Unit("1") ), w_mask=None, ws_gain=True, eps=0.0001, name=None, param_type=<class 'brainstate.ParamState'>)#