diff --git a/pytensor_ml/layers/embedding.py b/pytensor_ml/layers/embedding.py index 132cfd0..9d16064 100644 --- a/pytensor_ml/layers/embedding.py +++ b/pytensor_ml/layers/embedding.py @@ -1,11 +1,8 @@ -import numpy as np import pytensor.tensor as pt -from pytensor import config - from pytensor_ml.base import Layer, UnaryLayerOp from pytensor_ml.params import trainable -from pytensor_ml.state import Initializer +from pytensor_ml.state import Initializer, XavierNormalInitializer class EmbeddingLayer(UnaryLayerOp): @@ -50,8 +47,16 @@ def __init__( 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", initializer=weight_initializer) + # Drawn here, as in Linear, and for the same reason. The fallback is deliberately not declared, so + # a network-wide scheme still reaches the table. + W_initializer = ( + XavierNormalInitializer() if weight_initializer is None else weight_initializer + ) + self.W = trainable( + W_initializer.initial_value((n_embeddings, n_features)), + f"{self.name}_W", + initializer=weight_initializer, + ) def __call__(self, ids: pt.TensorLike) -> pt.TensorVariable: ids = pt.as_tensor(ids) diff --git a/pytensor_ml/layers/linear.py b/pytensor_ml/layers/linear.py index c2e5504..17e0eb3 100644 --- a/pytensor_ml/layers/linear.py +++ b/pytensor_ml/layers/linear.py @@ -1,11 +1,8 @@ -import numpy as np import pytensor.tensor as pt -from pytensor import config - from pytensor_ml.base import Layer, UnaryLayerOp from pytensor_ml.params import trainable -from pytensor_ml.state import Initializer, ZeroInitializer +from pytensor_ml.state import Initializer, XavierNormalInitializer, ZeroInitializer def shape_to_str(shape): @@ -47,10 +44,9 @@ class Linear(Layer): 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. + Both parameters are drawn when the layer is built, so a network trains without any further call. + :meth:`~pytensor_ml.model.Model.initialize` redraws them from a single seed, which is what makes a run + reproducible; :math:`W` follows its ``scheme`` there, since the draw here declares no initializer. """ def __init__( @@ -68,15 +64,25 @@ def __init__( 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", initializer=weight_initializer) + # Drawn here rather than left at zero: a weight and its value are one thing, and a stack of zero + # weights passes no gradient below its last layer. Note what is declared and what is not -- + # `weight_initializer` is, so it survives a network-wide scheme, while the fallback below must not + # be, or a scheme could never reach a weight matrix at all. + W_initializer = ( + XavierNormalInitializer() if weight_initializer is None else weight_initializer + ) + self.W = trainable( + W_initializer.initial_value((n_in, n_out)), + f"{self.name}_W", + initializer=weight_initializer, + ) if self.bias: - b_value = np.zeros(n_out, dtype=config.floatX) + b_initializer = ZeroInitializer() if bias_initializer is None else bias_initializer self.b = trainable( - b_value, + b_initializer.initial_value((n_out,)), f"{self.name}_b", - initializer=ZeroInitializer() if bias_initializer is None else bias_initializer, + initializer=b_initializer, ) def __call__(self, X: pt.TensorLike) -> pt.TensorVariable: diff --git a/pytensor_ml/layers/norm.py b/pytensor_ml/layers/norm.py index dd4ef75..56465d8 100644 --- a/pytensor_ml/layers/norm.py +++ b/pytensor_ml/layers/norm.py @@ -89,15 +89,11 @@ def _affine_parameters( 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() if loc_initializer is None else loc_initializer, - ) + resolved_loc = ZeroInitializer() if loc_initializer is None else loc_initializer + resolved_scale = OneInitializer() if scale_initializer is None else scale_initializer + loc = trainable(resolved_loc.initial_value((n_in,)), f"{name}_loc", initializer=resolved_loc) scale = trainable( - np.ones(n_in, dtype=config.floatX), - f"{name}_scale", - initializer=OneInitializer() if scale_initializer is None else scale_initializer, + resolved_scale.initial_value((n_in,)), f"{name}_scale", initializer=resolved_scale ) return loc, scale diff --git a/pytensor_ml/model.py b/pytensor_ml/model.py index e110d11..33c14bb 100644 --- a/pytensor_ml/model.py +++ b/pytensor_ml/model.py @@ -11,7 +11,11 @@ from pytensor_ml.loss import Loss, supervised_loss from pytensor_ml.params import TrainableParameter from pytensor_ml.pytensorf import collect_trainable_params, compile_predict -from pytensor_ml.state import InitializationSchemeLike, initialize_params +from pytensor_ml.state import ( + InitializationSchemeLike, + initialize_params, + require_varying_scheme, +) class Model: @@ -49,10 +53,13 @@ def initialize( ---------- scheme : str or Initializer 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'. + scheme, or an :class:`~pytensor_ml.state.Initializer` instance. A constant is refused here, + since identical weights leave a layer with no symmetry to break; declare one on the parameter + that wants it. Default 'xavier_normal'. seed : int or numpy Generator, optional Seed for reproducible initialization. """ + require_varying_scheme(scheme) parameters = self.weights for parameter, value in zip(parameters, initialize_params(parameters, scheme, rng=seed)): parameter.set_value(value) diff --git a/pytensor_ml/state.py b/pytensor_ml/state.py index 9eab637..5977cfc 100644 --- a/pytensor_ml/state.py +++ b/pytensor_ml/state.py @@ -4,6 +4,7 @@ import numpy as np +from pytensor import config from pytensor.compile.sharedvalue import SharedVariable from pytensor_ml.params import TrainableParameter @@ -36,6 +37,20 @@ def sample( self, shape: tuple[int, ...], dtype: str, rng: np.random.Generator ) -> np.ndarray: ... + def initial_value(self, shape: tuple[int, ...]) -> np.ndarray: + """ + Draw the value a parameter of ``shape`` is born holding, at the current ``floatX``. + + Uses fresh entropy: reproducibility comes from :meth:`~pytensor_ml.model.Model.initialize`, which + redraws every parameter from one seed and discards whatever this produced. + + Parameters + ---------- + shape : tuple of int + Shape of the parameter to draw. + """ + return self.sample(shape, config.floatX, np.random.default_rng()) + def _sample_like(self, param: SharedVariable, rng: RandomState | None = None) -> np.ndarray: rng = np.random.default_rng(rng) value = param.get_value() @@ -191,6 +206,34 @@ def _declared_initializer(param: SharedVariable, default: Initializer) -> Initia return default if declared is None else declared +def _resolve_scheme(scheme: InitializationSchemeLike) -> Initializer: + """The initializer a scheme names, or the instance itself when one was passed.""" + return scheme if isinstance(scheme, Initializer) else _INITIALIZERS[scheme]() + + +def require_varying_scheme(scheme: InitializationSchemeLike) -> None: + """ + Raise if ``scheme`` would give every weight matrix in a network the same value. + + A constant is the right initializer for a bias or a norm scale and never for a whole network: identical + weights leave no gradient to distinguish two units in a layer, so there is no symmetry for training to + break. Declare it on the parameter that wants it instead. + + Parameters + ---------- + scheme : str or Initializer + The scheme about to be applied to every parameter that declares nothing. + """ + if isinstance(_resolve_scheme(scheme), ZeroInitializer | OneInitializer): + name = scheme if isinstance(scheme, str) else type(scheme).__name__ + raise ValueError( + f"{name!r} gives every weight matrix the same value, so no gradient distinguishes two units in " + "a layer and training cannot break the symmetry. A constant belongs on one parameter rather " + "than on a network: declare it where that parameter is built, as in " + "`Linear(..., weight_initializer=ZeroInitializer())`." + ) + + def initialize_params( params: Sequence[SharedVariable], scheme: InitializationSchemeLike = "xavier_normal", @@ -223,7 +266,7 @@ def initialize_params( # Resolve once and share: a seed handed to each _sample_like call would repeat draws across parameters. rng = np.random.default_rng(rng) - initializer = scheme if isinstance(scheme, Initializer) else _INITIALIZERS[scheme]() + initializer = _resolve_scheme(scheme) return [ _declared_initializer(param, default=initializer)._sample_like(param, rng) for param in params diff --git a/tests/test_layers.py b/tests/test_layers.py index 7b2f75d..00e43b1 100644 --- a/tests/test_layers.py +++ b/tests/test_layers.py @@ -466,3 +466,69 @@ def test_a_bias_initializer_replaces_the_zero_declaration_rather_than_fighting_i values = dict(zip((p.name for p in params), initialize_params(params, scheme="zeros", rng=0))) assert not np.allclose(values["fc_b"], 0.0) + + +@pytest.mark.parametrize("layer_name", sorted(DECLARED_BY_LAYERS), ids=sorted(DECLARED_BY_LAYERS)) +def test_a_layer_is_born_holding_a_draw_from_its_own_initializer(layer_name): + """Creating a parameter and giving it a value are one event, as they are in torch, flax and keras. A + weight left at zero passes no gradient below the last layer of a stack, so a network nobody remembered to + initialize trains its output bias and nothing else.""" + build, layer_input, expected = DECLARED_BY_LAYERS[layer_name] + prediction = build()(layer_input) + + for param in collect_trainable_params(prediction): + declared = expected[param.name] + value = param.get_value() + if declared is None: + assert len(np.unique(value)) > 1, ( + f"{param.name} was born holding one repeated value rather than a draw" + ) + else: + np.testing.assert_allclose(value, declared, err_msg=param.name) + + +def test_a_keyword_initializer_reaches_the_value_and_not_only_the_declaration(): + """The bias case: a keyword used to be recorded as a declaration while the value stayed zero, so it took + effect only if someone happened to call initialize afterwards.""" + sentinel = 7.0 + layer = Linear( + "fc", + n_in=4, + n_out=4, + bias_initializer=CustomInitializer( + lambda shape, dtype, rng: np.full(shape, sentinel, dtype=dtype) + ), + ) + + np.testing.assert_allclose(layer.b.get_value(), sentinel) + + +def test_a_non_default_scheme_still_reaches_every_weight_matrix(): + """The trap in drawing at construction: recording the fallback as a declaration would outrank `scheme` + and silently disable it library-wide, while every test that initializes with the default still passed.""" + X = pt.tensor("features", shape=(None, 4), dtype=floatX) + prediction = Sequential(Linear("fc1", 4, 4), Linear("fc2", 4, 4))(X) + params = collect_trainable_params(prediction) + + values = dict(zip((p.name for p in params), initialize_params(params, scheme="ones", rng=0))) + + np.testing.assert_allclose(values["fc1_W"], 1.0) + np.testing.assert_allclose(values["fc2_W"], 1.0) + np.testing.assert_allclose(values["fc1_b"], 0.0) # its declaration still outranks the scheme + + +def test_construction_draws_do_not_leak_into_a_seeded_initialize(): + """Construction draws from fresh entropy, so two identically seeded runs must still agree: the seeded + initialize has to overwrite every value rather than build on what construction happened to produce.""" + + def seeded_values(): + X = pt.tensor("features", shape=(None, 4), dtype=floatX) + prediction = Sequential(Linear("fc1", 4, 4), Linear("fc2", 4, 4))(X) + params = collect_trainable_params(prediction) + return dict(zip((p.name for p in params), initialize_params(params, rng=0))) + + first, second = seeded_values(), seeded_values() + + assert set(first) == set(second) + for name in first: + np.testing.assert_array_equal(first[name], second[name], err_msg=name) diff --git a/tests/test_model.py b/tests/test_model.py index d480178..895b04c 100644 --- a/tests/test_model.py +++ b/tests/test_model.py @@ -11,6 +11,7 @@ from pytensor_ml.loss import SquaredError from pytensor_ml.model import Model from pytensor_ml.optim import sgd +from pytensor_ml.state import OneInitializer, ZeroInitializer, initialize_params class TestModelPredict: @@ -149,3 +150,48 @@ def test_compile_train_reduces_loss(): history = [float(step(X_batch, target)) for _ in range(50)] assert history[-1] < history[0] + + +@pytest.mark.parametrize( + "scheme", + ["zeros", "ones", ZeroInitializer(), OneInitializer()], + ids=["zeros_by_name", "ones_by_name", "zeros_instance", "ones_instance"], +) +def test_initialize_refuses_a_constant_scheme(scheme): + """A constant is the right initializer for one parameter and never for a network: identical weights leave + no gradient to tell two units in a layer apart. Reachable by name and by instance, so both are refused.""" + X = pt.tensor("X", shape=(None, 8)) + y = Sequential(Linear("fc1", 8, 4), ReLU(), Linear("fc2", 4, 2))(X) + + with pytest.raises(ValueError, match="same value"): + Model(X, y).initialize(scheme, seed=0) + + +def test_the_refusal_is_deliberately_not_symmetric_with_initialize_params(): + """`initialize_params` takes the parameters it is handed, so "set these to zeros" is a coherent request + and several tests rely on a constant scheme as a distinguishable value. `Model.initialize` means the whole + network, where a constant is always wrong. Moving the check down to be consistent would break both this + contrast and about fifteen tests that use it.""" + X = pt.tensor("X", shape=(None, 8)) + y = Sequential(Linear("fc1", 8, 4), ReLU(), Linear("fc2", 4, 2))(X) + model = Model(X, y) + + values = initialize_params(model.weights, scheme="zeros") + + assert all(np.all(value == 0) for value in values) + with pytest.raises(ValueError, match="same value"): + model.initialize("zeros") + + +def test_initialize_still_accepts_a_constant_declared_on_one_parameter(): + """The escape hatch the refusal points at: a constant on the parameter you mean, rather than on all of + them. A zeroed output head is a real technique, and a declaration is how it is said.""" + X = pt.tensor("X", shape=(None, 8)) + first = Linear("fc1", 8, 4) + head = Linear("head", 4, 2, weight_initializer=ZeroInitializer()) + y = Sequential(first, ReLU(), head)(X) + + Model(X, y).initialize("xavier_normal", seed=0) + + np.testing.assert_array_equal(head.W.get_value(), 0) # its declaration outranks the scheme + assert np.abs(first.W.get_value()).min() > 0 # and the scheme still reached its sibling diff --git a/tests/test_state.py b/tests/test_state.py index daa1517..09fa781 100644 --- a/tests/test_state.py +++ b/tests/test_state.py @@ -217,3 +217,13 @@ def test_the_normal_scheme_is_reachable_by_name(): [value] = initialize_params([parameter], scheme="normal", rng=0) assert value.std() == pytest.approx(0.01, rel=0.05) # the default std + + +def test_initial_value_draws_at_floatx_in_the_requested_shape(): + """What every layer relies on when it creates a parameter: the dtype matches the graph's, so nothing has + to remember `config.floatX` at each of the five call sites.""" + value = XavierNormalInitializer().initial_value((6, 4)) + + assert value.shape == (6, 4) + assert str(value.dtype) == pytensor.config.floatX + assert len(np.unique(value)) > 1 # drawn, not filled