From 2ebf2c940f1a8e935771132e618eddb50b892a7e Mon Sep 17 00:00:00 2001
From: jessegrabowski <jessegrabowski@gmail.com>
Date: Mon, 10 Aug 2026 21:15:04 -0500
Subject: [PATCH 1/3] Let initialize_params take per-parameter initializers

---
 pytensor_ml/state.py | 40 +++++++++++++++++--------
 tests/test_state.py  | 71 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 98 insertions(+), 13 deletions(-)

diff --git a/pytensor_ml/state.py b/pytensor_ml/state.py
index 5977cfc..650dcb5 100644
--- a/pytensor_ml/state.py
+++ b/pytensor_ml/state.py
@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from collections.abc import Callable, Sequence
+from collections.abc import Callable, Mapping, Sequence
 from typing import Literal
 
 import numpy as np
@@ -200,8 +200,14 @@ 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."""
+def _initializer_for(
+    param: SharedVariable,
+    initializers: Mapping[SharedVariable, Initializer],
+    default: Initializer,
+) -> Initializer:
+    """The initializer to draw ``param`` with: one named for it, else its own, else ``default``."""
+    if param in initializers:
+        return initializers[param]
     declared = param.initializer if isinstance(param, TrainableParameter) else None
     return default if declared is None else declared
 
@@ -230,7 +236,8 @@ def require_varying_scheme(scheme: InitializationSchemeLike) -> None:
             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())`."
+            "`Linear(..., weight_initializer=ZeroInitializer())`, or name that parameter in "
+            "`initializers={parameter: ZeroInitializer()}`."
         )
 
 
@@ -238,25 +245,31 @@ def initialize_params(
     params: Sequence[SharedVariable],
     scheme: InitializationSchemeLike = "xavier_normal",
     rng: RandomState | None = None,
+    initializers: Mapping[SharedVariable, Initializer] | None = None,
 ) -> list[np.ndarray]:
     """
-    Initialize parameter values using the specified scheme.
+    Initialize parameter values, resolving each parameter's initializer in three steps.
 
-    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.
+    An ``initializers`` entry for a parameter wins; failing that, a
+    :class:`~pytensor_ml.params.TrainableParameter` that declares its own ``initializer`` uses it; failing
+    that, ``scheme`` applies. The middle step is what leaves batch norm at its unit scale while the weight
+    matrices around it are drawn from the requested scheme, and the first is how a caller overrules that
+    for a parameter they hold.
 
     Parameters
     ----------
     params
         SharedVariables to initialize values for.
     scheme
-        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).
+        Initialization scheme for the parameters that ``initializers`` does not name and that declare none
+        of their own: 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.
+    initializers
+        Mapping from a parameter to the initializer to draw it with, taking precedence over both its own
+        declaration and ``scheme``. Keyed by the parameter itself rather than its name, so nothing here
+        depends on how a layer names what it builds.
 
     Returns
     -------
@@ -267,7 +280,8 @@ def initialize_params(
     rng = np.random.default_rng(rng)
 
     initializer = _resolve_scheme(scheme)
+    named = {} if initializers is None else initializers
     return [
-        _declared_initializer(param, default=initializer)._sample_like(param, rng)
+        _initializer_for(param, named, default=initializer)._sample_like(param, rng)
         for param in params
     ]
diff --git a/tests/test_state.py b/tests/test_state.py
index 09fa781..01597ee 100644
--- a/tests/test_state.py
+++ b/tests/test_state.py
@@ -21,6 +21,11 @@
 )
 
 
+def unit_constant(value: float = 7.0) -> CustomInitializer:
+    """An initializer whose draws are unmistakable, so a value says which of the three sources produced it."""
+    return CustomInitializer(lambda shape, dtype, rng: np.full(shape, value, dtype=dtype))
+
+
 def test_scheme_names_match_the_initializer_registry():
     assert set(get_args(InitializationScheme)) == set(_INITIALIZERS)
 
@@ -227,3 +232,69 @@ def test_initial_value_draws_at_floatx_in_the_requested_shape():
     assert value.shape == (6, 4)
     assert str(value.dtype) == pytensor.config.floatX
     assert len(np.unique(value)) > 1  # drawn, not filled
+
+
+class TestPerParameterInitializers:
+    """Naming a parameter is how a caller overrules its declaration, which is otherwise impossible: the
+    declaration always wins, and the only escape was assigning `param.initializer` between construction and
+    initialization."""
+
+    def test_a_named_initializer_beats_a_declaration(self):
+        # A norm scale declares ones precisely so a scheme cannot move it; naming the parameter is the one
+        # thing that should, since the caller picked this parameter rather than every parameter.
+        scale = trainable(np.zeros(4, dtype="float64"), "scale", initializer=OneInitializer())
+
+        [value] = initialize_params(
+            [scale], scheme="zeros", initializers={scale: unit_constant()}, rng=0
+        )
+
+        np.testing.assert_allclose(value, 7.0)
+
+    def test_a_declaration_beats_the_scheme_when_nothing_names_it(self):
+        scale = trainable(np.zeros(4, dtype="float64"), "scale", initializer=OneInitializer())
+        weight = trainable(np.zeros((4, 4), dtype="float64"), "w")
+
+        scale_value, weight_value = initialize_params(
+            [scale, weight], scheme="zeros", initializers={weight: unit_constant()}, rng=0
+        )
+
+        np.testing.assert_allclose(
+            scale_value, 1.0
+        )  # its declaration, untouched by the entry below
+        np.testing.assert_allclose(weight_value, 7.0)
+
+    def test_the_scheme_applies_where_neither_a_named_initializer_nor_a_declaration_exists(self):
+        weight = trainable(np.zeros((4, 4), dtype="float64"), "w")
+        other = trainable(np.zeros((4, 4), dtype="float64"), "other")
+
+        weight_value, other_value = initialize_params(
+            [weight, other], scheme="ones", initializers={weight: unit_constant()}, rng=0
+        )
+
+        np.testing.assert_allclose(weight_value, 7.0)
+        np.testing.assert_allclose(other_value, 1.0)
+
+    def test_an_entry_is_keyed_by_the_parameter_not_its_name(self):
+        """Two parameters can share a name -- nothing prevents it, and optimizer state collides on it rather
+        than parameters -- so an entry keyed by name would reach both. Identity reaches one."""
+        first = trainable(np.zeros((4, 4), dtype="float64"), "w")
+        second = trainable(np.zeros((4, 4), dtype="float64"), "w")
+
+        first_value, second_value = initialize_params(
+            [first, second], scheme="zeros", initializers={first: unit_constant()}, rng=0
+        )
+
+        np.testing.assert_allclose(first_value, 7.0)
+        np.testing.assert_allclose(second_value, 0.0)
+
+    def test_naming_a_parameter_that_is_not_being_initialized_does_nothing(self):
+        """The mapping is consulted per parameter rather than iterated, so an entry for something outside
+        `params` is inert instead of an error -- one dict can serve several initialize calls."""
+        weight = trainable(np.zeros((4, 4), dtype="float64"), "w")
+        elsewhere = trainable(np.zeros((4, 4), dtype="float64"), "elsewhere")
+
+        [value] = initialize_params(
+            [weight], scheme="zeros", initializers={elsewhere: unit_constant()}, rng=0
+        )
+
+        np.testing.assert_allclose(value, 0.0)

