From c67d9af1709a0e6af5623c4fabd3e6624ff3b650 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Mon, 10 Aug 2026 08:04:28 -0500 Subject: [PATCH 1/2] Let attention take an out-projection initializer --- pytensor_ml/layers/attention.py | 21 ++++++++++++++++++++- tests/test_transformer.py | 17 +++++++++++++++++ 2 files changed, 37 insertions(+), 1 deletion(-) diff --git a/pytensor_ml/layers/attention.py b/pytensor_ml/layers/attention.py index cfabd9b..da1cb8a 100644 --- a/pytensor_ml/layers/attention.py +++ b/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, ) diff --git a/tests/test_transformer.py b/tests/test_transformer.py index 68e8d5c..06512b5 100644 --- a/tests/test_transformer.py +++ b/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,18 @@ 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_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) From 0b231c487c9f2fd9ba0eb869d0186ff001c43d17 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Mon, 10 Aug 2026 08:05:13 -0500 Subject: [PATCH 2/2] Forward a residual initializer through TransformerBlock --- pytensor_ml/layers/transformer.py | 30 +++++++++++++++-- tests/test_transformer.py | 55 +++++++++++++++++++++++++++++++ 2 files changed, 83 insertions(+), 2 deletions(-) diff --git a/pytensor_ml/layers/transformer.py b/pytensor_ml/layers/transformer.py index 2faf060..a4b3b05 100644 --- a/pytensor_ml/layers/transformer.py +++ b/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,6 +56,8 @@ 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 @@ -58,7 +65,13 @@ def __init__( 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 diff --git a/tests/test_transformer.py b/tests/test_transformer.py index 06512b5..fc998aa 100644 --- a/tests/test_transformer.py +++ b/tests/test_transformer.py @@ -144,6 +144,61 @@ def test_block_forwards_attention_config(rng): 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."""