Merge pull request #6 from stef1949/codex/evaluate-and-improve-neural-network-design

Author: stef1949

models/ae.py

@@ -7,6 +7,8 @@
 """
 from __future__ import annotations
 
+from typing import Sequence
+
 import torch
 import torch.nn as nn
 
@@ -44,27 +46,96 @@ def set_lambda(self, lambda_: float):
 # ----------------------------
 
 
+def _make_activation(name: str | None) -> nn.Module | None:
+    if name is None:
+        return None
+    name = name.lower()
+    if name == "relu":
+        return nn.ReLU()
+    if name == "tanh":
+        return nn.Tanh()
+    if name == "sigmoid":
+        return nn.Sigmoid()
+    if name in {"silu", "swish"}:
+        return nn.SiLU()
+    if name == "gelu":
+        return nn.GELU()
+    raise ValueError(f"Unsupported activation: {name}")
+
+
 class ResidualBlock(nn.Module):
-    """
-    Residual block with skip connection: x + FFN(x) + LayerNorm
-    """
+    """Feed-forward residual block with optional projection skip."""
 
-    def __init__(self, hidden_size: int, dropout: float = 0.0):
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        *,
+        hidden_features: int | None = None,
+        dropout: float = 0.0,
+        activation: str = "silu",
+        use_layer_norm: bool = True,
+    ) -> None:
         super().__init__()
-        self.ffn = nn.Sequential(
-            nn.Linear(hidden_size, hidden_size),
-            nn.SiLU(),
-            nn.Dropout(dropout),
-            nn.Linear(hidden_size, hidden_size),
-            nn.Dropout(dropout),
-        )
-        self.layer_norm = nn.LayerNorm(hidden_size)
+        hidden_features = out_features if hidden_features is None else hidden_features
+
+        ff_layers: list[nn.Module] = [nn.Linear(in_features, hidden_features)]
+        act = _make_activation(activation)
+        if act is not None:
+            ff_layers.append(act)
+        if dropout > 0:
+            ff_layers.append(nn.Dropout(dropout))
+        ff_layers.append(nn.Linear(hidden_features, out_features))
+        if dropout > 0:
+            ff_layers.append(nn.Dropout(dropout))
+        self.ffn = nn.Sequential(*ff_layers)
+
+        if in_features == out_features:
+            self.shortcut: nn.Module = nn.Identity()
+        else:
+            self.shortcut = nn.Linear(in_features, out_features, bias=False)
 
-    def forward(self, x):
-        return self.layer_norm(x + self.ffn(x))
+        self.layer_norm = nn.LayerNorm(out_features) if use_layer_norm else nn.Identity()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.layer_norm(self.shortcut(x) + self.ffn(x))
 
 
-def make_mlp(sizes, dropout=0.0, last_activation=None, use_residual=False):
+class ResidualStack(nn.Module):
+    """Stack of residual blocks that can change feature dimensionality."""
+
+    def __init__(
+        self,
+        sizes: Sequence[int],
+        *,
+        dropout: float = 0.0,
+        activation: str = "silu",
+    ) -> None:
+        super().__init__()
+        if len(sizes) < 2:
+            raise ValueError("ResidualStack requires at least two layer sizes")
+
+        blocks = [
+            ResidualBlock(
+                in_features=in_f,
+                out_features=out_f,
+                dropout=dropout,
+                activation=activation,
+            )
+            for in_f, out_f in zip(sizes[:-1], sizes[1:])
+        ]
+        self.blocks = nn.Sequential(*blocks)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.blocks(x)
+
+
+def make_mlp(
+    sizes: Sequence[int],
+    dropout: float = 0.0,
+    last_activation: str | None = None,
+    use_residual: bool = False,
+):
     """
     Hidden layers: Linear -> LayerNorm -> SiLU -> Dropout
     Output layer: optional activation per arg