From dc860de5cab9953bfd98d64bcc171128ce92908f Mon Sep 17 00:00:00 2001
From: jessegrabowski <jessegrabowski@gmail.com>
Date: Mon, 10 Aug 2026 21:18:00 -0500
Subject: [PATCH 2/3] Thread per-parameter initializers through
 Model.initialize

---
 pytensor_ml/model.py | 24 ++++++++++++++++--------
 tests/test_model.py  | 38 +++++++++++++++++++++++++++++++++++++-
 2 files changed, 53 insertions(+), 9 deletions(-)

diff --git a/pytensor_ml/model.py b/pytensor_ml/model.py
index 33c14bb..0bd00c0 100644
--- a/pytensor_ml/model.py
+++ b/pytensor_ml/model.py
@@ -1,8 +1,8 @@
-from collections.abc import Sequence
+from collections.abc import Mapping, Sequence
 
 import numpy as np
 
-from pytensor.compile import Function
+from pytensor.compile import Function, SharedVariable
 from pytensor.graph.basic import Variable
 from pytensor.printing import debugprint
 from pytensor.tensor.variable import TensorVariable
@@ -13,6 +13,7 @@
 from pytensor_ml.pytensorf import collect_trainable_params, compile_predict
 from pytensor_ml.state import (
     InitializationSchemeLike,
+    Initializer,
     initialize_params,
     require_varying_scheme,
 )
@@ -41,27 +42,34 @@ def initialize(
         self,
         scheme: InitializationSchemeLike = "xavier_normal",
         seed: int | np.random.Generator | None = None,
+        initializers: Mapping[SharedVariable, Initializer] | None = None,
     ) -> "Model":
         """
         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.
+        identity transform, and biases stay at zero. An ``initializers`` entry outranks both.
 
         Parameters
         ----------
         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. 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'.
+            Initialization scheme for the weights that ``initializers`` does not name and that declare
+            none of their own: the name of a built-in 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; put one on the parameter that wants it.
+            Default 'xavier_normal'.
         seed : int or numpy Generator, optional
             Seed for reproducible initialization.
+        initializers : dict mapping parameter to Initializer, optional
+            How to draw specific parameters, in place of their declaration and of ``scheme``. Keyed by the
+            parameter object, reached through the layer that owns it -- ``{block.ff.fc_in.b: ...}`` -- so a
+            parameter no constructor keyword exposes is still addressable.
         """
         require_varying_scheme(scheme)
         parameters = self.weights
-        for parameter, value in zip(parameters, initialize_params(parameters, scheme, rng=seed)):
+        values = initialize_params(parameters, scheme, rng=seed, initializers=initializers)
+        for parameter, value in zip(parameters, values):
             parameter.set_value(value)
         return self
 
diff --git a/tests/test_model.py b/tests/test_model.py
index 895b04c..d251eba 100644
--- a/tests/test_model.py
+++ b/tests/test_model.py
@@ -11,7 +11,12 @@
 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
+from pytensor_ml.state import (
+    CustomInitializer,
+    OneInitializer,
+    ZeroInitializer,
+    initialize_params,
+)
 
 
 class TestModelPredict:
@@ -195,3 +200,34 @@ def test_initialize_still_accepts_a_constant_declared_on_one_parameter():
 
     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
+
+
+def test_initialize_takes_per_parameter_initializers():
+    """The model-level entry point a user actually calls. Reaching a parameter through the layer that owns
+    it is what makes a bias inside a composed layer addressable, since no constructor keyword exposes one."""
+    X = pt.tensor("X", shape=(None, 8))
+    fc1 = Linear("fc1", 8, 4)
+    norm = BatchNorm2D("norm", n_in=4)
+    y = Sequential(fc1, norm, ReLU(), Linear("fc2", 4, 2))(X)
+    drawn = CustomInitializer(lambda shape, dtype, rng: np.full(shape, 7.0, dtype=dtype))
+
+    Model(X, y).initialize("xavier_normal", seed=0, initializers={fc1.b: drawn, norm.scale: drawn})
+
+    np.testing.assert_allclose(fc1.b.get_value(), 7.0)  # outranks its zero declaration
+    np.testing.assert_allclose(norm.scale.get_value(), 7.0)  # and the norm's unit declaration
+    np.testing.assert_allclose(norm.loc.get_value(), 0.0)  # unnamed, so still its declaration
+    assert len(np.unique(fc1.W.get_value())) > 1  # the scheme, since nothing named it
+
+
+def test_a_constant_reaches_one_parameter_through_initializers():
+    """The second escape hatch the constant-scheme refusal names, alongside a declaration. A zeroed output
+    head is the same technique either way; this is the route for a parameter you did not construct."""
+    X = pt.tensor("X", shape=(None, 8))
+    first = Linear("fc1", 8, 4)
+    head = Linear("head", 4, 2)
+    y = Sequential(first, ReLU(), head)(X)
+
+    Model(X, y).initialize("xavier_normal", seed=0, initializers={head.W: ZeroInitializer()})
+
+    np.testing.assert_array_equal(head.W.get_value(), 0)
+    assert np.abs(first.W.get_value()).min() > 0  # the scheme still reached its sibling

From 91952ca102b8dd4098135d96c018cc26ab60f7dc Mon Sep 17 00:00:00 2001
From: jessegrabowski <jessegrabowski@gmail.com>
Date: Mon, 10 Aug 2026 21:47:10 -0500
Subject: [PATCH 3/3] Remove the network-wide initialization scheme

Every parameter now declares the law it is drawn from, so `initialize` means redraw from those declarations under one seed. Linear and Embedding could not declare their own default while a scheme existed, since a declaration shadowed it -- that trap, the refusal of a constant scheme, and the three-way precedence all go with it.
---
 pytensor_ml/layers/attention.py   |   7 +-
 pytensor_ml/layers/embedding.py   |  13 ++-
 pytensor_ml/layers/linear.py      |  15 ++--
 pytensor_ml/layers/norm.py        |  20 ++---
 pytensor_ml/layers/transformer.py |   7 +-
 pytensor_ml/model.py              |  35 +++-----
 pytensor_ml/params.py             |  15 ++--
 pytensor_ml/state.py              |  92 +++++++-------------
 tests/test_layers.py              |  57 +++++--------
 tests/test_model.py               |  85 ++++++++-----------
 tests/test_state.py               | 134 +++++++++++++++---------------
 tests/test_transformer.py         |  32 +++----
 12 files changed, 204 insertions(+), 308 deletions(-)

