fix(nn): initialize random-signs batch-ensemble scaling in place

deeptab/configs/models/tabm_config.py

@@ -32,7 +32,9 @@ class TabMConfig(BaseModelConfig):
     ensemble_bias : bool, default=True
         Whether to use a unique bias term for each ensemble member.
     scaling_init : Literal['ones', 'random-signs', 'normal'], default='ones'
-        Initialization method for scaling weights.
+        Initialization of the ensemble scaling factors: ``'ones'`` sets them
+        to 1, ``'random-signs'`` draws each from {-1, +1} with equal
+        probability, and ``'normal'`` draws them from N(0, 1).
     average_ensembles : bool, default=False
         Whether to average the outputs of the ensembles.
     model_type : Literal['mini', 'full'], default='mini'

deeptab/nn/blocks/common.py

@@ -885,10 +885,40 @@ def forward(self, x):
         return x
 
 
+def _random_signs_(tensor: torch.Tensor) -> torch.Tensor:
+    """Fill *tensor* in place with random ``-1``/``+1`` values.
+
+    Used as the ``"random-signs"`` scaling initialization of the batch-ensemble
+    layers. Like the ``torch.nn.init`` functions it modifies its argument, so the
+    ensemble scaling parameters do not keep the memory they were allocated with.
+    """
+    with torch.no_grad():
+        return tensor.copy_(torch.bernoulli(torch.full_like(tensor, 0.5)) * 2 - 1)
+
+
 class LinearBatchEnsembleLayer(nn.Module):
     """A configurable BatchEnsemble layer that supports optional input scaling, output scaling,
     and output bias terms as per the 'BatchEnsemble' paper.
     It provides initialization options for scaling terms to diversify ensemble members.
+
+    Parameters
+    ----------
+    in_features : int
+        Number of input features.
+    out_features : int
+        Number of output features.
+    ensemble_size : int
+        Number of ensemble members.
+    ensemble_scaling_in : bool, default=True
+        Whether to use input scaling for each ensemble member.
+    ensemble_scaling_out : bool, default=True
+        Whether to use output scaling for each ensemble member.
+    ensemble_bias : bool, default=False
+        Whether to use a unique bias term for each ensemble member.
+    scaling_init : {'ones', 'random-signs', 'normal'}, default='ones'
+        Initialization of the scaling factors: ``'ones'`` sets them to 1,
+        ``'random-signs'`` draws each from {-1, +1} with equal probability,
+        and ``'normal'`` draws them from N(0, 1).
     """
 
     def __init__(
@@ -899,7 +929,7 @@ def __init__(
         ensemble_scaling_in: bool = True,
         ensemble_scaling_out: bool = True,
         ensemble_bias: bool = False,
-        scaling_init: Literal["ones", "random-signs"] = "ones",
+        scaling_init: Literal["ones", "random-signs", "normal"] = "ones",
     ):
         super().__init__()
         self.in_features = in_features
@@ -928,7 +958,7 @@ def reset_parameters(self, scaling_init: Literal["ones", "random-signs", "normal
         # Initialize scaling factors r and s based on selected initialization
         scaling_init_fn = {
             "ones": nn.init.ones_,
-            "random-signs": lambda x: torch.sign(torch.randn_like(x)),
+            "random-signs": _random_signs_,
             "normal": lambda x: nn.init.normal_(x, mean=0.0, std=1.0),
         }
 
@@ -1002,6 +1032,10 @@ def __init__(
             Whether to use output scaling for each ensemble member.
         ensemble_bias : bool, default=False
             Whether to use a unique bias term for each ensemble member.
+        scaling_init : {'ones', 'random-signs', 'normal'}, default='ones'
+            Initialization of the scaling factors: ``'ones'`` sets them to 1,
+            ``'random-signs'`` draws each from {-1, +1} with equal probability,
+            and ``'normal'`` draws them from N(0, 1).
         """
         super().__init__()
         self.input_size = input_size
@@ -1028,7 +1062,7 @@ def reset_parameters(self, scaling_init: Literal["ones", "random-signs", "normal
         # Initialize scaling factors r and s based on selected initialization
         scaling_init_fn = {
             "ones": nn.init.ones_,
-            "random-signs": lambda x: torch.sign(torch.randn_like(x)),
+            "random-signs": _random_signs_,
             "normal": lambda x: nn.init.normal_(x, mean=0.0, std=1.0),
         }
 
@@ -1287,7 +1321,7 @@ def reset_parameters(self, scaling_init: Literal["ones", "random-signs", "normal
         # Initialize scaling factors r and s based on selected initialization
         scaling_init_fn = {
             "ones": nn.init.ones_,
-            "random-signs": lambda x: torch.sign(torch.randn_like(x)),
+            "random-signs": _random_signs_,
             "normal": lambda x: nn.init.normal_(x, mean=0.0, std=1.0),
         }
 

docs/model_zoo/stable/tabm.md

@@ -73,13 +73,13 @@ model = TabMClassifier(
 
 Key settings:
 
-| Setting             | Typical range                          | Effect                                           |
-| ------------------- | -------------------------------------- | ------------------------------------------------ |
-| `ensemble_size`     | `8` to `64`                            | Number of virtual ensemble members.              |
-| `layer_sizes`       | `[128, 128]` to `[512, 256, 128]`      | Shared MLP capacity.                             |
-| `model_type`        | `"mini"` or `"full"`                   | Amount of member-specific scaling.               |
-| `average_ensembles` | `False` or `True`                      | Return per-member outputs or average internally. |
-| `scaling_init`      | `"ones"`, `"random-signs"`, `"normal"` | Diversity initialization for scaling factors.    |
+| Setting             | Typical range                          | Effect                                                   |
+| ------------------- | -------------------------------------- | -------------------------------------------------------- |
+| `ensemble_size`     | `8` to `64`                            | Number of virtual ensemble members.                      |
+| `layer_sizes`       | `[128, 128]` to `[512, 256, 128]`      | Shared MLP capacity.                                     |
+| `model_type`        | `"mini"` or `"full"`                   | Amount of member-specific scaling.                       |
+| `average_ensembles` | `False` or `True`                      | Return per-member outputs or average internally.         |
+| `scaling_init`      | `"ones"`, `"random-signs"`, `"normal"` | Scaling-factor initialization: 1, random ±1, or N(0, 1). |
 
 ## When To Use
 

tests/architectures/test_architecture_contracts.py

@@ -148,6 +148,23 @@ def test_tabm_normalization_keeps_ensemble_members_independent(norm):
     torch.testing.assert_close(perturbed[:, 1:], baseline[:, 1:])
 
 
+def test_tabm_random_signs_scaling_init_trains_first_layer():
+    model_class = next(model_class for model_class in model_classes if model_class.__name__ == "TabM")
+    config = get_model_config(model_class)
+    config.layer_sizes = [16, 8]
+    config.ensemble_size = 4
+    config.dropout = 0.0
+    config.scaling_init = "random-signs"
+    model = model_class(({f"value_{index}": {"dimension": 1} for index in range(3)}, {}, {}), config=config)
+    first_layer = model.layers[0]
+    # Every factor is -1 or +1 (a 4 x 3 tensor can draw a single sign by chance, so only |r| is checked).
+    assert torch.all(first_layer.r.detach().abs() == 1)
+    assert torch.all(first_layer.s.detach().abs() == 1)
+    model([torch.randn(8, 1) for _ in range(3)], [], []).sum().backward()
+    assert first_layer.W.grad is not None
+    assert first_layer.W.grad.abs().sum() > 0
+
+
 @pytest.mark.parametrize("norm", ["InstanceNorm", "GroupNorm"])
 def test_tabm_rejects_normalization_across_ensemble_members(norm):
     from deeptab.architectures.tabm import TabM

tests/nn/blocks/test_common_blocks.py

@@ -627,6 +627,50 @@ def test_ensemble_bias(self):
         assert out.shape == (B, S, E, 16)
 
 
+class TestRandomSignsScalingInit:
+    """``scaling_init="random-signs"`` must fill the ensemble scaling factors in
+    place instead of leaving the memory they were allocated with."""
+
+    @staticmethod
+    def _scaling_tensors(layer):
+        tensors = []
+        for scaling in (layer.r, layer.s):
+            if isinstance(scaling, nn.ParameterDict):
+                tensors.extend(scaling.values())
+            elif scaling is not None:
+                tensors.append(scaling)
+        return tensors
+
+    @pytest.mark.parametrize(
+        "make_layer",
+        [
+            lambda: LinearBatchEnsembleLayer(
+                in_features=64, out_features=64, ensemble_size=E, scaling_init="random-signs"
+            ),
+            lambda: RNNBatchEnsembleLayer(input_size=64, hidden_size=64, ensemble_size=E, scaling_init="random-signs"),
+            lambda: MultiHeadAttentionBatchEnsemble(
+                embed_dim=64, num_heads=2, ensemble_size=E, scaling_init="random-signs"
+            ),
+        ],
+        ids=["linear", "rnn", "attention"],
+    )
+    def test_scaling_factors_are_random_signs(self, make_layer, monkeypatch):
+        # Allocate float tensors as NaN so uninitialized memory cannot pass as a valid init.
+        empty = torch.empty
+
+        def nan_empty(*size, **kwargs):
+            tensor = empty(*size, **kwargs)
+            return tensor.fill_(float("nan")) if tensor.is_floating_point() else tensor
+
+        monkeypatch.setattr(torch, "empty", nan_empty)
+        layer = make_layer()
+        tensors = self._scaling_tensors(layer)
+        assert tensors
+        for tensor in tensors:
+            values = set(tensor.detach().unique().tolist())
+            assert values == {-1.0, 1.0}
+
+
 class TestmLSTMblock:
     def test_forward_shape(self):
         # hidden_size and num_layers: BlockDiagonal needs hidden_size % num_layers == 0