From f813beba12a1d19add0f084771e2ce4e7e1ecd0c Mon Sep 17 00:00:00 2001
From: ChrisW09 <50968720+ChrisW09@users.noreply.github.com>
Date: Wed, 21 Oct 2026 15:46:22 +0200
Subject: [PATCH 1/2] fix(nn): initialize random-signs batch-ensemble scaling
 in place

The "random-signs" entry of the scaling-init table built a new tensor with
torch.sign(torch.randn_like(x)) and discarded it, so the r/s scaling
parameters of LinearBatchEnsembleLayer, RNNBatchEnsembleLayer and
MultiHeadAttentionBatchEnsemble kept the uninitialized torch.empty memory
(often zeros, which blocks the gradient of the shared weight matrix).

Add an in-place _random_signs_ initializer that fills the tensor with
-1/+1 values, matching the other torch.nn.init entries.

Fixes #497
---
 deeptab/nn/blocks/common.py                   | 17 +++++--
 .../test_architecture_contracts.py            | 16 +++++++
 tests/nn/blocks/test_common_blocks.py         | 44 +++++++++++++++++++
 3 files changed, 74 insertions(+), 3 deletions(-)

diff --git a/deeptab/nn/blocks/common.py b/deeptab/nn/blocks/common.py
index 48fcb058..7d072558 100644
--- a/deeptab/nn/blocks/common.py
+++ b/deeptab/nn/blocks/common.py
@@ -885,6 +885,17 @@ 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.
@@ -928,7 +939,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),
         }
 
@@ -1028,7 +1039,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 +1298,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),
         }
 
diff --git a/tests/architectures/test_architecture_contracts.py b/tests/architectures/test_architecture_contracts.py
index ec9090b5..cf6dc98b 100644
--- a/tests/architectures/test_architecture_contracts.py
+++ b/tests/architectures/test_architecture_contracts.py
@@ -148,6 +148,22 @@ 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]
+    assert set(first_layer.r.detach().unique().tolist()) == {-1.0, 1.0}
+    assert set(first_layer.s.detach().unique().tolist()) == {-1.0, 1.0}
+    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
diff --git a/tests/nn/blocks/test_common_blocks.py b/tests/nn/blocks/test_common_blocks.py
index 6c30d620..82b4b595 100644
--- a/tests/nn/blocks/test_common_blocks.py
+++ b/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

From 0baced039d49771aa029e25b35f8c0de410a46a0 Mon Sep 17 00:00:00 2001
From: ChrisW09 <50968720+ChrisW09@users.noreply.github.com>
Date: Wed, 21 Oct 2026 23:50:04 +0200
Subject: [PATCH 2/2] docs: describe the scaling_init options and avoid a flaky
 sign check

TabMConfig and the batch-ensemble layers now say what 'ones',
'random-signs' and 'normal' initialize the scaling factors to, and the
TabM page names the values. The TabM test checks |r| == 1 instead of
expecting both signs in a 4 x 3 tensor, which can draw a single sign.

Refs #497
---
 deeptab/configs/models/tabm_config.py         |  4 ++-
 deeptab/nn/blocks/common.py                   | 25 ++++++++++++++++++-
 docs/model_zoo/stable/tabm.md                 | 14 +++++------
 .../test_architecture_contracts.py            |  5 ++--
 4 files changed, 37 insertions(+), 11 deletions(-)

diff --git a/deeptab/configs/models/tabm_config.py b/deeptab/configs/models/tabm_config.py
index f1d7b738..625ea3ed 100644
--- a/deeptab/configs/models/tabm_config.py
+++ b/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'
diff --git a/deeptab/nn/blocks/common.py b/deeptab/nn/blocks/common.py
index 7d072558..d5a3e53d 100644
--- a/deeptab/nn/blocks/common.py
+++ b/deeptab/nn/blocks/common.py
@@ -900,6 +900,25 @@ 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__(
@@ -910,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
@@ -1013,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
diff --git a/docs/model_zoo/stable/tabm.md b/docs/model_zoo/stable/tabm.md
index 2fe42d33..7fb75c80 100644
--- a/docs/model_zoo/stable/tabm.md
+++ b/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
 
diff --git a/tests/architectures/test_architecture_contracts.py b/tests/architectures/test_architecture_contracts.py
index cf6dc98b..006e81b3 100644
--- a/tests/architectures/test_architecture_contracts.py
+++ b/tests/architectures/test_architecture_contracts.py
@@ -157,8 +157,9 @@ def test_tabm_random_signs_scaling_init_trains_first_layer():
     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]
-    assert set(first_layer.r.detach().unique().tolist()) == {-1.0, 1.0}
-    assert set(first_layer.s.detach().unique().tolist()) == {-1.0, 1.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