diff --git a/pytensor_ml/layers/attention.py b/pytensor_ml/layers/attention.py
index da1cb8a..b4516ec 100644
--- a/pytensor_ml/layers/attention.py
+++ b/pytensor_ml/layers/attention.py
@@ -154,10 +154,9 @@ class MultiheadAttention(Layer):
     is_causal : bool, optional
         Apply a causal mask in the attention. Default is False.
     out_proj_initializer : Initializer, optional
-        How the output projection's weight is drawn, in place of whatever scheme
-        :meth:`~pytensor_ml.model.Model.initialize` is given. The three input projections are unaffected,
-        which is what a scaling applied only to the projection writing back into a residual stream needs.
-        Left to the scheme when omitted.
+        How the output projection's weight is drawn. The three input projections are unaffected, which is
+        what a scaling applied only to the projection writing back into a residual stream needs. Xavier
+        normal when omitted, as for any other weight.
     """
 
     def __init__(
diff --git a/pytensor_ml/layers/embedding.py b/pytensor_ml/layers/embedding.py
index 9d16064..92f36b9 100644
--- a/pytensor_ml/layers/embedding.py
+++ b/pytensor_ml/layers/embedding.py
@@ -29,10 +29,10 @@ class Embedding(Layer):
     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.
+        How the table is drawn, at construction and on every redraw. Xavier normal when omitted, which puts
+        the vocabulary size in the denominator -- 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__(
@@ -47,15 +47,14 @@ def __init__(
         self.n_embeddings = n_embeddings
         self.n_features = n_features
 
-        # 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.
+        # Drawn here, as in Linear, and for the same reason.
         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,
+            initializer=W_initializer,
         )
 
     def __call__(self, ids: pt.TensorLike) -> pt.TensorVariable:
diff --git a/pytensor_ml/layers/linear.py b/pytensor_ml/layers/linear.py
index 17e0eb3..c78ab2a 100644
--- a/pytensor_ml/layers/linear.py
+++ b/pytensor_ml/layers/linear.py
@@ -33,11 +33,9 @@ class Linear(Layer):
     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.
+        Add the learned shift :math:`b`, which starts at zero and is redrawn there. 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.
+        How :math:`W` is drawn, at construction and on every redraw. Xavier normal when omitted.
     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}})`.
@@ -46,7 +44,7 @@ class Linear(Layer):
     -----
     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.
+    reproducible.
     """
 
     def __init__(
@@ -65,16 +63,15 @@ def __init__(
         self.bias = bias
 
         # 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.
+        # weights passes no gradient below its last layer. The initializer is declared as well as used, so
+        # `initialize(seed=...)` redraws from the same law.
         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,
+            initializer=W_initializer,
         )
 
         if self.bias:
diff --git a/pytensor_ml/layers/norm.py b/pytensor_ml/layers/norm.py
index 56465d8..92f0d59 100644
--- a/pytensor_ml/layers/norm.py
+++ b/pytensor_ml/layers/norm.py
@@ -85,9 +85,9 @@ def _affine_parameters(
     """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. A caller who wants
-    something else says so, and their choice becomes the declaration.
+    Both declare their initializer, so a redraw returns them to the identity transform -- 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.
     """
     resolved_loc = ZeroInitializer() if loc_initializer is None else loc_initializer
     resolved_scale = OneInitializer() if scale_initializer is None else scale_initializer
@@ -167,11 +167,10 @@ class BatchNorm2D(Layer):
         0.1.
     affine : bool, optional
         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.
+        transform :math:`\gamma = 1`, :math:`\beta = 0`, which a redraw returns them to. 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.
+        How :math:`\gamma` is drawn. Ones when omitted, which is the identity transform; drawing a random
+        factor to rescale a normalized activation by would defeat the layer.
     loc_initializer : Initializer, optional
         How :math:`\beta` is drawn. Zeros when omitted.
     track_running_stats : bool, optional
@@ -316,11 +315,10 @@ class LayerNorm(Layer):
         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`, starting from the identity
-        transform :math:`\gamma = 1`, :math:`\beta = 0`, which a network-wide initialization scheme
-        leaves in place. Default is True.
+        transform :math:`\gamma = 1`, :math:`\beta = 0`, which a redraw returns them to. 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.
+        How :math:`\gamma` is drawn. Ones when omitted, which is the identity transform; drawing a random
+        factor to rescale a normalized activation by would defeat the layer.
     loc_initializer : Initializer, optional
         How :math:`\beta` is drawn. Zeros when omitted.
     """
diff --git a/pytensor_ml/layers/transformer.py b/pytensor_ml/layers/transformer.py
index a4b3b05..805def8 100644
--- a/pytensor_ml/layers/transformer.py
+++ b/pytensor_ml/layers/transformer.py
@@ -43,9 +43,8 @@ class FeedForward(Layer):
     bias : bool, optional
         Include bias terms in both linear layers. Default is True.
     fc_out_initializer : Initializer, optional
-        How the second layer's weight is drawn, in place of whatever scheme
-        :meth:`~pytensor_ml.model.Model.initialize` is given. The hidden layer is unaffected. Left to the
-        scheme when omitted.
+        How the second layer's weight is drawn. The hidden layer is unaffected. Xavier normal when omitted,
+        as for any other weight.
     """
 
     def __init__(
@@ -125,7 +124,7 @@ class TransformerBlock(Layer):
         Constant added to the layer-norm variance for numerical stability. Default is 1e-5.
     residual_initializer : Initializer, optional
         How the two projections that write back into the residual stream are drawn -- attention's output
-        projection and the feed-forward's second layer. Their siblings are left to the scheme, which is what
+        projection and the feed-forward's second layer. Their siblings keep the default, which is what
         GPT-style initialization asks for: a residual stream accumulates one contribution per block, so those
         projections are scaled by :math:`1/\sqrt{2 n_\text{layer}}` while the rest are not. The depth is not
         known here, so the scaling belongs in the initializer the caller passes.
diff --git a/pytensor_ml/model.py b/pytensor_ml/model.py
index 0bd00c0..7e63257 100644
--- a/pytensor_ml/model.py
+++ b/pytensor_ml/model.py
@@ -2,7 +2,8 @@
 
 import numpy as np
 
-from pytensor.compile import Function, SharedVariable
+from pytensor.compile import Function
+from pytensor.compile.sharedvalue import SharedVariable
 from pytensor.graph.basic import Variable
 from pytensor.printing import debugprint
 from pytensor.tensor.variable import TensorVariable
@@ -11,12 +12,7 @@
 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,
-    Initializer,
-    initialize_params,
-    require_varying_scheme,
-)
+from pytensor_ml.state import Initializer, initialize_params
 
 
 class Model:
@@ -40,35 +36,28 @@ def weights(self) -> list[TrainableParameter]:
 
     def initialize(
         self,
-        scheme: InitializationSchemeLike = "xavier_normal",
         seed: int | np.random.Generator | None = None,
         initializers: Mapping[SharedVariable, Initializer] | None = None,
     ) -> "Model":
         """
-        Initialize the trainable weights in place and return self.
+        Redraw every trainable weight from its own initializer, 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. An ``initializers`` entry outranks both.
+        Each parameter was already built holding a draw, so this is what makes a run reproducible rather
+        than what makes it trainable: one seed regenerates all of them. Every layer declares how each
+        parameter it builds is drawn, so a batch norm layer redraws to its identity transform and a bias
+        to zero.
 
         Parameters
         ----------
-        scheme : str or Initializer
-            Initialization scheme for the weights that ``initializers`` does not name and that declare
-            none of their own: the name of a built-in 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; put one on the parameter that wants it.
-            Default 'xavier_normal'.
         seed : int or numpy Generator, optional
             Seed for reproducible initialization.
         initializers : dict mapping parameter to Initializer, optional
-            How to draw specific parameters, in place of their declaration and of ``scheme``. Keyed by the
-            parameter object, reached through the layer that owns it -- ``{block.ff.fc_in.b: ...}`` -- so a
-            parameter no constructor keyword exposes is still addressable.
+            How to draw specific parameters, in place of what they declare. Keyed by the parameter object,
+            reached through the layer that owns it -- ``{block.ff.fc_in.b: ...}`` -- so a parameter no
+            constructor keyword exposes is still addressable.
         """
-        require_varying_scheme(scheme)
         parameters = self.weights
-        values = initialize_params(parameters, scheme, rng=seed, initializers=initializers)
+        values = initialize_params(parameters, rng=seed, initializers=initializers)
         for parameter, value in zip(parameters, values):
             parameter.set_value(value)
         return self
diff --git a/pytensor_ml/params.py b/pytensor_ml/params.py
index 51efb13..92d22b7 100644
--- a/pytensor_ml/params.py
+++ b/pytensor_ml/params.py
@@ -17,9 +17,9 @@ class TrainableParameter(TensorSharedVariable):
     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.
+        The law this parameter's value is drawn from, which :func:`~pytensor_ml.state.initialize_params`
+        redraws it from. Every layer declares one for each parameter it builds. None means a redraw has
+        nothing to go on and raises.
     """
 
     initializer: "Initializer | None" = None
@@ -59,8 +59,8 @@ def trainable(
     Create a shared variable marked as a trainable parameter.
 
     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.
+    optimizer updates exactly these. A parameter also declares the initializer it is drawn from, which is
+    what lets one seed reproduce every value in a network.
 
     Parameters
     ----------
@@ -74,9 +74,8 @@ def trainable(
     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.
+        The law to redraw this parameter from, whether that is a unit scale, a zero bias, or a fan-scaled
+        draw. Default None, which leaves the parameter with no law and raises on a redraw.
     **kwargs
         Additional arguments passed to the SharedVariable constructor.
     """
diff --git a/pytensor_ml/state.py b/pytensor_ml/state.py
index 650dcb5..205d6cc 100644
--- a/pytensor_ml/state.py
+++ b/pytensor_ml/state.py
@@ -126,8 +126,7 @@ def fans(shape: tuple[int, ...]) -> tuple[int, int]:
         raise ValueError(
             f"A fan-scaled initializer needs a parameter of at least two dimensions to size its draws, but "
             f"got shape {shape}. A bias or a norm scale has no fans; give it an initializer of its own -- "
-            "`trainable(value, name, initializer=ZeroInitializer())` -- or initialize it with the 'zeros' "
-            "or 'ones' scheme."
+            "`trainable(value, name, initializer=ZeroInitializer())`."
         )
     receptive_field = int(np.prod(shape[2:]))
     return shape[0] * receptive_field, shape[1] * receptive_field
@@ -187,101 +186,66 @@ def sample(self, shape: tuple[int, ...], dtype: str, rng: np.random.Generator) -
         return self._sample_fn(shape, dtype, rng)
 
 
+# The name of each initializer, for the config a network is serialized to: an initializer crosses that
+# boundary as a name, since a class cannot. Every entry must construct with no arguments to come back.
 _INITIALIZERS: dict[str, type[Initializer]] = {
     "zeros": ZeroInitializer,
     "ones": OneInitializer,
     "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
-
 
 def _initializer_for(
-    param: SharedVariable,
-    initializers: Mapping[SharedVariable, Initializer],
-    default: Initializer,
+    param: SharedVariable, initializers: Mapping[SharedVariable, Initializer]
 ) -> Initializer:
-    """The initializer to draw ``param`` with: one named for it, else its own, else ``default``."""
+    """The initializer to draw ``param`` with: one named for it, else the one it declares."""
     if param in initializers:
         return initializers[param]
-    declared = param.initializer if isinstance(param, TrainableParameter) else None
-    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__
+    declared = param.initializer if isinstance(param, TrainableParameter) else None
+    if declared is None:
         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())`, or name that parameter in "
-            "`initializers={parameter: ZeroInitializer()}`."
+            f"{param.name!r} declares no initializer, so there is no law to redraw it from. Every layer in "
+            "this library declares one for each parameter it builds; a parameter built by hand needs the "
+            "same -- `trainable(value, name, initializer=XavierNormalInitializer())` -- or an entry naming "
+            "it in `initializers={parameter: XavierNormalInitializer()}`."
         )
+    return declared
 
 
 def initialize_params(
     params: Sequence[SharedVariable],
-    scheme: InitializationSchemeLike = "xavier_normal",
     rng: RandomState | None = None,
     initializers: Mapping[SharedVariable, Initializer] | None = None,
 ) -> list[np.ndarray]:
     """
-    Initialize parameter values, resolving each parameter's initializer in three steps.
+    Redraw each parameter from its own initializer, or from the one ``initializers`` names for it.
 
-    An ``initializers`` entry for a parameter wins; failing that, a
-    :class:`~pytensor_ml.params.TrainableParameter` that declares its own ``initializer`` uses it; failing
-    that, ``scheme`` applies. The middle step is what leaves batch norm at its unit scale while the weight
-    matrices around it are drawn from the requested scheme, and the first is how a caller overrules that
-    for a parameter they hold.
+    Every parameter carries the law its value comes from, so there is nothing to decide here beyond the
+    seed: a batch norm layer redraws to its unit scale, a bias to zero, and a weight matrix to a fresh
+    Xavier draw. Pass ``initializers`` to draw a particular parameter some other way.
 
     Parameters
     ----------
-    params
-        SharedVariables to initialize values for.
-    scheme
-        Initialization scheme for the parameters that ``initializers`` does not name and that declare none
-        of their own: 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.
-    initializers
-        Mapping from a parameter to the initializer to draw it with, taking precedence over both its own
-        declaration and ``scheme``. Keyed by the parameter itself rather than its name, so nothing here
-        depends on how a layer names what it builds.
+    params : sequence of SharedVariable
+        Parameters to draw values for. Each must declare an initializer or be named in ``initializers``.
+    rng : int or numpy Generator, optional
+        Random number generator. A fresh one is created when omitted.
+    initializers : dict mapping parameter to Initializer, optional
+        How to draw specific parameters, taking precedence over what they declare. Keyed by the parameter
+        itself rather than its name, so nothing here depends on how a layer names what it builds.
 
     Returns
     -------
-    list of np.ndarray
-        Initialized values matching the shapes and dtypes of params.
+    list of ndarray
+        Drawn values, matching the shapes and dtypes of ``params``.
     """
     # Resolve once and share: a seed handed to each _sample_like call would repeat draws across parameters.
     rng = np.random.default_rng(rng)
+    if initializers is None:
+        initializers = {}
 
-    initializer = _resolve_scheme(scheme)
-    named = {} if initializers is None else initializers
-    return [
-        _initializer_for(param, named, default=initializer)._sample_like(param, rng)
-        for param in params
-    ]
+    return [_initializer_for(param, initializers)._sample_like(param, rng) for param in params]
diff --git a/tests/test_layers.py b/tests/test_layers.py
index 00e43b1..10d1df5 100644
--- a/tests/test_layers.py
+++ b/tests/test_layers.py
@@ -347,9 +347,10 @@ def test_batch_norm_variants_agree_on_output_arity():
     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.
+# Every parameter a layer owns, and the value a redraw must give it. A constant is part of the layer's
+# definition -- a batch-norm scale drawn as a random factor would rescale a normalized activation and defeat
+# the layer, and a drawn bias would undo the zero start these layers are documented to have. None means the
+# declaration is a real draw rather than a constant, so the assertion is that the value varies.
 FEATURES = pt.tensor("features", shape=(None, 4), dtype=floatX)
 IDS = pt.tensor("ids", shape=(None, 4), dtype="int32")
 
@@ -366,26 +367,21 @@ def test_batch_norm_variants_agree_on_output_arity():
 
 
 @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."""
+def test_a_layer_declares_an_initializer_for_every_parameter_it_builds(layer_name):
+    """A redraw consults nothing but the declaration, so a parameter declaring none cannot be redrawn at all.
+    Every parameter here has to come back with its layer's definition intact."""
     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)