@@ -78,52 +149,34 @@ def make_mlp(sizes, dropout=0.0, last_activation=None, use_residual=False):
                 "Residual networks need at least input, hidden, and output layers"
             )
 
-        # Check that all hidden layers have the same size for residual connections
-        hidden_sizes = sizes[1:-1]
-        if len(set(hidden_sizes)) > 1:
-            raise ValueError(
-                f"For residual connections, all hidden layer sizes must be the same. Got: {hidden_sizes}"
-            )
-
-        hidden_size = hidden_sizes[0]
-        num_hidden_layers = len(hidden_sizes)
-
-        layers = []
-
-        # Input projection to hidden size
-        layers.append(nn.Linear(sizes[0], hidden_size))
-        layers.extend([nn.LayerNorm(hidden_size), nn.SiLU(), nn.Dropout(dropout)])
-
-        # Residual blocks
-        for _ in range(num_hidden_layers):
-            layers.append(ResidualBlock(hidden_size, dropout))
-
-        # Output layer
-        layers.append(nn.Linear(hidden_size, sizes[-1]))
-        if last_activation == "relu":
-            layers.append(nn.ReLU())
-        elif last_activation == "tanh":
-            layers.append(nn.Tanh())
-        elif last_activation == "sigmoid":
-            layers.append(nn.Sigmoid())
-
+        trunk_sizes = sizes[:-1]
+        layers: list[nn.Module] = [
+            ResidualStack(trunk_sizes, dropout=dropout, activation="silu")
+        ]
+        layers.append(nn.Linear(trunk_sizes[-1], sizes[-1]))
+        final_act = _make_activation(last_activation)
+        if final_act is not None:
+            layers.append(final_act)
         return nn.Sequential(*layers)
 
     else:
         # Original implementation
-        layers = []
+        layers: list[nn.Module] = []
         for i in range(len(sizes) - 1):
             in_f, out_f = sizes[i], sizes[i + 1]
             layers.append(nn.Linear(in_f, out_f))
-            if i < len(sizes) - 2:
-                layers += [nn.LayerNorm(out_f), nn.SiLU(), nn.Dropout(dropout)]
+            is_last = i == len(sizes) - 2
+            if not is_last:
+                layers.append(nn.LayerNorm(out_f))
+                act = _make_activation("silu")
+                if act is not None:
+                    layers.append(act)
+                if dropout > 0:
+                    layers.append(nn.Dropout(dropout))
             else:
-                if last_activation == "relu":
-                    layers += [nn.ReLU()]
-                elif last_activation == "tanh":
-                    layers += [nn.Tanh()]
-                elif last_activation == "sigmoid":
-                    layers += [nn.Sigmoid()]
+                final_act = _make_activation(last_activation)
+                if final_act is not None:
+                    layers.append(final_act)
         return nn.Sequential(*layers)
 
 

tests/test_residual_and_inference.py

@@ -16,9 +16,9 @@ def test_residual_block_forward(self):
         hidden_size = 128
         batch_size = 16
 
-        block = ResidualBlock(hidden_size, dropout=0.1)
+        block = ResidualBlock(hidden_size, hidden_size, dropout=0.1)
         x = torch.randn(batch_size, hidden_size)
-        
+
         output = block(x)
 
         # Output should have same shape as input
@@ -29,11 +29,18 @@ def test_residual_block_forward(self):
 
     def test_residual_block_dimensions(self):
         """Test ResidualBlock with various dimensions."""
-        for hidden_size in [64, 128, 256]:
-            block = ResidualBlock(hidden_size, dropout=0.0)
-            x = torch.randn(10, hidden_size)
+        for in_dim, out_dim in [(64, 64), (64, 128), (128, 32)]:
+            block = ResidualBlock(in_dim, out_dim, dropout=0.0)
+            x = torch.randn(10, in_dim)
             output = block(x)
