Let parameters declare their own initializer

pytensor_ml/layers/linear.py

@@ -5,6 +5,7 @@
 
 from pytensor_ml.base import Layer, UnaryLayerOp
 from pytensor_ml.params import trainable
+from pytensor_ml.state import ZeroInitializer
 
 
 def shape_to_str(shape):
@@ -23,6 +24,29 @@ def build_inner_graph(self, X, W, b=None):
 
 
 class Linear(Layer):
+    r"""
+    Affine map :math:`y = x W + b`.
+
+    Parameters
+    ----------
+    name : str or None
+        Name prefix for the layer's parameters. Defaults to "Linear" when None.
+    n_in : int
+        Size of the input feature axis.
+    n_out : int
+        Size of the output feature axis.
+    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.
+
+    Notes
+    -----
+    The weight matrix :math:`W` starts at zero, so in a stack every activation below the first layer is
+    zero and every weight matrix receives a zero gradient: an uninitialized network can fit only its
+    output bias, and predicts a constant. Call :meth:`~pytensor_ml.model.Model.initialize`, or assign a
+    value yourself, before training.
+    """
+
     def __init__(self, name: str | None, n_in: int, n_out: int, bias: bool = True):
         self.name = name if name else "Linear"
         self.n_in = n_in
@@ -34,7 +58,7 @@ def __init__(self, name: str | None, n_in: int, n_out: int, bias: bool = True):
 
         if self.bias:
             b_value = np.zeros(n_out, dtype=config.floatX)
-            self.b = trainable(b_value, f"{self.name}_b")
+            self.b = trainable(b_value, f"{self.name}_b", initializer=ZeroInitializer())
 
     def __call__(self, X: pt.TensorLike) -> pt.TensorVariable:
         X = pt.as_tensor(X)

pytensor_ml/layers/norm.py

@@ -5,6 +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
 
 
 def _standardize(X, epsilon, axis, keepdims=False):
@@ -77,9 +78,17 @@ def _resolve_n_in(name: str, n_in: int | None, X: pt.TensorVariable | None) -> i
 
 def _affine_parameters(name: str, n_in: int) -> 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."""
-    loc = trainable(np.zeros(n_in, dtype=config.floatX), f"{name}_loc")
-    scale = trainable(np.ones(n_in, dtype=config.floatX), f"{name}_scale")
+    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.
+    """
+    loc = trainable(
+        np.zeros(n_in, dtype=config.floatX), f"{name}_loc", initializer=ZeroInitializer()
+    )
+    scale = trainable(
+        np.ones(n_in, dtype=config.floatX), f"{name}_scale", initializer=OneInitializer()
+    )
     return loc, scale
 
 
@@ -151,7 +160,9 @@ class BatchNorm2D(Layer):
         Weight :math:`m` of the current batch statistic in the running-average update. Default is
         0.1.
     affine : bool, optional
-        Apply the learned scale :math:`\gamma` and shift :math:`\beta`. Default is True.
+        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.
     track_running_stats : bool, optional
         Maintain running mean and variance for use at prediction time. Default is True.
 
@@ -283,7 +294,9 @@ class LayerNorm(Layer):
     epsilon : float, optional
         Constant :math:`\epsilon` added to the variance for numerical stability. Default is 1e-5.
     affine : bool, optional
-        Apply the learned scale :math:`\gamma` and shift :math:`\beta`. Default is True.
+        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.
     """
 
     def __init__(

pytensor_ml/model.py

@@ -41,11 +41,15 @@ def initialize(
         """
         Initialize the trainable weights in place and return self.
 
+        A parameter that declares its own initializer keeps it, so ``scheme`` reaches the weights whose
+        starting value is a free choice and leaves the rest alone: a batch norm layer stays at its
+        identity transform, and biases stay at zero.
+
         Parameters
         ----------
         scheme : str or Initializer