+
+    values = initialize_params(params, 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
+            assert len(np.unique(value)) > 1, f"{param.name} redrew to one repeated value"
         else:
-            np.testing.assert_allclose(value, expected[param.name])
+            np.testing.assert_allclose(value, expected[param.name], err_msg=param.name)
 
 
 # One case per keyword: the layer to build with it, the parameter it must reach, and the parameters it must
@@ -402,9 +398,7 @@ def test_a_layer_declares_the_initializers_its_parameters_need(layer_name):
         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
+        {"fc_W": None},  # a bias initializer reaching the weight would show up as the sentinel here
     ),
     "Embedding.weight": (
         lambda init: Embedding("emb", n_embeddings=6, n_features=4, weight_initializer=init),
@@ -449,21 +443,24 @@ def test_an_initializer_keyword_reaches_only_the_parameter_it_names(case):
     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)))
+    values = dict(zip((p.name for p in params), initialize_params(params, rng=0)))
 
     np.testing.assert_allclose(values[target], sentinel)
     for name, expected in siblings.items():
-        np.testing.assert_allclose(values[name], expected)
+        if expected is None:
+            assert len(np.unique(values[name])) > 1, name
+        else:
+            np.testing.assert_allclose(values[name], expected, err_msg=name)
 
 
 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."""
+    """The keyword becomes the declaration, so a redraw uses it rather than the zero it replaced. 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)))
+    values = dict(zip((p.name for p in params), initialize_params(params, rng=0)))
 
     assert not np.allclose(values["fc_b"], 0.0)
 
