diff --git a/pytensor_ml/layers/embedding.py b/pytensor_ml/layers/embedding.py index 0772b92..132cfd0 100644 --- a/pytensor_ml/layers/embedding.py +++ b/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) diff --git a/pytensor_ml/layers/linear.py b/pytensor_ml/layers/linear.py index 029aaa8..c2e5504 100644 --- a/pytensor_ml/layers/linear.py +++ b/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) diff --git a/pytensor_ml/layers/norm.py b/pytensor_ml/layers/norm.py index a3ed750..dd4ef75 100644 --- a/pytensor_ml/layers/norm.py +++ b/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,6 +201,9 @@ 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 @@ -193,6 +211,8 @@ def __init__( 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 diff --git a/pytensor_ml/state.py b/pytensor_ml/state.py index 3b4d2b4..9eab637 100644 --- a/pytensor_ml/state.py +++ b/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,6 +119,18 @@ 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)) @@ -100,6 +138,17 @@ def sample(self, shape: tuple[int, ...], dtype: str, rng: np.random.Generator) - 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 diff --git a/tests/test_layers.py b/tests/test_layers.py index 6dd1ea7..7b2f75d 100644 --- a/tests/test_layers.py +++ b/tests/test_layers.py @@ -16,6 +16,7 @@ collect_trainable_params, rewrite_for_prediction, ) +from pytensor_ml.state import CustomInitializer, NormalInitializer, initialize_params floatX = pytensor.config.floatX @@ -344,3 +345,124 @@ def test_batch_norm_variants_agree_on_output_arity(): # batch statistics but must not write them anywhere. assert len(tracked.owner.outputs) == len(untracked.owner.outputs) assert collect_non_trainable_updates(untracked) == {} + + +# Every parameter a layer owns, and the value a network-wide scheme must leave it at. A scheme reaching a +# batch-norm scale would rescale a normalized activation by a random factor, defeating the layer; reaching a +# bias would undo the zero start these layers are documented to have. +FEATURES = pt.tensor("features", shape=(None, 4), dtype=floatX) +IDS = pt.tensor("ids", shape=(None, 4), dtype="int32") + +DECLARED_BY_LAYERS = { + "Linear": (lambda: Linear("fc", n_in=4, n_out=4), FEATURES, {"fc_W": None, "fc_b": 0.0}), + "Embedding": ( + lambda: Embedding("emb", n_embeddings=6, n_features=4), + IDS, + {"emb_W": None}, + ), + "BatchNorm2D": (lambda: BatchNorm2D("bn", n_in=4), FEATURES, {"bn_scale": 1.0, "bn_loc": 0.0}), + "LayerNorm": (lambda: LayerNorm("ln", n_in=4), FEATURES, {"ln_scale": 1.0, "ln_loc": 0.0}), +} + + +@pytest.mark.parametrize("layer_name", sorted(DECLARED_BY_LAYERS), ids=sorted(DECLARED_BY_LAYERS)) +def test_a_layer_declares_the_initializers_its_parameters_need(layer_name): + """A parameter whose starting value is part of the layer's definition declares an initializer, so a + network-wide scheme passes it by. The scheme here returns a sentinel, so any parameter it reaches is + obvious and any declaration that stopped working shows up as that sentinel.""" + build, layer_input, expected = DECLARED_BY_LAYERS[layer_name] + prediction = build()(layer_input) + + sentinel = 7.0 + reached_by_the_scheme = CustomInitializer( + lambda shape, dtype, rng: np.full(shape, sentinel, dtype=dtype) + ) + params = collect_trainable_params(prediction) + values = initialize_params(params, scheme=reached_by_the_scheme, rng=0) + + assert {p.name for p in params} == set(expected) + for param, value in zip(params, values): + if expected[param.name] is None: + np.testing.assert_allclose(value, sentinel) # no declaration: the scheme applies + else: + np.testing.assert_allclose(value, expected[param.name]) + + +# One case per keyword: the layer to build with it, the parameter it must reach, and the parameters it must +# leave alone. A keyword that hits the wrong parameter, or that quietly strips a sibling's declaration, is +# the failure this is aimed at. +INITIALIZER_KEYWORDS = { + "Linear.weight": ( + lambda init: Linear("fc", n_in=4, n_out=4, weight_initializer=init), + FEATURES, + "fc_W", + {"fc_b": 0.0}, + ), + "Linear.bias": ( + lambda init: Linear("fc", n_in=4, n_out=4, bias_initializer=init), + FEATURES, + "fc_b", + { + "fc_W": 0.0 + }, # the scheme's value: a bias initializer reaching the weight would show up here + ), + "Embedding.weight": ( + lambda init: Embedding("emb", n_embeddings=6, n_features=4, weight_initializer=init), + IDS, + "emb_W", + {}, + ), + "BatchNorm2D.scale": ( + lambda init: BatchNorm2D("bn", n_in=4, scale_initializer=init), + FEATURES, + "bn_scale", + {"bn_loc": 0.0}, + ), + "BatchNorm2D.loc": ( + lambda init: BatchNorm2D("bn", n_in=4, loc_initializer=init), + FEATURES, + "bn_loc", + {"bn_scale": 1.0}, + ), + "LayerNorm.scale": ( + lambda init: LayerNorm("ln", n_in=4, scale_initializer=init), + FEATURES, + "ln_scale", + {"ln_loc": 0.0}, + ), + "LayerNorm.loc": ( + lambda init: LayerNorm("ln", n_in=4, loc_initializer=init), + FEATURES, + "ln_loc", + {"ln_scale": 1.0}, + ), +} + + +@pytest.mark.parametrize("case", sorted(INITIALIZER_KEYWORDS), ids=sorted(INITIALIZER_KEYWORDS)) +def test_an_initializer_keyword_reaches_only_the_parameter_it_names(case): + build, layer_input, target, siblings = INITIALIZER_KEYWORDS[case] + sentinel = 7.0 + layer = build( + CustomInitializer(lambda shape, dtype, rng: np.full(shape, sentinel, dtype=dtype)) + ) + prediction = layer(layer_input) + + params = collect_trainable_params(prediction) + values = dict(zip((p.name for p in params), initialize_params(params, scheme="zeros", rng=0))) + + np.testing.assert_allclose(values[target], sentinel) + for name, expected in siblings.items(): + np.testing.assert_allclose(values[name], expected) + + +def test_a_bias_initializer_replaces_the_zero_declaration_rather_than_fighting_it(): + """The keyword becomes the declaration, so it survives a network-wide scheme the same way the zero it + replaced did. torch draws biases from a fan-scaled uniform, and this is how you say that here.""" + layer = Linear("fc", n_in=4, n_out=4, bias_initializer=NormalInitializer(0.0, 1.0)) + prediction = layer(pt.tensor("features", shape=(None, 4), dtype=floatX)) + + params = collect_trainable_params(prediction) + 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) diff --git a/tests/test_state.py b/tests/test_state.py index 775ed5c..daa1517 100644 --- a/tests/test_state.py +++ b/tests/test_state.py @@ -11,6 +11,7 @@ _INITIALIZERS, CustomInitializer, InitializationScheme, + NormalInitializer, OneInitializer, XavierNormalInitializer, XavierUniformInitializer, @@ -183,3 +184,36 @@ def test_a_fan_scaled_initializer_rejects_a_parameter_with_no_fans(initializer): instead should say so rather than draw something arbitrary.""" with pytest.raises(ValueError, match="at least two dimensions"): initializer.sample((768,), "float64", np.random.default_rng(0)) + + +def test_a_normal_initializer_draws_at_the_standard_deviation_it_was_given(): + """The fan-scaled initializers derive their spread from the shape; this one is told it, so the same + standard deviation has to come out whatever the shape's fans work out to. GPT-2 applies 0.02 to a + 50257x768 embedding and a 768x768 weight alike, where Xavier gives 0.006 and 0.036. Both orientations + are checked at a size large enough to estimate a standard deviation from; four samples cannot.""" + initializer = NormalInitializer(0.0, 0.02) + + for shape in [(1000, 64), (64, 1000)]: + value = initializer.sample(shape, "float64", np.random.default_rng(0)) + assert value.std() == pytest.approx(0.02, rel=0.05) + assert value.mean() == pytest.approx(0.0, abs=0.001) + + +def test_a_normal_initializer_has_no_fans_to_satisfy(): + """Unlike Xavier it accepts a 1-D parameter, which is the point: a bias drawn from a normal is torch's + convention and needs no fan computation. Asserting the spread rather than only the shape, so this says + the draw was right and not merely that nothing raised.""" + value = NormalInitializer(0.0, 1.0).sample((4096,), "float64", np.random.default_rng(0)) + + assert value.shape == (4096,) + assert value.std() == pytest.approx(1.0, rel=0.05) + + +def test_the_normal_scheme_is_reachable_by_name(): + """Its arguments both have defaults, which is what lets it into a registry whose entries are built with + no arguments. Anything parameterized differently has to be passed as an instance.""" + parameter = trainable(np.zeros((100, 100), dtype="float64"), "w") + + [value] = initialize_params([parameter], scheme="normal", rng=0) + + assert value.std() == pytest.approx(0.01, rel=0.05) # the default std