-            assert output.shape == (10, hidden_size)
+            assert output.shape == (10, out_dim)
+
+    def test_residual_block_projection_learns(self):
+        """Projection branch should learn non-trivial mappings when dims change."""
+        block = ResidualBlock(32, 64, dropout=0.0)
+        x = torch.randn(4, 32)
+        out = block(x)
+        assert out.shape == (4, 64)
 
 
 class TestMakeMLP:
@@ -47,24 +54,16 @@ def test_make_mlp_without_residual(self):
 
         assert output.shape == (8, 10)
 
-    def test_make_mlp_with_residual_uniform_hidden(self):
-        """Test make_mlp with residual connections and uniform hidden sizes."""
-        # All hidden layers must be the same size for residual connections
-        sizes = [100, 64, 64, 64, 10]  # input, hidden1, hidden2, hidden3, output
+    def test_make_mlp_with_residual_handles_dim_changes(self):
+        """Residual builder should accept non-uniform hidden sizes via projections."""
+        sizes = [100, 128, 64, 32, 10]
         mlp = make_mlp(sizes, dropout=0.1, use_residual=True)
-        
+
         x = torch.randn(8, 100)
         output = mlp(x)
-        
+
         assert output.shape == (8, 10)
-    
-    def test_make_mlp_with_residual_non_uniform_hidden_raises_error(self):
-        """Test that non-uniform hidden sizes raise ValueError with residual=True."""
-        sizes = [100, 64, 32, 16, 10]  # Non-uniform hidden sizes
-        
-        with pytest.raises(ValueError, match="all hidden layer sizes must be the same"):
-            make_mlp(sizes, dropout=0.1, use_residual=True)
-    
+
     def test_make_mlp_residual_minimum_layers(self):
         """Test that residual networks need at least 3 layers."""
         sizes = [100, 10]  # Only input and output
@@ -110,14 +109,14 @@ def test_ae_with_residual_uniform_hidden(self):
         model = AEBatchCorrector(
             n_genes=1000,
             latent_dim=32,
-            enc_hidden=(128, 128, 128),  # Uniform hidden sizes
-            dec_hidden=(128, 128, 128),  # Uniform hidden sizes
+            enc_hidden=(128, 128, 128),
+            dec_hidden=(128, 128, 128),
             n_batches=3,
             n_labels=2,
             dropout=0.1,
             use_residual=True,
         )
-        
+
         x = torch.randn(16, 1000)
         x_hat, b_logits, l_logits, z = model(x, adv_lambda=1.0)
 
@@ -126,17 +125,22 @@ def test_ae_with_residual_uniform_hidden(self):
         assert l_logits.shape == (16, 2)
         assert z.shape == (16, 32)
 
-    def test_ae_with_residual_non_uniform_hidden_raises_error(self):
-        """Test that non-uniform hidden sizes raise error with residual=True."""
-        with pytest.raises(ValueError, match="all hidden layer sizes must be the same"):
-            model = AEBatchCorrector(
-                n_genes=1000,
-                latent_dim=32,
-                enc_hidden=(256, 128),  # Non-uniform
-                dec_hidden=(128, 256),  # Non-uniform
-                n_batches=3,
-                use_residual=True,
-            )
+    def test_ae_with_residual_non_uniform_hidden(self):
+        """Non-uniform hidden sizes should be supported via projection skips."""
+        model = AEBatchCorrector(
+            n_genes=1000,
+            latent_dim=32,
+            enc_hidden=(256, 128),
+            dec_hidden=(128, 256),
+            n_batches=3,
+            use_residual=True,
+        )
+
+        x = torch.randn(8, 1000)
+        x_hat, b_logits, l_logits, z = model(x, adv_lambda=1.0)
+        assert x_hat.shape == (8, 1000)
+        assert b_logits.shape == (8, 3)
+        assert z.shape == (8, 32)
 
     def test_ae_reconstruct_method(self):
         """Test the reconstruct method for inference."""