Let every layer take its initializers

pytensor_ml/layers/embedding.py

@@ -5,6 +5,7 @@
 
 from pytensor_ml.base import Layer, UnaryLayerOp
 from pytensor_ml.params import trainable
+from pytensor_ml.state import Initializer
 
 
 class EmbeddingLayer(UnaryLayerOp):
@@ -30,15 +31,27 @@ class Embedding(Layer):
         Number of rows in the table -- the number of distinct indices it can map.
     n_features : int
         Size of each embedding row.
+    weight_initializer : Initializer, optional
+        How the table is drawn, in place of whatever scheme :meth:`~pytensor_ml.model.Model.initialize` is
+        given. Left to the scheme when omitted. A fan-scaled scheme puts the vocabulary size in the
+        denominator, which is correct Xavier and much tighter than the ``NormalInitializer(0.0, 0.02)`` that
+        reference implementations of GPT-2 use, so this is the keyword to reach for when matching one.
     """
 
-    def __init__(self, name: str | None, n_embeddings: int, n_features: int):
+    def __init__(
+        self,
+        name: str | None,
+        n_embeddings: int,
+        n_features: int,
+        *,
+        weight_initializer: Initializer | None = None,
+    ):
         self.name = name if name else "Embedding"
         self.n_embeddings = n_embeddings
         self.n_features = n_features
 
         W_value = np.zeros((n_embeddings, n_features), dtype=config.floatX)
-        self.W = trainable(W_value, f"{self.name}_W")
+        self.W = trainable(W_value, f"{self.name}_W", initializer=weight_initializer)
 
     def __call__(self, ids: pt.TensorLike) -> pt.TensorVariable:
         ids = pt.as_tensor(ids)

pytensor_ml/layers/linear.py

@@ -5,7 +5,7 @@
 
 from pytensor_ml.base import Layer, UnaryLayerOp
 from pytensor_ml.params import trainable
-from pytensor_ml.state import ZeroInitializer
+from pytensor_ml.state import Initializer, ZeroInitializer
 
 
 def shape_to_str(shape):
@@ -38,6 +38,12 @@ class Linear(Layer):
     bias : bool, optional
         Add the learned shift :math:`b`, which starts at zero and stays there under a network-wide
         initialization scheme. Default is True.
+    weight_initializer : Initializer, optional
+        How :math:`W` is drawn, in place of whatever scheme :meth:`~pytensor_ml.model.Model.initialize` is
+        given. Left to the scheme when omitted, which is what makes a network-wide choice reach the weights.
+    bias_initializer : Initializer, optional
+        How :math:`b` is drawn. Zeros when omitted, following Keras and flax rather than torch, which draws
+        the bias from :math:`\mathcal{U}(\pm 1/\sqrt{\text{fan\_in}})`.
 
     Notes
     -----
@@ -47,18 +53,31 @@ class Linear(Layer):
     value yourself, before training.
     """
 
-    def __init__(self, name: str | None, n_in: int, n_out: int, bias: bool = True):
+    def __init__(
+        self,
+        name: str | None,
+        n_in: int,
+        n_out: int,
+        bias: bool = True,
+        *,
+        weight_initializer: Initializer | None = None,
+        bias_initializer: Initializer | None = None,
+    ):
         self.name = name if name else "Linear"
         self.n_in = n_in
         self.n_out = n_out
         self.bias = bias
 
         W_value = np.zeros((n_in, n_out), dtype=config.floatX)
-        self.W = trainable(W_value, f"{self.name}_W")
+        self.W = trainable(W_value, f"{self.name}_W", initializer=weight_initializer)
 
         if self.bias:
             b_value = np.zeros(n_out, dtype=config.floatX)
-            self.b = trainable(b_value, f"{self.name}_b", initializer=ZeroInitializer())
+            self.b = trainable(
+                b_value,
+                f"{self.name}_b",
+                initializer=ZeroInitializer() if bias_initializer is None else bias_initializer,
+            )
 
     def __call__(self, X: pt.TensorLike) -> pt.TensorVariable:
         X = pt.as_tensor(X)

pytensor_ml/layers/norm.py

@@ -5,7 +5,7 @@
 
 from pytensor_ml.base import Layer, LayerOp, UnaryLayerOp
 from pytensor_ml.params import NonTrainableParameter, TrainableParameter, non_trainable, trainable
-from pytensor_ml.state import OneInitializer, ZeroInitializer
+from pytensor_ml.state import Initializer, OneInitializer, ZeroInitializer
 
 
 def _standardize(X, epsilon, axis, keepdims=False):
