diff --git a/pytensor_ml/layers/linear.py b/pytensor_ml/layers/linear.py index ab1b766..029aaa8 100644 --- a/pytensor_ml/layers/linear.py +++ b/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) diff --git a/pytensor_ml/layers/norm.py b/pytensor_ml/layers/norm.py index ed264ca..a3ed750 100644 --- a/pytensor_ml/layers/norm.py +++ b/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__( diff --git a/pytensor_ml/model.py b/pytensor_ml/model.py index 65b88f6..e110d11 100644 --- a/pytensor_ml/model.py +++ b/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. """ diff --git a/pytensor_ml/params.py b/pytensor_ml/params.py index 8e4bfe2..51efb13 100644 --- a/pytensor_ml/params.py +++ b/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: diff --git a/pytensor_ml/state.py b/pytensor_ml/state.py index 42e9c0d..00e4958 100644 --- a/pytensor_ml/state.py +++ b/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 + ] diff --git a/tests/test_model.py b/tests/test_model.py index 1d922a0..d480178 100644 --- a/tests/test_model.py +++ b/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. diff --git a/tests/test_state.py b/tests/test_state.py index 9fbc370..80766f3 100644 --- a/tests/test_state.py +++ b/tests/test_state.py @@ -1,6 +1,7 @@ from typing import get_args import numpy as np +import pytensor import pytest from pytensor_ml.params import trainable @@ -9,6 +10,7 @@ _INITIALIZERS, CustomInitializer, InitializationScheme, + OneInitializer, ZeroInitializer, initialize_params, ) @@ -72,9 +74,11 @@ def test_parameters_do_not_all_receive_the_same_draws(self): assert not np.array_equal(values[0], values[1]) - def test_accepts_a_custom_initializer(self, simple_network): - X, y = simple_network - params = collect_trainable_params(y) + def test_accepts_a_custom_initializer(self): + params = [ + trainable(np.zeros((4, 4), dtype="float64"), "first"), + trainable(np.zeros((4, 2), dtype="float64"), "second"), + ] constant = CustomInitializer(lambda shape, dtype, rng: np.full(shape, 7.0, dtype=dtype)) values = initialize_params(params, scheme=constant) @@ -84,6 +88,39 @@ def test_accepts_a_custom_initializer(self, simple_network): np.testing.assert_array_equal(val, 7.0) +class TestDeclaredInitializers: + def test_a_declared_initializer_wins_over_the_scheme(self): + # Starting from zeros, so the assertion only holds if the declared initializer actually ran. + scale = trainable(np.zeros(4, dtype="float64"), "scale", initializer=OneInitializer()) + + [value] = initialize_params([scale], scheme="xavier_normal", rng=0) + + np.testing.assert_array_equal(value, 1) + + def test_undeclared_parameters_still_follow_the_scheme(self): + weight = trainable(np.zeros((4, 4), dtype="float64"), "weight") + bias = trainable(np.zeros(4, dtype="float64"), "bias", initializer=ZeroInitializer()) + + weight_value, bias_value = initialize_params([weight, bias], scheme="unit_uniform", rng=0) + + assert weight_value.min() > 0 + np.testing.assert_array_equal(bias_value, 0) + + def test_shared_variables_without_the_marker_class_follow_the_scheme(self): + state = pytensor.shared(np.zeros(4, dtype="float64"), name="state") + + [value] = initialize_params([state], scheme="unit_uniform", rng=0) + + assert value.min() > 0 + + def test_calling_an_initializer_overrides_a_declaration(self): + scale = trainable(np.ones(3, dtype="float64"), "scale", initializer=OneInitializer()) + + ZeroInitializer()(scale) + + np.testing.assert_array_equal(scale.get_value(), 0) + + def test_calling_an_initializer_assigns_the_parameter_in_place(): param = trainable(np.ones(3, dtype="float64"), "w")