Fix native init: resolve initializer_range sentinel, thread gain into xavier/kaiming

tests/unit/model_bridge/test_boot_native.py

@@ -1,4 +1,5 @@
 """Tests for ``TransformerBridge.boot_native`` classmethod."""
+
 from __future__ import annotations
 
 import sys
@@ -489,3 +490,40 @@ def test_boot_native_supports_training_step():
     ), "No non-zero gradients after backward"
     optimizer.step()
     optimizer.zero_grad()
+
+
+def test_boot_native_resolves_initializer_range_sentinel():
+    """Regression for #1568 — the resolved initializer range must be used by boot_native."""
+    import math
+
+    cfg = _cfg(init_mode="gpt2")
+    expected = 0.8 / math.sqrt(cfg.d_model)
+
+    assert cfg.initializer_range == pytest.approx(expected)
+
+    bridge = TransformerBridge.boot_native(cfg)
+    assert bridge.W_E.std().item() == pytest.approx(expected, rel=0.15)
+
+
+def test_boot_native_resolves_non_gpt2_initializer_range_sentinel():
+    cfg = _cfg(init_mode="kaiming_normal")
+
+    assert cfg.initializer_range == pytest.approx(1.0)
+
+
+def test_boot_native_kaiming_gain_scales_weights():
+    """Regression for #1568 — xavier/kaiming init must use initializer_range
+    as a multiplicative gain. Without this, the config value is silently
+    ignored and every kaiming/xavier model gets the same fixed scale
+    regardless of what the caller asked for."""
+    cfg_gain_1 = _cfg(init_mode="kaiming_normal", initializer_range=1.0, seed=0)
+    cfg_gain_2 = _cfg(init_mode="kaiming_normal", initializer_range=2.0, seed=0)
+
+    bridge_1 = TransformerBridge.boot_native(cfg_gain_1)
+    bridge_2 = TransformerBridge.boot_native(cfg_gain_2)
+
+    std_1 = bridge_1.W_E.std().item()
+    std_2 = bridge_2.W_E.std().item()
+
+    # Same seed, only gain differs -> std should scale ~proportionally.
+    assert std_2 / std_1 == pytest.approx(2.0)

transformer_lens/config/transformer_bridge_config.py

@@ -2,6 +2,7 @@
 
 from typing import Optional
 
+import numpy as np
 import torch
 
 from transformer_lens.utilities.activation_functions import SOFTCAP_DISABLED
@@ -219,6 +220,18 @@ def __post_init__(self):
         ):
             raise ValueError(f"architecture must be a string, got {type(self.architecture)}")
 
+        # Resolve the initializer_range sentinel (-1.0 means "not set by the user").
+        # Mirrors HookedTransformerConfig.__post_init__ (hooked_transformer_config.py).
+        # Guarded with getattr: this method also runs once from the dataclass
+        # parent's __init__, before self.initializer_range is assigned below.
+        if getattr(self, "initializer_range", None) is not None:
+            if self.initializer_range < 0 and self.init_mode == "gpt2":
+                # Roughly copy the GPT-2 value, but proportional to sqrt(1/d_model)
+                self.initializer_range = 0.8 / np.sqrt(self.d_model)
+            if self.initializer_range < 0 and self.init_mode != "gpt2":
+                # This is the gain parameter for the weight initialisation
+                self.initializer_range = 1.0
+
         # Call parent's __post_init__ after our validation
         if hasattr(super(), "__post_init__"):
             super().__post_init__()

transformer_lens/model_bridge/sources/native/init.py

@@ -72,7 +72,7 @@ def initialize_native_model(
         generator = None
 
     def _staged(
-        fn: Callable[[torch.Tensor], torch.Tensor]
+        fn: Callable[[torch.Tensor], torch.Tensor],
     ) -> Callable[[torch.Tensor], torch.Tensor]:
         def apply(t: torch.Tensor) -> torch.Tensor:
             staging = torch.empty(t.shape, dtype=torch.float32)
@@ -97,6 +97,14 @@ def apply(t: torch.Tensor) -> torch.Tensor:
         # std = 0.8/sqrt(d_model), not GPT-2's paper 0.02 — toy-model training
         # dynamics (e.g. the grokking demo) depend on this scale.
         std = cfg.initializer_range if cfg.initializer_range > 0 else 0.8 / math.sqrt(cfg.d_model)
+
+        # NOTE: this residual output scaling (1/sqrt(2*n_layers), applied only
+        # to output projections below) is NOT present in HookedTransformer's
+        # _init_weights_gpt2 (see transformer_lens/HookedTransformer.py).
+        # Intentional delta for NativeModel: kept because it follows the
+        # original GPT-2 paper's residual-scaling convention and improves
+        # training stability at init for deeper models. Flagged in issue #1568
+        # as a maintainer call; kept + documented rather than removed.
         residual_scale = 1.0 / math.sqrt(2 * cfg.n_layers)
         weight_init = lambda t: nn.init.normal_(
             t, mean=0.0, std=std, generator=generator