@@ -76,18 +76,28 @@ def _resolve_n_in(name: str, n_in: int | None, X: pt.TensorVariable | None) -> i
     return inferred
 
 
-def _affine_parameters(name: str, n_in: int) -> tuple[TrainableParameter, TrainableParameter]:
+def _affine_parameters(
+    name: str,
+    n_in: int,
+    loc_initializer: Initializer | None = None,
+    scale_initializer: Initializer | None = None,
+) -> tuple[TrainableParameter, TrainableParameter]:
     """Build the learned shift and scale. Returns them in the ``(loc, scale)`` order that every norm
     op unpacks its inputs in, so the two cannot drift apart.
 
     Both declare their initializer, so a network-wide scheme leaves the identity transform in place --
-    normalizing and then rescaling by a random factor defeats the point of the layer.
+    normalizing and then rescaling by a random factor defeats the point of the layer. A caller who wants
+    something else says so, and their choice becomes the declaration.
     """
     loc = trainable(
-        np.zeros(n_in, dtype=config.floatX), f"{name}_loc", initializer=ZeroInitializer()
+        np.zeros(n_in, dtype=config.floatX),
+        f"{name}_loc",
+        initializer=ZeroInitializer() if loc_initializer is None else loc_initializer,
     )
     scale = trainable(
-        np.ones(n_in, dtype=config.floatX), f"{name}_scale", initializer=OneInitializer()
+        np.ones(n_in, dtype=config.floatX),
+        f"{name}_scale",
+        initializer=OneInitializer() if scale_initializer is None else scale_initializer,
     )
     return loc, scale
 
@@ -163,6 +173,11 @@ class BatchNorm2D(Layer):
         Apply the learned scale :math:`\gamma` and shift :math:`\beta`, starting from the identity
         transform :math:`\gamma = 1`, :math:`\beta = 0`, which a network-wide initialization scheme
         leaves in place. Default is True.
+    scale_initializer : Initializer, optional
+        How :math:`\gamma` is drawn. Ones when omitted, which is the identity transform; a scheme that
+        rescaled a normalized activation by a random factor would defeat the layer.
+    loc_initializer : Initializer, optional
+        How :math:`\beta` is drawn. Zeros when omitted.
     track_running_stats : bool, optional
         Maintain running mean and variance for use at prediction time. Default is True.
 
@@ -186,13 +201,18 @@ def __init__(
         momentum: float = 0.1,
         affine: bool = True,
         track_running_stats: bool = True,
+        *,
+        scale_initializer: Initializer | None = None,
+        loc_initializer: Initializer | None = None,
     ):
         self.name = name if name else "BatchNorm"
         self.n_in = n_in
         self.epsilon = epsilon
         self.momentum = momentum
         self.affine = affine
         self.track_running_stats = track_running_stats
+        self._scale_initializer = scale_initializer
+        self._loc_initializer = loc_initializer
 
         self.scale: TrainableParameter | None = None
         self.loc: TrainableParameter | None = None
@@ -213,7 +233,12 @@ def _initialize_params(self, X: pt.TensorVariable | None):
             return
 
         if self.affine:
-            self.loc, self.scale = _affine_parameters(self.name, n_in)
+            self.loc, self.scale = _affine_parameters(
+                self.name,
+                n_in,
+                loc_initializer=self._loc_initializer,
+                scale_initializer=self._scale_initializer,
+            )
 
         if self.track_running_stats:
             zeros = np.zeros(n_in, dtype=config.floatX)
@@ -297,6 +322,11 @@ class LayerNorm(Layer):
         Apply the learned scale :math:`\gamma` and shift :math:`\beta`, starting from the identity
         transform :math:`\gamma = 1`, :math:`\beta = 0`, which a network-wide initialization scheme
         leaves in place. Default is True.
