feat: add BatchEnsemble for parameter-efficient ensembling (TabM)

BatchEnsembleLinear: shared backbone with per-member rank-1 scaling.
Provides ensemble-level performance at single-model inference cost.
Based on TabM (Gorishniy et al., ICLR 2025).

Author: w4nderlust

ludwig/modules/batch_ensemble.py

@@ -0,0 +1,77 @@
+"""TabM-style BatchEnsemble for parameter-efficient ensembling.
+
+Implements the BatchEnsemble technique from Wen et al. (ICLR 2020) adapted for
+tabular deep learning as described in TabM (Gorishniy et al., ICLR 2025).
+
+A single MLP efficiently imitates an ensemble by sharing a backbone and using
+per-member affine transforms (non-shared scaling vectors). This provides
+ensemble-level performance at single-model inference cost.
+
+Usage:
+    from ludwig.modules.batch_ensemble import BatchEnsembleLinear
+
+    # Replace nn.Linear with BatchEnsembleLinear
+    layer = BatchEnsembleLinear(in_features=128, out_features=64, num_members=4)
+"""
+
+import torch
+import torch.nn as nn
+
+
+class BatchEnsembleLinear(nn.Module):
+    """Linear layer with BatchEnsemble for parameter-efficient ensembling.
+
+    Shares the main weight matrix across ensemble members, but each member has
+    its own rank-1 scaling factors (r_i and s_i):
+        output_i = (s_i * (W @ (r_i * x))) + b
+
+    This adds only O(in + out) parameters per member instead of O(in * out).
+    """
+
+    def __init__(self, in_features: int, out_features: int, num_members: int = 4, bias: bool = True):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.num_members = num_members
+
+        # Shared backbone
+        self.weight = nn.Parameter(torch.randn(out_features, in_features) / in_features**0.5)
+        if bias:
+            self.bias = nn.Parameter(torch.zeros(out_features))
+        else:
+            self.bias = None
+
+        # Per-member scaling vectors (rank-1 perturbations)
+        self.r = nn.Parameter(torch.ones(num_members, in_features))  # input scaling
+        self.s = nn.Parameter(torch.ones(num_members, out_features))  # output scaling
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass with implicit ensemble.
+
+        During training, randomly selects an ensemble member per sample.
+        During eval, averages predictions across all members.
+
+        Args:
+            x: [batch, in_features]
+
+        Returns:
+            [batch, out_features]
+        """
+        if self.training:
+            # Random member assignment per sample
+            member_idx = torch.randint(0, self.num_members, (x.shape[0],), device=x.device)
+            r = self.r[member_idx]  # [batch, in_features]
+            s = self.s[member_idx]  # [batch, out_features]
+
+            # Apply: s * (W @ (r * x)) + b
+            x_scaled = x * r
+            out = torch.nn.functional.linear(x_scaled, self.weight, self.bias)
+            return out * s
+        else:
+            # Average over all members at eval time
+            outputs = []
+            for i in range(self.num_members):
+                x_scaled = x * self.r[i]
+                out = torch.nn.functional.linear(x_scaled, self.weight, self.bias)
+                outputs.append(out * self.s[i])
+            return torch.stack(outputs).mean(dim=0)

tests/ludwig/modules/test_batch_ensemble.py

@@ -0,0 +1,48 @@
+"""Tests for BatchEnsemble."""
+
+import torch
+
+from ludwig.modules.batch_ensemble import BatchEnsembleLinear
+
+
+class TestBatchEnsembleLinear:
+    def test_output_shape(self):
+        layer = BatchEnsembleLinear(128, 64, num_members=4)
+        x = torch.randn(16, 128)
+        out = layer(x)
+        assert out.shape == (16, 64)
+
+    def test_train_vs_eval(self):
+        layer = BatchEnsembleLinear(32, 16, num_members=4)
+        x = torch.randn(8, 32)
+
+        layer.train()
+        out_train = layer(x)
+
+        layer.eval()
+        out_eval = layer(x)
+
+        assert out_train.shape == out_eval.shape
+
+    def test_gradient_flow(self):
+        layer = BatchEnsembleLinear(32, 16, num_members=4)
+        layer.train()
+        x = torch.randn(8, 32, requires_grad=True)
+        out = layer(x)
+        out.sum().backward()
+        assert x.grad is not None
+        assert layer.weight.grad is not None
+        assert layer.r.grad is not None
+        assert layer.s.grad is not None
+
+    def test_per_member_params(self):
+        layer = BatchEnsembleLinear(32, 16, num_members=4)
+        assert layer.r.shape == (4, 32)
+        assert layer.s.shape == (4, 16)
+
+    def test_no_bias(self):
+        layer = BatchEnsembleLinear(32, 16, num_members=2, bias=False)
+        assert layer.bias is None
+        x = torch.randn(4, 32)
+        out = layer(x)
+        assert out.shape == (4, 16)