@@ -503,20 +500,6 @@ def test_a_keyword_initializer_reaches_the_value_and_not_only_the_declaration():
     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."""
diff --git a/tests/test_model.py b/tests/test_model.py
index d251eba..c38807d 100644
--- a/tests/test_model.py
+++ b/tests/test_model.py
@@ -11,12 +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 (
-    CustomInitializer,
-    OneInitializer,
-    ZeroInitializer,
-    initialize_params,
-)
+from pytensor_ml.state import CustomInitializer, ZeroInitializer
 
 
 class TestModelPredict:
@@ -88,7 +83,7 @@ def test_leaves_a_norm_layer_at_the_identity_transform(self, norm_layer):
         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)
+        Model(X, y).initialize(seed=0)
 
         np.testing.assert_array_equal(norm.scale.get_value(), 1)
         np.testing.assert_array_equal(norm.loc.get_value(), 0)
@@ -98,11 +93,28 @@ def test_draws_weight_matrices_and_leaves_biases_at_zero(self):
         fc1 = Linear("fc1", 8, 4)
         y = Sequential(fc1, ReLU(), Linear("fc2", 4, 2))(X)
 
-        Model(X, y).initialize("xavier_normal", seed=0)
+        Model(X, y).initialize(seed=0)
 
         assert np.abs(fc1.W.get_value()).min() > 0
         np.testing.assert_array_equal(fc1.b.get_value(), 0)
 
+    def test_the_seed_is_what_a_reproducible_run_rests_on(self):
+        """The whole remaining job of `initialize`: same seed, same network, same values. The differing-seed
+        half is what fails if the seed is dropped on the way to the draw, which same-seed alone would miss --
+        construction draws from fresh entropy, so two networks agree only because the seed made them."""
+
+        def values_for(seed):
+            X = pt.tensor("X", shape=(None, 8))
+            y = Sequential(Linear("fc1", 8, 4), ReLU(), Linear("fc2", 4, 2))(X)
+            model = Model(X, y).initialize(seed=seed)
+            return {p.name: p.get_value() for p in model.weights}
+
+        first, again, other = values_for(0), values_for(0), values_for(1)
+
+        for name, value in first.items():
+            np.testing.assert_array_equal(value, again[name], err_msg=name)
+        assert not np.array_equal(first["fc1_W"], other["fc1_W"])
+
     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
@@ -111,7 +123,7 @@ def test_a_declared_initializer_does_not_freeze_the_parameter(self):
         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)
+        model = Model(X, y).initialize(seed=0)
 
         target = pt.matrix("target")
         step = model.compile_train(
@@ -157,49 +169,18 @@ def test_compile_train_reduces_loss():
     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."""