+    scale_initializer : Initializer, optional
+        How :math:`\gamma` is drawn. Ones when omitted, which is the identity transform; a scheme that
+        rescaled a normalized activation by a random factor would defeat the layer.
+    loc_initializer : Initializer, optional
+        How :math:`\beta` is drawn. Zeros when omitted.
     """
 
     def __init__(
@@ -305,11 +335,16 @@ def __init__(
         n_in: int | None = None,
         epsilon: float = 1e-5,
         affine: bool = True,
+        *,
+        scale_initializer: Initializer | None = None,
+        loc_initializer: Initializer | None = None,
     ):
         self.name = name if name else "LayerNorm"
         self.n_in = n_in
         self.epsilon = epsilon
         self.affine = affine
+        self._scale_initializer = scale_initializer
+        self._loc_initializer = loc_initializer
 
         self.scale: TrainableParameter | None = None
         self.loc: TrainableParameter | None = None
@@ -326,7 +361,12 @@ def _initialize_params(self, X: pt.TensorVariable | None):
             return
 
         if self.affine:
-            self.loc, self.scale = _affine_parameters(self.name, n_in)
+            self.loc, self.scale = _affine_parameters(
+                self.name,
+                n_in,
+                loc_initializer=self._loc_initializer,
+                scale_initializer=self._scale_initializer,
+            )
 
         self.initialized = True
 

pytensor_ml/state.py

@@ -11,7 +11,9 @@
 
 RandomState = RandomSeed | np.random.RandomState | np.random.Generator
 
-InitializationScheme = Literal["zeros", "ones", "xavier_uniform", "xavier_normal", "unit_uniform"]
+InitializationScheme = Literal[
+    "zeros", "ones", "xavier_uniform", "xavier_normal", "unit_uniform", "normal"
+]
 
 SamplingFunction = Callable[[tuple[int, ...], str, np.random.Generator], np.ndarray]
 
@@ -55,6 +57,30 @@ def sample(self, shape: tuple[int, ...], dtype: str, rng: np.random.Generator) -
         return rng.uniform(0.0, 1.0, size=shape).astype(dtype)
 
 
+class NormalInitializer(Initializer):
+    r"""
+    Draw every element from :math:`\mathcal{N}(\mu, \sigma^2)`, independent of the parameter's shape.
+
+    The fan-scaled initializers derive their spread from the shape; this one is told it, which is what a
+    reference implementation quoting a specific standard deviation needs -- GPT-2 initializes its embeddings
+    and weights from ``NormalInitializer(0.0, 0.02)`` whatever their fans work out to.
+
+    Parameters
+    ----------
+    mean : float
+        Center of the distribution :math:`\mu`. Default 0.0.
+    std : float
+        Standard deviation :math:`\sigma`. Default 0.01.
+    """
+
+    def __init__(self, mean: float = 0.0, std: float = 0.01):
+        self.mean = mean
+        self.std = std
+
+    def sample(self, shape: tuple[int, ...], dtype: str, rng: np.random.Generator) -> np.ndarray:
+        return rng.normal(self.mean, self.std, size=shape).astype(dtype)
+
+
 def fans(shape: tuple[int, ...]) -> tuple[int, int]:
     r"""
     Return the number of units feeding into and out of one position of a parameter of ``shape``.
@@ -93,13 +119,36 @@ def fans(shape: tuple[int, ...]) -> tuple[int, int]:
 
 
 class XavierUniformInitializer(Initializer):
+    r"""
+    Draw from :math:`\mathcal{U}(\pm\sqrt{6 / (\text{fan\_in} + \text{fan\_out})})`.
+
+    The bound is chosen so the variance of the activations, and of the gradients flowing back, stays roughly
+    constant through depth. Also called Glorot initialization.
+
+    References
+    ----------
+    .. [1] Glorot, X. and Bengio, Y. (2010). Understanding the difficulty of training deep feedforward
+           neural networks. Proceedings of AISTATS, 249-256.
+    """
+
     def sample(self, shape: tuple[int, ...], dtype: str, rng: np.random.Generator) -> np.ndarray:
         fan_in, fan_out = fans(shape)
         scale = np.sqrt(6.0 / (fan_in + fan_out))
         return rng.uniform(-scale, scale, size=shape).astype(dtype)
 
 
 class XavierNormalInitializer(Initializer):
+    r"""
+    Draw from :math:`\mathcal{N}(0, 2 / (\text{fan\_in} + \text{fan\_out}))`.
+
+    The normal counterpart of :class:`XavierUniformInitializer`, targeting the same variance.
+
+    References
+    ----------
+    .. [1] Glorot, X. and Bengio, Y. (2010). Understanding the difficulty of training deep feedforward
+           neural networks. Proceedings of AISTATS, 249-256.
+    """
+
     def sample(self, shape: tuple[int, ...], dtype: str, rng: np.random.Generator) -> np.ndarray:
         fan_in, fan_out = fans(shape)
         scale = np.sqrt(2.0 / (fan_in + fan_out))
@@ -129,6 +178,8 @@ def sample(self, shape: tuple[int, ...], dtype: str, rng: np.random.Generator) -
     "xavier_uniform": XavierUniformInitializer,
     "xavier_normal": XavierNormalInitializer,
     "unit_uniform": UnitUniformInitializer,
+    # Reachable by name because both of its arguments have defaults; pass an instance for anything else.
+    "normal": NormalInitializer,
 }
 
 InitializationSchemeLike = InitializationScheme | Initializer