-            Initialization scheme for the weights: the name of a built-in scheme, or an
-            :class:`~pytensor_ml.state.Initializer` instance. Default 'xavier_normal'.
+            Initialization scheme for the weights that do not declare one: the name of a built-in
+            scheme, or an :class:`~pytensor_ml.state.Initializer` instance. Default 'xavier_normal'.
         seed : int or numpy Generator, optional
             Seed for reproducible initialization.
         """

pytensor_ml/params.py

@@ -1,12 +1,28 @@
+from typing import TYPE_CHECKING
+
 import numpy as np
 
 from pytensor.tensor.sharedvar import TensorSharedVariable
 from pytensor.tensor.type import TensorType
 from pytensor.tensor.variable import TensorVariable
 
+if TYPE_CHECKING:
+    from pytensor_ml.state import Initializer
+
 
 class TrainableParameter(TensorSharedVariable):
-    """Marker class for trainable parameters (weights, biases)."""
+    """
+    Marker class for trainable parameters (weights, biases).
+
+    Attributes
+    ----------
+    initializer : Initializer or None
+        The parameter's own initializer, which wins over any scheme passed to
+        :func:`~pytensor_ml.state.initialize_params`. A layer declares one when the starting value is
+        part of its definition, such as batch norm's unit scale. None defers to the caller's scheme.
+    """
+
+    initializer: "Initializer | None" = None
 
 
 class NonTrainableParameter(TensorSharedVariable):
@@ -31,12 +47,20 @@ def _make_parameter[T: TensorSharedVariable](
     return parameter_type(name=name, type=ttype, value=value, strict=strict, **kwargs)
 
 
-def trainable(value, name=None, shape=None, strict=False, **kwargs) -> TrainableParameter:
+def trainable(
+    value,
+    name=None,
+    shape=None,
+    strict=False,
+    initializer: "Initializer | None" = None,
+    **kwargs,
+) -> TrainableParameter:
     """
     Create a shared variable marked as a trainable parameter.
 
-    The marker class is the only difference from a plain pytensor shared variable. It exists so that graph
-    traversal can tell parameters apart from other shared state; it adds no behavior of its own.
+    The marker class lets graph traversal tell parameters apart from other shared state, so that an
+    optimizer updates exactly these. A parameter may also declare its own initializer, which protects a
+    meaningful starting value from being overwritten by a network-wide initialization scheme.
 
     Parameters
     ----------
@@ -49,10 +73,16 @@ def trainable(value, name=None, shape=None, strict=False, **kwargs) -> Trainable
         entries for dynamic dimensions, e.g. ``(None, None)`` for a fully dynamic matrix.
     strict : bool, optional
         If True, the value must exactly match the dtype.