+def test_a_layer_keyword_survives_a_redraw():
+    """A zeroed output head is a real technique, and the constructor keyword is how it is asked for. The
+    redraw has to honor it rather than treat every weight matrix as a fresh Xavier draw."""
     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)
+    Model(X, y).initialize(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
+    np.testing.assert_array_equal(head.W.get_value(), 0)  # redrawn from its declaration
+    assert np.abs(first.W.get_value()).min() > 0  # and its sibling from the layer default
 
 
 def test_initialize_takes_per_parameter_initializers():
@@ -211,23 +192,23 @@ def test_initialize_takes_per_parameter_initializers():
     y = Sequential(fc1, norm, ReLU(), Linear("fc2", 4, 2))(X)
     drawn = CustomInitializer(lambda shape, dtype, rng: np.full(shape, 7.0, dtype=dtype))
 
-    Model(X, y).initialize("xavier_normal", seed=0, initializers={fc1.b: drawn, norm.scale: drawn})
+    Model(X, y).initialize(seed=0, initializers={fc1.b: drawn, norm.scale: drawn})
 
     np.testing.assert_allclose(fc1.b.get_value(), 7.0)  # outranks its zero declaration
     np.testing.assert_allclose(norm.scale.get_value(), 7.0)  # and the norm's unit declaration
     np.testing.assert_allclose(norm.loc.get_value(), 0.0)  # unnamed, so still its declaration
-    assert len(np.unique(fc1.W.get_value())) > 1  # the scheme, since nothing named it
+    assert len(np.unique(fc1.W.get_value())) > 1  # its own declaration, since nothing named it
 
 
 def test_a_constant_reaches_one_parameter_through_initializers():
-    """The second escape hatch the constant-scheme refusal names, alongside a declaration. A zeroed output
-    head is the same technique either way; this is the route for a parameter you did not construct."""
+    """The same zeroed head as above, asked for at initialize time rather than at construction -- the route
+    for a parameter whose layer you did not build, such as one inside a loaded network."""
     X = pt.tensor("X", shape=(None, 8))
     first = Linear("fc1", 8, 4)
     head = Linear("head", 4, 2)
     y = Sequential(first, ReLU(), head)(X)
 
-    Model(X, y).initialize("xavier_normal", seed=0, initializers={head.W: ZeroInitializer()})
+    Model(X, y).initialize(seed=0, initializers={head.W: ZeroInitializer()})
 
     np.testing.assert_array_equal(head.W.get_value(), 0)
-    assert np.abs(first.W.get_value()).min() > 0  # the scheme still reached its sibling
+    assert np.abs(first.W.get_value()).min() > 0  # its sibling still drew from its declaration
diff --git a/tests/test_state.py b/tests/test_state.py
index 01597ee..405bf88 100644
--- a/tests/test_state.py
+++ b/tests/test_state.py
@@ -11,8 +11,10 @@
     _INITIALIZERS,
     CustomInitializer,
     InitializationScheme,
+    Initializer,
     NormalInitializer,
     OneInitializer,
+    UnitUniformInitializer,
     XavierNormalInitializer,
     XavierUniformInitializer,
     ZeroInitializer,
@@ -21,22 +23,31 @@
 )
 
 
-def unit_constant(value: float = 7.0) -> CustomInitializer:
-    """An initializer whose draws are unmistakable, so a value says which of the three sources produced it."""
-    return CustomInitializer(lambda shape, dtype, rng: np.full(shape, value, dtype=dtype))
+def unit_constant() -> CustomInitializer:
+    """An initializer whose draws are unmistakable, so a value says which source produced it."""
+    return CustomInitializer(lambda shape, dtype, rng: np.full(shape, 7.0, dtype=dtype))
 
 
-def test_scheme_names_match_the_initializer_registry():
+def test_every_registered_name_maps_to_a_class_that_needs_no_arguments():
+    """The registry is how an initializer crosses into a serialized config, as a name. A class needing a
+    constructor argument could be written out and not read back, so the Literal and the registry have to
+    agree and every entry has to rebuild from its name alone."""
     assert set(get_args(InitializationScheme)) == set(_INITIALIZERS)
 
+    for name, initializer_class in _INITIALIZERS.items():
+        assert isinstance(initializer_class(), Initializer), name
+
 
 class TestInitializeParams:
-    @pytest.mark.parametrize("scheme", sorted(_INITIALIZERS))
-    def test_values_match_parameter_shapes_and_dtypes(self, scheme, simple_network):
+    @pytest.mark.parametrize("initializer_class", sorted(_INITIALIZERS.values(), key=repr))
+    def test_values_match_parameter_shapes_and_dtypes(self, initializer_class, simple_network):
         X, y = simple_network
         params = collect_trainable_params(y)
+        # Only the weight matrices, since a fan-scaled initializer cannot draw a 1-D bias. The biases come
+        # from their own declarations, and the assertion covers every parameter either way.
+        drawn_by = {p: initializer_class() for p in params if p.get_value().ndim > 1}
 
-        values = initialize_params(params, scheme=scheme, rng=np.random.default_rng(0))
+        values = initialize_params(params, rng=np.random.default_rng(0), initializers=drawn_by)
 
         assert len(values) == len(params)
         for param, val in zip(params, values):
@@ -44,27 +55,17 @@ def test_values_match_parameter_shapes_and_dtypes(self, scheme, simple_network):
             assert str(val.dtype) == str(param.get_value().dtype)
 
     def test_xavier_normal_scales_by_fan_in_plus_fan_out(self, simple_network):
+        # Drawn from fc1_W's own declaration, which is what a Linear layer records for its weight.
         X, y = simple_network
         weight = next(p for p in collect_trainable_params(y) if p.name == "fc1_W")
         fan_sum = sum(weight.get_value().shape)
 
-        [value] = initialize_params([weight], scheme="xavier_normal", rng=np.random.default_rng(42))
+        [value] = initialize_params([weight], rng=np.random.default_rng(42))
 
         assert value.std() == pytest.approx(np.sqrt(2.0 / fan_sum), rel=0.05)
         # Xavier uniform targets this same variance; only a normal draw crosses its hard bound.
         assert np.abs(value).max() > np.sqrt(6.0 / fan_sum)
 
-    @pytest.mark.parametrize("scheme", ["zeros", ZeroInitializer()], ids=["by_name", "by_instance"])
-    def test_zeros(self, scheme, simple_network):
-        X, y = simple_network
-        params = collect_trainable_params(y)
-
-        values = initialize_params(params, scheme=scheme)
-
-        assert len(values) == len(params)
-        for val in values:
-            np.testing.assert_array_equal(val, 0)
-
     def test_reproducible_with_seed(self, simple_network):
         X, y = simple_network
         params = collect_trainable_params(y)
