Let a transformer block initialize its residual projections

pytensor_ml/layers/attention.py

@@ -6,6 +6,7 @@
 
 from pytensor_ml.base import Layer, UnaryLayerOp
 from pytensor_ml.layers.linear import Linear
+from pytensor_ml.state import Initializer
 
 
 class AttentionLayer(UnaryLayerOp):
@@ -152,6 +153,11 @@ class MultiheadAttention(Layer):
         Include bias terms in the projections. Default is True.
     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.
     """
 
     def __init__(
@@ -162,6 +168,8 @@ def __init__(
         n_kv_head: int | None = None,
         bias: bool = True,
         is_causal: bool = False,
+        *,
+        out_proj_initializer: Initializer | None = None,
     ):
         if n_embd % n_head != 0:
             raise ValueError(f"n_embd ({n_embd}) must be divisible by n_head ({n_head})")
@@ -179,7 +187,13 @@ def __init__(
         self.q_proj = Linear(f"{self.name}_q_proj", n_embd, n_head * self.head_dim, bias)
         self.k_proj = Linear(f"{self.name}_k_proj", n_embd, n_kv_head * self.head_dim, bias)
         self.v_proj = Linear(f"{self.name}_v_proj", n_embd, n_kv_head * self.head_dim, bias)
-        self.out_proj = Linear(f"{self.name}_out_proj", n_head * self.head_dim, n_embd, bias)
+        self.out_proj = Linear(
+            f"{self.name}_out_proj",
+            n_head * self.head_dim,
+            n_embd,
+            bias,
+            weight_initializer=out_proj_initializer,
+        )
 
     def _split_heads(self, x: pt.TensorVariable, n_head: int) -> pt.TensorVariable:
         # (..., seq, n_head * head_dim) -> (..., n_head, seq, head_dim). split_dims keeps the static
@@ -225,6 +239,8 @@ class CausalSelfAttention(MultiheadAttention):
         Number of key/value heads, for grouped-query attention. Defaults to ``n_head``.
     bias : bool, optional
         Include bias terms in the projections. Default is True.
+    out_proj_initializer : Initializer, optional
+        How the output projection's weight is drawn. See :class:`MultiheadAttention`.
     """
 
     def __init__(
@@ -234,6 +250,8 @@ def __init__(
         n_head: int,
         n_kv_head: int | None = None,
         bias: bool = True,
+        *,
+        out_proj_initializer: Initializer | None = None,
     ):
         super().__init__(
             name if name else "CausalSelfAttention",
@@ -242,6 +260,7 @@ def __init__(
             n_kv_head=n_kv_head,
             bias=bias,
             is_causal=True,
+            out_proj_initializer=out_proj_initializer,
         )
 
 

pytensor_ml/layers/transformer.py

@@ -6,6 +6,7 @@
 from pytensor_ml.layers.dropout import Dropout
 from pytensor_ml.layers.linear import Linear
 from pytensor_ml.layers.norm import LayerNorm
+from pytensor_ml.state import Initializer
 
 
 def _identity(x: pt.TensorVariable) -> pt.TensorVariable:
@@ -41,6 +42,10 @@ class FeedForward(Layer):
         Activation applied to the hidden layer. Default is :class:`GELU`.
     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.
     """
 
     def __init__(
@@ -51,14 +56,22 @@ def __init__(
         mlp_ratio: int = 4,
         activation: Activation | None = None,
         bias: bool = True,
+        *,
+        fc_out_initializer: Initializer | None = None,
     ):
         self.name = name if name else "FeedForward"
         self.d_model = d_model
         self.hidden_dim = hidden_dim if hidden_dim is not None else mlp_ratio * d_model
         self.activation = activation if activation is not None else GELU()
 
         self.fc_in = Linear(f"{self.name}_fc_in", d_model, self.hidden_dim, bias)
-        self.fc_out = Linear(f"{self.name}_fc_out", self.hidden_dim, d_model, bias)
+        self.fc_out = Linear(
+            f"{self.name}_fc_out",
+            self.hidden_dim,
+            d_model,
+            bias,
+            weight_initializer=fc_out_initializer,
+        )
 
     def __call__(self, x: pt.TensorLike) -> pt.TensorVariable:
         x = pt.as_tensor(x)
@@ -110,6 +123,12 @@ class TransformerBlock(Layer):
         Number of key/value heads for grouped-query attention. Defaults to ``n_head``.
     epsilon : float, optional
         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
+        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.
     """
 
     def __init__(
@@ -126,6 +145,7 @@ def __init__(
         is_causal: bool = False,
         n_kv_head: int | None = None,
         epsilon: float = 1e-5,
+        residual_initializer: Initializer | None = None,
     ):
         self.name = name if name else "TransformerBlock"
         self.d_model = d_model
@@ -140,9 +160,15 @@ def __init__(
             n_kv_head=n_kv_head,
             bias=bias,
             is_causal=is_causal,
+            out_proj_initializer=residual_initializer,
         )
         self.ff = FeedForward(
-            f"{self.name}_ff", d_model, mlp_ratio=mlp_ratio, activation=activation, bias=bias
+            f"{self.name}_ff",
+            d_model,
+            mlp_ratio=mlp_ratio,
+            activation=activation,
+            bias=bias,
+            fc_out_initializer=residual_initializer,
         )
         self.attn_dropout = (
             Dropout(f"{self.name}_attn_dropout", p=dropout) if dropout > 0 else _identity

tests/test_transformer.py

@@ -8,8 +8,10 @@
 from pytensor.tensor.random.op import RandomVariable
 
 from pytensor_ml.activations import ReLU
+from pytensor_ml.layers.attention import CausalSelfAttention
 from pytensor_ml.layers.transformer import FeedForward, TransformerBlock
 from pytensor_ml.pytensorf import collect_trainable_params
+from pytensor_ml.state import CustomInitializer, initialize_params
 
 floatX = pytensor.config.floatX
 # The python linker evaluates the same graph without a slow numba compile of these block-sized graphs.
@@ -140,3 +142,73 @@ def test_block_forwards_attention_config(rng):
 
     X_np = rng.normal(size=(1, 4, 8)).astype(floatX)
     assert not np.allclose(out.eval({X: X_np}, mode=FAST), masked_out.eval({X: X_np}, mode=FAST))
+
+
+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."""
+    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)))
+
+    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")
+    for name in [
+        "blk_attn_q_proj_W",
+        "blk_attn_k_proj_W",
+        "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")
+
+
+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)
+    )
+    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)
+        )
+    )
+
+    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),
+        ("blk_norm1_loc", 0.0),
+        ("blk_norm2_scale", 1.0),
+        ("blk_norm2_loc", 0.0),
+    ]:
+        np.testing.assert_allclose(values[name], expected, err_msg=name)
+
+
+def test_a_causal_attention_forwards_the_out_projection_initializer():
+    """CausalSelfAttention builds its projections through its base class, so the keyword has to survive the
+    super().__init__ call rather than being silently dropped by the subclass."""
+    sentinel = 7.0
+    scaled = CustomInitializer(lambda shape, dtype, rng: np.full(shape, sentinel, dtype=dtype))
+    attention = CausalSelfAttention("attn", n_embd=8, n_head=2, out_proj_initializer=scaled)
+    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)))
+
+    np.testing.assert_allclose(values["attn_out_proj_W"], sentinel)
+    np.testing.assert_allclose(values["attn_q_proj_W"], 0.0)