+    initializer : Initializer, optional
+        Initializer that reinitializing this parameter must use, whatever scheme the caller asks for.
+        Declare one when the starting value belongs to the layer's definition, such as a unit scale or a
+        zero bias. Default None, which defers to the caller's scheme.
     **kwargs
         Additional arguments passed to the SharedVariable constructor.
     """
-    return _make_parameter(TrainableParameter, value, name, shape, strict, **kwargs)
+    parameter = _make_parameter(TrainableParameter, value, name, shape, strict, **kwargs)
+    parameter.initializer = initializer
+    return parameter
 
 
 def non_trainable(value, name=None, shape=None, strict=False, **kwargs) -> NonTrainableParameter:

pytensor_ml/state.py

@@ -6,11 +6,12 @@
 
 from pytensor.compile.sharedvalue import SharedVariable
 
+from pytensor_ml.params import TrainableParameter
 from pytensor_ml.pytensorf import RandomSeed
 
 RandomState = RandomSeed | np.random.RandomState | np.random.Generator
 
-InitializationScheme = Literal["zeros", "xavier_uniform", "xavier_normal", "unit_uniform"]
+InitializationScheme = Literal["zeros", "ones", "xavier_uniform", "xavier_normal", "unit_uniform"]
 
 SamplingFunction = Callable[[tuple[int, ...], str, np.random.Generator], np.ndarray]
 
@@ -44,6 +45,11 @@ def sample(self, shape: tuple[int, ...], dtype: str, rng: np.random.Generator) -
         return np.zeros(shape, dtype=dtype)
 
 
+class OneInitializer(Initializer):
+    def sample(self, shape: tuple[int, ...], dtype: str, rng: np.random.Generator) -> np.ndarray:
+        return np.ones(shape, dtype=dtype)
+
+
 class UnitUniformInitializer(Initializer):
     def sample(self, shape: tuple[int, ...], dtype: str, rng: np.random.Generator) -> np.ndarray:
         return rng.uniform(0.0, 1.0, size=shape).astype(dtype)
@@ -80,6 +86,7 @@ def sample(self, shape: tuple[int, ...], dtype: str, rng: np.random.Generator) -
 
 _INITIALIZERS: dict[str, type[Initializer]] = {
     "zeros": ZeroInitializer,
+    "ones": OneInitializer,
     "xavier_uniform": XavierUniformInitializer,
     "xavier_normal": XavierNormalInitializer,
     "unit_uniform": UnitUniformInitializer,
@@ -88,6 +95,12 @@ def sample(self, shape: tuple[int, ...], dtype: str, rng: np.random.Generator) -
 InitializationSchemeLike = InitializationScheme | Initializer
 
 
+def _declared_initializer(param: SharedVariable, default: Initializer) -> Initializer:
+    """The parameter's own initializer, or ``default`` when it does not declare one."""
+    declared = param.initializer if isinstance(param, TrainableParameter) else None
+    return default if declared is None else declared
+
+
 def initialize_params(
     params: Sequence[SharedVariable],
     scheme: InitializationSchemeLike = "xavier_normal",
@@ -96,13 +109,19 @@ def initialize_params(
     """
     Initialize parameter values using the specified scheme.
 
+    A :class:`~pytensor_ml.params.TrainableParameter` that declares its own ``initializer`` uses it
+    instead of ``scheme``, leaving batch norm at its unit scale while the weight matrices around it are
+    drawn from the requested scheme. Call an :class:`Initializer` on a parameter directly to overwrite a
+    declared value anyway.
+
     Parameters
     ----------
     params
         SharedVariables to initialize values for.
     scheme
-        Initialization scheme to use: the name of a built-in scheme, or any :class:`Initializer`
-        instance (including a :class:`CustomInitializer` wrapping your own sampling function).
+        Initialization scheme for parameters that do not declare one: the name of a built-in scheme, or
+        any :class:`Initializer` instance (including a :class:`CustomInitializer` wrapping your own
+        sampling function).
     rng
         Random number generator. If None, a new one is created.
 
@@ -115,4 +134,7 @@ def initialize_params(
     rng = np.random.default_rng(rng)
 
     initializer = scheme if isinstance(scheme, Initializer) else _INITIALIZERS[scheme]()
-    return [initializer._sample_like(param, rng) for param in params]
+    return [
+        _declared_initializer(param, default=initializer)._sample_like(param, rng)
+        for param in params
+    ]

tests/test_model.py

@@ -1,11 +1,13 @@
 import numpy as np
 import pytensor.tensor as pt
+import pytest
 
 from pytensor import config
 
 import pytensor_ml.model
 
-from pytensor_ml.layers import BatchNorm2D, Linear, Sequential
+from pytensor_ml.activations import ReLU
+from pytensor_ml.layers import BatchNorm2D, LayerNorm, Linear, Sequential
 from pytensor_ml.loss import SquaredError
 from pytensor_ml.model import Model
 from pytensor_ml.optim import sgd
@@ -71,6 +73,53 @@ def counting_compile_predict(*args, **kwargs):
         np.testing.assert_array_equal(first, second)
 
 
+class TestModelInitialize:
+    @pytest.mark.parametrize(
+        "norm_layer", [BatchNorm2D, LayerNorm], ids=["batch_norm", "layer_norm"]
+    )
+    def test_leaves_a_norm_layer_at_the_identity_transform(self, norm_layer):
+        X = pt.tensor("X", shape=(None, 8))
+        norm = norm_layer("norm", n_in=4)
+        y = Sequential(Linear("fc1", 8, 4), norm, ReLU(), Linear("fc2", 4, 2))(X)
+
+        Model(X, y).initialize("xavier_normal", seed=0)
+
+        np.testing.assert_array_equal(norm.scale.get_value(), 1)
+        np.testing.assert_array_equal(norm.loc.get_value(), 0)
+
+    def test_draws_weight_matrices_and_leaves_biases_at_zero(self):
+        X = pt.tensor("X", shape=(None, 8))
+        fc1 = Linear("fc1", 8, 4)
+        y = Sequential(fc1, ReLU(), Linear("fc2", 4, 2))(X)
+
+        Model(X, y).initialize("xavier_normal", seed=0)
+
+        assert np.abs(fc1.W.get_value()).min() > 0
+        np.testing.assert_array_equal(fc1.b.get_value(), 0)
+
+    def test_a_declared_initializer_does_not_freeze_the_parameter(self):
+        """Declaring an initializer protects a starting value, not the parameter. Excluding declared
+        parameters from training instead would leave batch norm's scale pinned at one and still satisfy
+        every assertion above."""
+        rng = np.random.default_rng(0)
+        X = pt.tensor("X", shape=(None, 8))
+        norm = BatchNorm2D("norm", n_in=4)
+        y = Sequential(Linear("fc1", 8, 4), norm, ReLU(), Linear("fc2", 4, 2))(X)
+        model = Model(X, y).initialize("xavier_normal", seed=0)
+
+        target = pt.matrix("target")
+        step = model.compile_train(
+            sgd(learning_rate=0.1), loss=((y - target) ** 2).mean(), inputs=[X, target]
+        )
+        step(
+            rng.normal(size=(8, 8)).astype(config.floatX),
+            rng.normal(size=(8, 2)).astype(config.floatX),
+        )
+
+        assert not np.array_equal(norm.scale.get_value(), np.ones(4))
+        assert not np.array_equal(norm.loc.get_value(), np.zeros(4))
+
+
 def test_compile_train_accepts_a_prebuilt_loss():
     # An autoencoder reconstructs its own input, so there is no target separate from X and the supervised
     # path cannot express it. The step takes one argument, not two.