@@ -77,10 +78,11 @@ def test_reproducible_with_seed(self, simple_network):
 
     def test_parameters_do_not_all_receive_the_same_draws(self):
         # The seed must be an int: passing a Generator masks the regression this guards.
-        first = trainable(np.zeros((4, 4), dtype="float64"), "first")
-        second = trainable(np.zeros((4, 4), dtype="float64"), "second")
+        drawn = UnitUniformInitializer()
+        first = trainable(np.zeros((4, 4), dtype="float64"), "first", initializer=drawn)
+        second = trainable(np.zeros((4, 4), dtype="float64"), "second", initializer=drawn)
 
-        values = initialize_params([first, second], scheme="unit_uniform", rng=0)
+        values = initialize_params([first, second], rng=0)
 
         assert not np.array_equal(values[0], values[1])
 
@@ -91,7 +93,7 @@ def test_accepts_a_custom_initializer(self):
         ]
         constant = CustomInitializer(lambda shape, dtype, rng: np.full(shape, 7.0, dtype=dtype))
 
-        values = initialize_params(params, scheme=constant)
+        values = initialize_params(params, initializers=dict.fromkeys(params, constant))
 
         assert len(values) == len(params)
         for val in values:
@@ -99,29 +101,39 @@ def test_accepts_a_custom_initializer(self):
 
 
 class TestDeclaredInitializers:
-    def test_a_declared_initializer_wins_over_the_scheme(self):
+    def test_a_declaration_is_what_a_redraw_draws_from(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)
+        [value] = initialize_params([scale], rng=0)
 
         np.testing.assert_array_equal(value, 1)
 
-    def test_undeclared_parameters_still_follow_the_scheme(self):
+    def test_a_parameter_declaring_nothing_refuses_to_be_redrawn(self):
+        """There is no law to draw it from, and the alternative is worse than an error: a hand-built weight
+        left at the value it was given is a stack of zeros that trains only its last bias."""
         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)
+        with pytest.raises(ValueError, match="'weight' declares no initializer"):
+            initialize_params([weight], rng=0)
+
+    def test_naming_a_parameter_that_declares_nothing_is_enough(self):
+        weight = trainable(np.zeros((4, 4), dtype="float64"), "weight")
 
-        assert weight_value.min() > 0
-        np.testing.assert_array_equal(bias_value, 0)
+        [value] = initialize_params([weight], rng=0, initializers={weight: unit_constant()})
 
-    def test_shared_variables_without_the_marker_class_follow_the_scheme(self):
+        np.testing.assert_array_equal(value, 7.0)
+
+    def test_a_shared_variable_without_the_marker_class_has_to_be_named(self):
+        """Plain shared state carries no ``initializer`` attribute at all, so the declaration lookup has to
+        tolerate its absence rather than raise an AttributeError on the way past."""
         state = pytensor.shared(np.zeros(4, dtype="float64"), name="state")
 
-        [value] = initialize_params([state], scheme="unit_uniform", rng=0)
+        with pytest.raises(ValueError, match="'state' declares no initializer"):
+            initialize_params([state], rng=0)
 
-        assert value.min() > 0
+        [value] = initialize_params([state], rng=0, initializers={state: unit_constant()})
+        np.testing.assert_array_equal(value, 7.0)
 
     def test_calling_an_initializer_overrides_a_declaration(self):
         scale = trainable(np.ones(3, dtype="float64"), "scale", initializer=OneInitializer())
@@ -185,8 +197,8 @@ def test_the_fan_in_is_the_dimension_the_layers_treat_as_input():
 )
 def test_a_fan_scaled_initializer_rejects_a_parameter_with_no_fans(initializer):
     """A bias or a norm scale has no fan-in and fan-out, and scaling by the length of a vector is a number
-    with no meaning behind it. Such parameters declare their own initializer; reaching one with a scheme
-    instead should say so rather than draw something arbitrary."""
+    with no meaning behind it. Pointing one of these at such a parameter -- through a layer keyword or an
+    `initializers` entry -- 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))
 
@@ -214,14 +226,12 @@ def test_a_normal_initializer_has_no_fans_to_satisfy():
     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)
+def test_a_normal_initializer_built_from_its_registry_name_has_a_usable_default_spread():
+    """Both arguments default, which is what lets it into the registry at all. The default has to be a
+    spread something could train from, since a config naming 'normal' rebuilds exactly this."""
+    value = _INITIALIZERS["normal"]().sample((100, 100), "float64", np.random.default_rng(0))
 
-    assert value.std() == pytest.approx(0.01, rel=0.05)  # the default std
+    assert value.std() == pytest.approx(0.01, rel=0.05)
 
 
 def test_initial_value_draws_at_floatx_in_the_requested_shape():
@@ -236,26 +246,24 @@ def test_initial_value_draws_at_floatx_in_the_requested_shape():
 
 class TestPerParameterInitializers:
     """Naming a parameter is how a caller overrules its declaration, which is otherwise impossible: the
-    declaration always wins, and the only escape was assigning `param.initializer` between construction and
-    initialization."""
+    declaration is the only thing a redraw consults, and the alternative is assigning `param.initializer`
+    between construction and initialization."""
 
     def test_a_named_initializer_beats_a_declaration(self):
-        # A norm scale declares ones precisely so a scheme cannot move it; naming the parameter is the one
+        # A norm scale declares ones precisely so nothing else can move it; naming the parameter is the one
         # thing that should, since the caller picked this parameter rather than every parameter.
         scale = trainable(np.zeros(4, dtype="float64"), "scale", initializer=OneInitializer())
 
-        [value] = initialize_params(
-            [scale], scheme="zeros", initializers={scale: unit_constant()}, rng=0
-        )
+        [value] = initialize_params([scale], initializers={scale: unit_constant()}, rng=0)
 
         np.testing.assert_allclose(value, 7.0)
 
-    def test_a_declaration_beats_the_scheme_when_nothing_names_it(self):
+    def test_an_unnamed_parameter_keeps_its_declaration(self):
         scale = trainable(np.zeros(4, dtype="float64"), "scale", initializer=OneInitializer())
-        weight = trainable(np.zeros((4, 4), dtype="float64"), "w")
+        weight = trainable(np.zeros((4, 4), dtype="float64"), "w", initializer=ZeroInitializer())
 
         scale_value, weight_value = initialize_params(
-            [scale, weight], scheme="zeros", initializers={weight: unit_constant()}, rng=0
+            [scale, weight], initializers={weight: unit_constant()}, rng=0
         )
 
         np.testing.assert_allclose(
@@ -263,25 +271,15 @@ def test_a_declaration_beats_the_scheme_when_nothing_names_it(self):
         )  # its declaration, untouched by the entry below
         np.testing.assert_allclose(weight_value, 7.0)
 
-    def test_the_scheme_applies_where_neither_a_named_initializer_nor_a_declaration_exists(self):
-        weight = trainable(np.zeros((4, 4), dtype="float64"), "w")
-        other = trainable(np.zeros((4, 4), dtype="float64"), "other")
-
-        weight_value, other_value = initialize_params(
-            [weight, other], scheme="ones", initializers={weight: unit_constant()}, rng=0
-        )
-
-        np.testing.assert_allclose(weight_value, 7.0)
-        np.testing.assert_allclose(other_value, 1.0)
-
     def test_an_entry_is_keyed_by_the_parameter_not_its_name(self):
         """Two parameters can share a name -- nothing prevents it, and optimizer state collides on it rather
         than parameters -- so an entry keyed by name would reach both. Identity reaches one."""
-        first = trainable(np.zeros((4, 4), dtype="float64"), "w")
-        second = trainable(np.zeros((4, 4), dtype="float64"), "w")
+        declared = ZeroInitializer()
+        first = trainable(np.zeros((4, 4), dtype="float64"), "w", initializer=declared)
+        second = trainable(np.zeros((4, 4), dtype="float64"), "w", initializer=declared)
 
         first_value, second_value = initialize_params(
-            [first, second], scheme="zeros", initializers={first: unit_constant()}, rng=0
+            [first, second], initializers={first: unit_constant()}, rng=0
         )
 
         np.testing.assert_allclose(first_value, 7.0)
@@ -290,11 +288,9 @@ def test_an_entry_is_keyed_by_the_parameter_not_its_name(self):
     def test_naming_a_parameter_that_is_not_being_initialized_does_nothing(self):
         """The mapping is consulted per parameter rather than iterated, so an entry for something outside
         `params` is inert instead of an error -- one dict can serve several initialize calls."""
-        weight = trainable(np.zeros((4, 4), dtype="float64"), "w")
+        weight = trainable(np.zeros((4, 4), dtype="float64"), "w", initializer=ZeroInitializer())
         elsewhere = trainable(np.zeros((4, 4), dtype="float64"), "elsewhere")
 
-        [value] = initialize_params(
-            [weight], scheme="zeros", initializers={elsewhere: unit_constant()}, rng=0
-        )
+        [value] = initialize_params([weight], initializers={elsewhere: unit_constant()}, rng=0)
 
         np.testing.assert_allclose(value, 0.0)
diff --git a/tests/test_transformer.py b/tests/test_transformer.py
index fc998aa..4b8037e 100644
--- a/tests/test_transformer.py
+++ b/tests/test_transformer.py
@@ -147,15 +147,15 @@ def test_block_forwards_attention_config(rng):
 def test_a_residual_initializer_reaches_both_projections_into_the_residual_stream():
     """The scaling GPT-2 applies targets the two projections that write back into the residual stream --
     attention's output projection and the feed-forward's second layer -- and must leave the other four weight
-    matrices to the scheme. A keyword that caught a sibling would inflate the very variance it exists to
-    control, and reading the constructors cannot tell you which it reached."""
+    matrices at the layer default. A keyword that caught a sibling would inflate the very variance it exists
+    to control, and reading the constructors cannot tell you which it reached."""
     sentinel = 7.0
     scaled = CustomInitializer(lambda shape, dtype, rng: np.full(shape, sentinel, dtype=dtype))
     block = TransformerBlock("blk", d_model=8, n_head=2, residual_initializer=scaled)
     out = block(pt.tensor("x", shape=(None, 4, 8), dtype=floatX))
 
     params = collect_trainable_params(out)
-    values = dict(zip((p.name for p in params), initialize_params(params, scheme="zeros", rng=0)))
+    values = dict(zip((p.name for p in params), initialize_params(params, rng=0)))
 
     for name in ["blk_attn_out_proj_W", "blk_ff_fc_out_W"]:
         np.testing.assert_allclose(values[name], sentinel, err_msg=f"{name} was not scaled")
@@ -165,30 +165,22 @@ def test_a_residual_initializer_reaches_both_projections_into_the_residual_strea
         "blk_attn_v_proj_W",
         "blk_ff_fc_in_W",
     ]:
-        np.testing.assert_allclose(values[name], 0.0, err_msg=f"{name} should follow the scheme")
+        assert len(np.unique(values[name])) > 1, f"{name} was scaled and should not have been"
 
 
 def test_a_block_built_without_a_residual_initializer_is_unchanged():
-    """Omitting the keyword has to leave every parameter exactly where it was: the weights drawn by the
-    scheme, the biases at zero, the norms at the identity transform. The scheme returns a sentinel rather
-    than ones, so a norm scale that lost its declaration is distinguishable from one that kept it."""
-    sentinel = 7.0
-    reached_by_the_scheme = CustomInitializer(
-        lambda shape, dtype, rng: np.full(shape, sentinel, dtype=dtype)
-    )
+    """Omitting the keyword has to leave every parameter where it was: the weights on a real draw, the biases
+    at zero, the norms at the identity transform. A norm scale that lost its declaration would come back as a
+    draw rather than as ones, which is the failure the constants here catch."""
     block = TransformerBlock("blk", d_model=8, n_head=2)
     out = block(pt.tensor("x", shape=(None, 4, 8), dtype=floatX))
 
     params = collect_trainable_params(out)
-    values = dict(
-        zip(
-            (p.name for p in params), initialize_params(params, scheme=reached_by_the_scheme, rng=0)
-        )
-    )
+    values = dict(zip((p.name for p in params), initialize_params(params, rng=0)))
 
+    for name in ["blk_attn_out_proj_W", "blk_ff_fc_out_W"]:
+        assert len(np.unique(values[name])) > 1, name
     for name, expected in [
-        ("blk_attn_out_proj_W", sentinel),
-        ("blk_ff_fc_out_W", sentinel),
         ("blk_attn_out_proj_b", 0.0),
         ("blk_ff_fc_out_b", 0.0),
         ("blk_norm1_scale", 1.0),
@@ -208,7 +200,7 @@ def test_a_causal_attention_forwards_the_out_projection_initializer():
     out = attention(pt.tensor("x", shape=(None, 4, 8), dtype=floatX))
 
     params = collect_trainable_params(out)
-    values = dict(zip((p.name for p in params), initialize_params(params, scheme="zeros", rng=0)))
+    values = dict(zip((p.name for p in params), initialize_params(params, rng=0)))
 
     np.testing.assert_allclose(values["attn_out_proj_W"], sentinel)
-    np.testing.assert_allclose(values["attn_q_proj_W"], 0.0)
+    assert len(np.unique(values["attn_q_proj_W"])) > 1