Repository: FilippoOlivo/SSM
Pull request #28: Fix S6 block, add ResidualBlock and improve S6
Author: login: FilippoOlivo
Labels: name: enhancement
description: New feature or request
Created: 2025-04-17T08:30:50Z | Merged: 2025-04-18T09:36:34Z

This PR align the initialization of dt to the one in the original GitHub repo of Mamba. In preliminary tests, this initialization (and making dt trainable) allow training convergence with S6. However, the test were done on very short sequences


Diff:
diff --git a/ssm/model/block/residual_block.py b/ssm/model/block/residual_block.py
new file mode 100644
index 0000000..4965cbd
--- /dev/null
+++ b/ssm/model/block/residual_block.py
@@ -0,0 +1,39 @@
+import torch
+
+
+class ResidualBlock(torch.nn.Module):
+    """
+    Implementation of a residual block used in the S4 model.
+
+    This block is designed to efficiently model long sequences using selective
+    state space models. Its selection mechanism allows it to focus on relevant
+    parts of the input sequence, making it suitable for tasks such as selective
+    copy.
+
+    The output is computed in an efficient manner by leveraging the parallel
+    scan algorithm.
+    """
+
+    def __init__(self, model):
+        """
+        Initialization of the Residual Block.
+
+        :param int input_dim: The input dimension.
+        :param int hid_dim: The hidden dimension.
+        """
+        super().__init__()
+        self.model = model
+
+    def forward(self, x):
+        """
+        Forward pass of the Residual Block.
+
+        :param torch.Tensor x: The input tensor.
+        :return: The output tensor.
+        :rtype: torch.Tensor
+        """
+        # Compute the output using the model
+        y = self.model(x)
+
+        # Add the residual connection
+        return x + y
diff --git a/ssm/model/block/s4_base_block.py b/ssm/model/block/s4_base_block.py
index 6807f0c..0d8cdc2 100644
--- a/ssm/model/block/s4_base_block.py
+++ b/ssm/model/block/s4_base_block.py
@@ -1,6 +1,6 @@
 import torch
 from .s4_block_interface import S4BlockInterface
-from ...utils import compute_hippo
+from ...utils import compute_hippo, initialize_dt
 
 
 class S4BaseBlock(S4BlockInterface):
@@ -35,7 +35,8 @@ def __init__(
         input_dim,
         hid_dim,
         method,
-        dt=0.1,
+        dt_min=0.001,
+        dt_max=0.01,
         hippo=False,
         **kwargs,
     ):
@@ -46,7 +47,10 @@ def __init__(
         :param int hid_dim: The hidden state dimension.
         :param str method: The forward computation method. Available options
             are: recurrent, convolutional.
-        :param float dt: The time step for discretization. Default is `0.1`.
+        :param float dt_min: The minimum time step for discretization.
+            Default is `0.001`.
+        :param float dt_max: The maximum time step for discretization.
+            Default is `0.01`.
         :param bool hippo: Whether to use the HIPPO matrix for initialization.
             Default is `False`.
         :param dict kwargs: Additional arguments for the class constructor.
@@ -59,10 +63,22 @@ def __init__(
         B = torch.rand(input_dim, hid_dim, 1)
         C = torch.rand(input_dim, 1, hid_dim)
 
+        # Initialize the time step dt
+        dt = (
+            initialize_dt(
+                input_dim=input_dim,
+                dt_min=dt_min,
+                dt_max=dt_max,
+                inverse_softplus=False,
+            )
+            .unsqueeze(-1)
+            .unsqueeze(-1)
+        )
+
         super().__init__(
             input_dim=input_dim,
             hid_dim=hid_dim,
-            dt=torch.Tensor([dt]),
+            dt=dt,
             A=A,
             B=B,
             C=C,
diff --git a/ssm/model/block/s4_block_interface.py b/ssm/model/block/s4_block_interface.py
index 07d79b2..bedffda 100644
--- a/ssm/model/block/s4_block_interface.py
+++ b/ssm/model/block/s4_block_interface.py
@@ -70,19 +70,19 @@ def __init__(self, input_dim, hid_dim, dt, A, B, C, method):
         # Initialize parameters
         self.input_dim = input_dim
         self.hid_dim = hid_dim
-        self.register_buffer("dt", dt)
 
         # Check method
         if method not in ["recurrent", "convolutional"]:
             raise ValueError(f"Unknown method: {method}")
         self.method = method
 
-        # Initialize matrices A, B, and C
-        self.B, self.C = (
+        # Initialize matrices A, B, C and the time step dt
+        self.A, self.B, self.C, self.dt = (
+            torch.nn.Parameter(A),
             torch.nn.Parameter(B),
             torch.nn.Parameter(C),
+            torch.nn.Parameter(dt),
         )
-        self.register_buffer("A", A)
 
     def forward_recurrent(self, x):
         """
diff --git a/ssm/model/block/s4_diagonal_block.py b/ssm/model/block/s4_diagonal_block.py
index ef5e75e..18b6305 100644
--- a/ssm/model/block/s4_diagonal_block.py
+++ b/ssm/model/block/s4_diagonal_block.py
@@ -5,6 +5,7 @@
     compute_S4DLin,
     compute_S4DQuad,
     compute_S4DReal,
+    initialize_dt,
 )
 
 
@@ -45,7 +46,8 @@ def __init__(
         input_dim,
         hid_dim,
         method,
-        dt=0.1,
+        dt_max=0.01,
+        dt_min=0.001,
         initialization="S4D-Inv",
         real_random=False,
         imag_random=False,
@@ -59,7 +61,10 @@ def __init__(
         :param int hid_dim: The hidden state dimension.
         :param str method: The forward computation method. Available options
             are: recurrent, convolutional.
-        :param float dt: The time step for discretization. Default is `0.1`.
+        :param float dt_max: The maximum time step for discretization.
+            Default is `0.01`.
+        :param float dt_min: The minimum time step for discretization.
+            Default is `0.001`.
         :param str initialization: The method for initializing the A matrix.
             Options are: S4D-Inv, S4D-Lin, S4D-Quad, S4D-Real.
             Default is `"S4D-Inv"`.
@@ -82,7 +87,13 @@ def __init__(
         B = torch.rand(input_dim, hid_dim)
         C = torch.rand(input_dim, hid_dim)
 
-        dt = torch.tensor(dt, dtype=torch.float32)
+        # Initialize the time step dt
+        dt = initialize_dt(
+            input_dim=input_dim,
+            dt_max=dt_max,
+            dt_min=dt_min,
+            inverse_softplus=False,
+        )
 
         super().__init__(
             input_dim=input_dim,
diff --git a/ssm/model/block/s4_low_rank_block.py b/ssm/model/block/s4_low_rank_block.py
index eecc9b4..2f02bf8 100644
--- a/ssm/model/block/s4_low_rank_block.py
+++ b/ssm/model/block/s4_low_rank_block.py
@@ -1,6 +1,6 @@
 import torch
 from .s4_block_interface import S4BlockInterface
-from ...utils import compute_hippo, compute_dplr
+from ...utils import compute_hippo, compute_dplr, initialize_dt
 
 
 class S4LowRankBlock(S4BlockInterface):
@@ -23,7 +23,8 @@ def __init__(
         input_dim,
         hid_dim,
         method,
-        dt=0.1,
+        dt_min=0.001,
+        dt_max=0.01,
         hippo=True,
         **kwargs,
     ):
@@ -34,7 +35,10 @@ def __init__(
         :param int hid_dim: The hidden state dimension.
         :param str method: The forward computation method. Low-rank S4
             block only supports the convolutional method.
-        :param float dt: The time step for discretization. Default is `0.1`.
+        :param float dt_min: Minimum time step for discretization. Default is
+            `0.001`.
+        :param float dt_max: Maximum time step for discretization. Default is
+            `0.01`.
         :param bool hippo: Whether to use the HIPPO matrix for initialization.
             Default is `True`.
         :param dict kwargs: Additional arguments for the class constructor.
@@ -43,10 +47,18 @@ def __init__(
         B = torch.nn.Parameter(torch.rand(input_dim, hid_dim))
         C = torch.nn.Parameter(torch.rand(input_dim, hid_dim))
 
+        # Initialize dt
+        dt = initialize_dt(
+            input_dim=input_dim,
+            dt_max=dt_max,
+            dt_min=dt_min,
+            inverse_softplus=False,
+        ).unsqueeze(-1)
+
         super().__init__(
             input_dim=input_dim,
             hid_dim=hid_dim,
-            dt=torch.tensor(dt),
+            dt=dt,
             A=torch.empty((1)),
             B=B,
             C=C,
@@ -131,12 +143,11 @@ def _compute_K(self, L):
 
         # Compute the denominator for the Cauchy product
         g = (
-            (2.0 / self.dt) * (1.0 - self.omega) / (1.0 + self.omega)
+            (2.0 / self.dt)
+            * (1.0 - self.omega.unsqueeze(0))
+            / (1.0 + self.omega.unsqueeze(0))
         ).unsqueeze(-1)
-        denominator = torch.cat(
-            [g - self.Lambda[i : i + 1, :, :] for i in range(self.input_dim)],
-            dim=0,
-        )
+        denominator = g - self.Lambda
 
         # Compute the Cauchy product
         k00, k01, k10, k11 = self._cauchy_dot(a0, a1, b0, b1, denominator)
diff --git a/ssm/model/block/s4_shift_block.py b/ssm/model/block/s4_shift_block.py
index 1c948ee..7265bb5 100644
--- a/ssm/model/block/s4_shift_block.py
+++ b/ssm/model/block/s4_shift_block.py
@@ -39,7 +39,8 @@ def __init__(
         input_dim,
         hid_dim,
         method,
-        dt=0.1,
+        dt_min=0.001,
+        dt_max=0.01,
         **kwargs,
     ):
         """
@@ -49,15 +50,20 @@ def __init__(
         :param int hid_dim: The hidden state dimension.
         :param str method: The forward computation method. Available options
             are: recurrent, convolutional.
-        :param float dt: The time step for discretization. Default is `0.1`.
+        :param float dt_min: The minimum time step for discretization.
+            Default is `0.001`.
+        :param float dt_max: The maximum time step for discretization.
+            Default is `0.01`.
         :param dict kwargs: Additional arguments for the class constructor.
         """
+
         super().__init__(
             input_dim=input_dim,
             hid_dim=hid_dim,
             method=method,
-            dt=dt,
             hippo=False,
+            dt_min=dt_min,
+            dt_max=dt_max,
             **kwargs,
         )
 
diff --git a/ssm/model/block/s6_block.py b/ssm/model/block/s6_block.py
index c13ef97..bd4f1e1 100644
--- a/ssm/model/block/s6_block.py
+++ b/ssm/model/block/s6_block.py
@@ -1,5 +1,6 @@
+import math
 import torch
-from ...utils import compute_S4DReal
+from ...utils import compute_S4DReal, initialize_dt
 
 
 class DeltaNetwork(torch.nn.Module):
@@ -11,22 +12,28 @@ class DeltaNetwork(torch.nn.Module):
     input dimension.
     """
 
-    def __init__(self, input_dim, dt):
+    def __init__(self, input_dim, dt_min, dt_max):
         """
         Initialization of the Delta Network.
 
         :param int input_dim: The input dimension.
-        :param float dt: The time step for discretization.
+        :param float dt_min: The minimum time step for discretization.
+        :param float dt_max: The maximum time step for discretization.
         """
         super().__init__()
-
-        # Initialize the parameters
-        self.dt = dt
         self.input_dim = input_dim
-
-        # Define the layers
+        self.linear = torch.nn.Linear(input_dim, 1, bias=False)
         self.activation = torch.nn.Softplus()
-        self.linear = torch.nn.Linear(input_dim, 1)
+
+        # Initialize the time step dt
+        self.dt = torch.nn.Parameter(
+            initialize_dt(
+                input_dim=input_dim,
+                dt_min=dt_min,
+                dt_max=dt_max,
+                inverse_softplus=True,
+            )
+        )
 
     def forward(self, x):
         """
@@ -36,8 +43,9 @@ def forward(self, x):
         :return: The output tensor.
         :rtype: torch.Tensor
         """
-        delta = self.activation(self.linear(x)) + self.dt
-        return delta.repeat(1, 1, self.input_dim)
+        x = self.linear(x)
+        x = x.expand(-1, -1, self.input_dim)
+        return self.activation(x + self.dt)
 
 
 class S6Block(torch.nn.Module):
@@ -68,7 +76,8 @@ def __init__(
         self,
         input_dim,
         hid_dim,
-        dt=0.1,
+        dt_min=0.001,
+        dt_max=0.01,
         real_random=False,
         **kwargs,
     ):
@@ -77,7 +86,10 @@ def __init__(
 
         :param int input_dim: The input dimension.
         :param int hid_dim: The hidden dimension.
-        :param float dt: The time step for discretization. Default is `0.1`.
+        :param float dt_min: The minimum time step for discretization.
+            Default is `0.001`.
+        :param float dt_max: The maximum time step for discretization.
+            Default is `0.01`.
         :param bool real_random: If `True`, the real part of the A matrix is
             initialized at random between 0 and 1. Default is `False`.
         :param dict kwargs: Additional keyword arguments.
@@ -87,7 +99,6 @@ def __init__(
         # Initialize parameters
         self.input_dim = input_dim
         self.hid_dim = hid_dim
-        self.dt = dt
 
         # Initialize the matrix A
         A = compute_S4DReal(hid_dim, real_random=real_random).unsqueeze(0)
@@ -96,7 +107,9 @@ def __init__(
         # Initialize the networks to compute matrices B and C
         self.linear_b = torch.nn.Linear(input_dim, hid_dim)
         self.linear_c = torch.nn.Linear(input_dim, hid_dim)
-        self.delta_net = DeltaNetwork(input_dim=input_dim, dt=dt)
+        self.delta_net = DeltaNetwork(
+            input_dim=input_dim, dt_min=dt_min, dt_max=dt_max
+        )
 
     def _discretize(self, A, B, dt):
         """
diff --git a/ssm/model/s4.py b/ssm/model/s4.py
index e528a30..41339b7 100644
--- a/ssm/model/s4.py
+++ b/ssm/model/s4.py
@@ -1,5 +1,6 @@
 import torch
 from .block import S4BaseBlock, S4LowRankBlock, S4DBlock
+from .block.residual_block import ResidualBlock
 
 
 class S4(torch.nn.Module):
@@ -40,6 +41,8 @@ def __init__(
         n_layers=2,
         block_type="S4",
         activation=torch.nn.ReLU,
+        layer_norm=True,
+        residual=True,
         **kwargs,
     ):
         """
@@ -55,6 +58,10 @@ def __init__(
             are: `"S4"`, `"S4D"`, `"S4LowRank"`. Default is `"S4"`.
         :param torch.nn.Module activation: The activation function.
             Default is `torch.nn.ReLU`.
+        :param bool layer_norm: If `True`, layer normalization is applied after
+            each S4 block. Default is `True`.
+        :param bool residual: If `True`, a residual connection is added to the
+            output of each S4 block. Default is `True`.
         :param dict kwargs: Additional keyword arguments used in the block.
         :raises ValueError: If the specified `block_type` is not valid.
         """
@@ -76,20 +83,21 @@ def __init__(
                 "Available options are: 'S4', 'S4D', 'S4LowRank'."
             )
 
-        # Initialize the layers
         layers = []
         for _ in range(n_layers):
-            layers.append(
+            tmp = torch.nn.Sequential(
                 block_class(
                     input_dim=input_dim,
                     hid_dim=hid_dim,
                     method=method,
                     **kwargs,
-                )
+                ),
+                activation(),
+                torch.nn.Linear(input_dim, input_dim),
+                # Conditionally add layer normalization
+                *([torch.nn.LayerNorm(input_dim)] if layer_norm else []),
             )
-            layers.append(activation())
-            layers.append(torch.nn.Linear(input_dim, input_dim))
-            layers.append(torch.nn.LayerNorm(input_dim))
+            layers.append(tmp if not residual else ResidualBlock(tmp))
         self.layers = torch.nn.Sequential(*layers)
 
         # Initialize the decoder to match the output dimension
diff --git a/ssm/model/s6.py b/ssm/model/s6.py
index db947a0..ade4103 100644
--- a/ssm/model/s6.py
+++ b/ssm/model/s6.py
@@ -1,5 +1,6 @@
 import torch
 from .block import S6Block
+from .block.residual_block import ResidualBlock
 
 
 class S6(torch.nn.Module):
@@ -34,8 +35,10 @@ def __init__(
         hid_dim,
         output_dim,
         n_layers=2,
-        activation=torch.nn.ReLU,
+        activation=torch.nn.GELU,
         real_random=False,
+        residual=True,
+        layer_norm=True,
         **kwargs,
     ):
         """
@@ -46,28 +49,34 @@ def __init__(
         :param int output_dim: The output dimension.
         :param int n_layers: Number of S6 blocks. Default is 2.
         :param torch.nn.Module activation: The activation function.
-            Default is `torch.nn.ReLU`.
+            Default is `torch.nn.GELU`.
         :param bool real_random: If `True`, the real part of the A matrix of the
             diagonal block is initialized at random between 0 and 1.
             Default is `False`.
+        :param bool residual: If `True`, a residual connection is added to the
+            output of each S6 block. Default is `True`.
+        :param bool layer_norm: If `True`, layer normalization is applied after
+            each S6 block. Default is `True`.
         :param dict kwargs: Additional keyword arguments used in the block.
         """
         super().__init__()
         self.input_dim = input_dim
+
         # Initialize the layers
         layers = []
         for _ in range(n_layers):
-            layers.append(
+            tmp = torch.nn.Sequential(
                 S6Block(
                     input_dim=input_dim,
                     hid_dim=hid_dim,
                     real_random=real_random,
                     **kwargs,
-                )
+                ),
+                activation(),
+                torch.nn.Linear(input_dim, input_dim),
+                *([torch.nn.LayerNorm(input_dim)] if layer_norm else []),
             )
-            layers.append(activation())
-            layers.append(torch.nn.Linear(input_dim, input_dim))
-            layers.append(torch.nn.LayerNorm(input_dim))
+            layers.append(tmp if not residual else ResidualBlock(tmp))
         self.layers = torch.nn.Sequential(*layers)
 
         # Initialize the decoder to match the output dimension
diff --git a/ssm/utils.py b/ssm/utils.py
index fc8759e..ba9790b 100644
--- a/ssm/utils.py
+++ b/ssm/utils.py
@@ -103,7 +103,7 @@ def compute_S4DReal(N, real_random=False):
     if real_random:
         real_part = torch.rand(N, dtype=torch.float32)
     else:
-        real_part = torch.arange(1, N + 1, dtype=torch.float32)
+        real_part = torch.arange(N, dtype=torch.float32)
     return -(real_part + 1)
 
 
@@ -131,3 +131,21 @@ def compute_dplr(A):
     p = Vc @ p.to(Vc.dtype)
     q = Vc @ q.to(Vc.dtype)
     return Lambda, p, q
+
+
+def initialize_dt(input_dim, dt_min, dt_max, inverse_softplus=False):
+    """
+    Initialize the time step dt for the S4 and S6 blocks.
+
+    :param float dt_min: The minimum time step for discretization.
+    :param float dt_max: The maximum time step for discretization.
+    :return: Initialized time step dt tensor of shape (input_dim,).
+    :rtype: torch.Tensor
+    """
+    # Sample dt from a uniform distribution on [dt_min, dt_max]
+    dt = torch.rand(input_dim) * (dt_max - dt_min) + dt_min
+
+    # Apply the inverse softplus to dt
+    dt = torch.log(torch.exp(dt) - 1.0 + 1e-6) if inverse_softplus else dt
+
+    return dt
diff --git a/tests/test_cli.py b/tests/test_cli.py
index f28c95a..4d272d1 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -29,8 +29,8 @@ def test_model():
     model = cli.trainer.model.model
 
     assert model.input_dim == 5
-    assert model.layers[0].hid_dim == 16
-    assert model.layers[0].method == "convolutional"
+    assert model.layers[0].model[0].hid_dim == 16
+    assert model.layers[0].model[0].method == "convolutional"
 
     shutil.rmtree("tests/logs/testing_logs")
 
diff --git a/tests/test_s6.py b/tests/test_s6.py
index 3e9ef55..80f889b 100644
--- a/tests/test_s6.py
+++ b/tests/test_s6.py
@@ -1,3 +1,4 @@
+import pytest
 import torch
 from ssm.model import S6
 
@@ -6,7 +7,9 @@
 out_dim = 2
 
 
-def test_s6_constructor():
+@pytest.mark.parametrize("residual", [True, False])
+@pytest.mark.parametrize("layer_norm", [True, False])
+def test_s6_constructor(residual, layer_norm):
 
     model = S6(
         input_dim=x.shape[2],
@@ -14,10 +17,14 @@ def test_s6_constructor():
         output_dim=out_dim,
         n_layers=3,
         activation=torch.nn.ReLU,
+        residual=residual,
+        layer_norm=layer_norm,
     )
 
 
-def test_s6_forward():
+@pytest.mark.parametrize("residual", [True, False])
+@pytest.mark.parametrize("layer_norm", [True, False])
+def test_s6_forward(residual, layer_norm):
 
     model = S6(
         input_dim=x.shape[2],
@@ -25,13 +32,17 @@ def test_s6_forward():
         output_dim=out_dim,
         n_layers=3,
         activation=torch.nn.ReLU,
+        residual=residual,
+        layer_norm=layer_norm,
     )
 
     y = model.forward(x)
     assert y.shape == (x.shape[0], x.shape[1], out_dim)
 
 
-def test_s6_backward():
+@pytest.mark.parametrize("residual", [True, False])
+@pytest.mark.parametrize("layer_norm", [True, False])
+def test_s6_backward(residual, layer_norm):
 
     model = S6(
         input_dim=x.shape[2],
@@ -39,6 +50,8 @@ def test_s6_backward():
         output_dim=out_dim,
         n_layers=3,
         activation=torch.nn.ReLU,
+        residual=residual,
+        layer_norm=layer_norm,
     )
 
     y = model.forward(x.requires_grad_())


Reviews:
user: login: GiovanniCanali
body: Hi @FilippoOlivo, thank you for the PR. 
I believe that residual block can be added to the models in a much simpler way: just add residual connections in the `forward` of interested models / blocks.
Be careful when using automatic substitutions.
state: CHANGES_REQUESTED
author_association: COLLABORATOR
submitted_at: 2025-04-17T08:54:44Z
user: login: FilippoOlivo
state: COMMENTED
author_association: OWNER
submitted_at: 2025-04-17T08:56:50Z
user: login: FilippoOlivo
state: COMMENTED
author_association: OWNER
submitted_at: 2025-04-17T09:00:07Z
user: login: copilot-pull-request-reviewer[bot]
body: 

Copilot reviewed 5 out of 5 changed files in this pull request and generated no comments.


<details>
<summary>Comments suppressed due to low confidence (2)</summary>

**ssm/utils.py:106**
* Verify that changing the starting index from 1 to 0 for torch.arange is intentional, as this adjustment shifts the resulting values and may affect downstream computations.
```
real_part = torch.arange(N, dtype=torch.float32)
```
**ssm/model/block/residual_block.py:21**
* The docstring parameters in ResidualBlock.__init__ do not match the actual function signature. Update the docstring to accurately describe the expected parameters.
```
        :param int input_dim: The input dimension.
```
</details>


state: COMMENTED
author_association: NONE
submitted_at: 2025-04-17T10:04:06Z
user: login: FilippoOlivo
state: COMMENTED
author_association: OWNER
submitted_at: 2025-04-17T11:46:45Z
user: login: GiovanniCanali
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-17T13:31:53Z
user: login: GiovanniCanali
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-17T13:33:58Z
user: login: FilippoOlivo
state: COMMENTED
author_association: OWNER
submitted_at: 2025-04-17T13:41:47Z
user: login: GiovanniCanali
body: Good to go!
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-18T09:36:12Z

Review comments:
diff_hunk: @@ -0,0 +1,39 @@
+import torch
+
+
+class ResidualBlock(torch.nn.Module):
path: ssm/model/block/residual_block.py
user: login: GiovanniCanali
body: Isn't it better to write something like `x + prediction` in the forward of the interested models, instead of creating a new class?
created_at: 2025-04-17T08:39:51Z
author_association: COLLABORATOR
diff_hunk: @@ -11,22 +12,26 @@ class DeltaNetwork(torch.nn.Module):
     input dimension.
     """
 
-    def __init__(self, input_dim, dt):
+    def __init__(self, input_dim, dt_min, dt_max, dt_init_floor):
         """
         Initialization of the Delta Network.
 
         :param int input_dim: The input dimension.
-        :param float dt: The time step for discretization.
+        :param float dt_min: The minimum time step.
+        :param float dt_max: The maximum time step.
+        :param float dt_init_floor: The minimum value for the time step.
         """
         super().__init__()
-
-        # Initialize the parameters
-        self.dt = dt
         self.input_dim = input_dim
-
-        # Define the layers
+        self.linear = torch.nn.Linear(input_dim, 1, bias=False)
         self.activation = torch.nn.Softplus()
-        self.linear = torch.nn.Linear(input_dim, 1)
+
+        dt = torch.exp(
path: ssm/model/block/s6_block.py
user: login: GiovanniCanali
body: Is there any reference for this change?
created_at: 2025-04-17T08:42:24Z
author_association: COLLABORATOR
diff_hunk: @@ -36,11 +41,13 @@ def forward(self, x):
         :return: The output tensor.
         :rtype: torch.Tensor
         """
-        delta = self.activation(self.linear(x)) + self.dt
-        return delta.repeat(1, 1, self.input_dim)
+
+        x = self.linear(x)
+        x = x.expand(-1, -1, self.input_dim)
+        return self.activation(x + self.bias)
 
 
-class S6Block(torch.nn.Module):
+class S6Block(torch.torch.nn.Module):
path: ssm/model/block/s6_block.py
user: login: GiovanniCanali
body: Duplicated `torch`
created_at: 2025-04-17T08:42:40Z
author_association: COLLABORATOR
diff_hunk: @@ -49,7 +56,7 @@ class S6Block(torch.nn.Module):
     parts of the input sequence, making it suitable for tasks such as selective
     copy.
 
-    The output is computed in an efficient manner by leveraging the parallel
+    The output is computed in an efficient matorch.nner by leveraging the parallel
path: ssm/model/block/s6_block.py
user: login: GiovanniCanali
body: Please, restore `manner`
created_at: 2025-04-17T08:42:53Z
author_association: COLLABORATOR
diff_hunk: @@ -87,16 +99,20 @@ def __init__(
         # Initialize parameters
         self.input_dim = input_dim
         self.hid_dim = hid_dim
-        self.dt = dt
 
         # Initialize the matrix A
         A = compute_S4DReal(hid_dim, real_random=real_random).unsqueeze(0)
-        self.A = torch.nn.Parameter(A.repeat(input_dim, 1))
+        self.A = torch.torch.nn.Parameter(A.repeat(input_dim, 1))
path: ssm/model/block/s6_block.py
user: login: GiovanniCanali
body: Duplicated `torch`
created_at: 2025-04-17T08:43:30Z
author_association: COLLABORATOR
diff_hunk: @@ -87,16 +99,20 @@ def __init__(
         # Initialize parameters
         self.input_dim = input_dim
         self.hid_dim = hid_dim
-        self.dt = dt
 
         # Initialize the matrix A
         A = compute_S4DReal(hid_dim, real_random=real_random).unsqueeze(0)
-        self.A = torch.nn.Parameter(A.repeat(input_dim, 1))
+        self.A = torch.torch.nn.Parameter(A.repeat(input_dim, 1))
 
         # Initialize the networks to compute matrices B and C
-        self.linear_b = torch.nn.Linear(input_dim, hid_dim)
-        self.linear_c = torch.nn.Linear(input_dim, hid_dim)
-        self.delta_net = DeltaNetwork(input_dim=input_dim, dt=dt)
+        self.linear_b = torch.torch.nn.Linear(input_dim, hid_dim)
+        self.linear_c = torch.torch.nn.Linear(input_dim, hid_dim)
path: ssm/model/block/s6_block.py
user: login: GiovanniCanali
body: Duplicated `torch` in both lines
created_at: 2025-04-17T08:44:36Z
author_association: COLLABORATOR
diff_hunk: @@ -177,7 +193,7 @@ def _complex_log(input_):
         a_star = torch.cumsum(log_a, dim=1)
 
         # Padding over the sequence length L
-        a_star = torch.nn.functional.pad(a_star, (0, 0, 0, 0, 1, 0))
+        a_star = torch.torch.nn.functional.pad(a_star, (0, 0, 0, 0, 1, 0))
path: ssm/model/block/s6_block.py
user: login: GiovanniCanali
body: Duplicated `torch`
created_at: 2025-04-17T08:44:52Z
author_association: COLLABORATOR
diff_hunk: @@ -11,22 +12,26 @@ class DeltaNetwork(torch.nn.Module):
     input dimension.
     """
 
-    def __init__(self, input_dim, dt):
+    def __init__(self, input_dim, dt_min, dt_max, dt_init_floor):
         """
         Initialization of the Delta Network.
 
         :param int input_dim: The input dimension.
-        :param float dt: The time step for discretization.
+        :param float dt_min: The minimum time step.
+        :param float dt_max: The maximum time step.
+        :param float dt_init_floor: The minimum value for the time step.
         """
         super().__init__()
-
-        # Initialize the parameters
-        self.dt = dt
         self.input_dim = input_dim
-
-        # Define the layers
+        self.linear = torch.nn.Linear(input_dim, 1, bias=False)
         self.activation = torch.nn.Softplus()
-        self.linear = torch.nn.Linear(input_dim, 1)
+
+        dt = torch.exp(
path: ssm/model/block/s6_block.py
user: login: FilippoOlivo
body: Nope, just they use this in they GitHub repo!
created_at: 2025-04-17T08:56:50Z
author_association: OWNER
diff_hunk: @@ -0,0 +1,39 @@
+import torch
+
+
+class ResidualBlock(torch.nn.Module):
path: ssm/model/block/residual_block.py
user: login: FilippoOlivo
body: I believe this solution is better due the fact that we use the residual on a sequence of modules 
SSM + mixing + activation (+ Layer Normalization). In this way adding the Residual is just a wrapping of a torch sequential. In the other case we need to break down the torch Sequential module in a torch.nn.Module list and call and iterate over the list of modules (and then adding the residual)
created_at: 2025-04-17T09:00:07Z
author_association: OWNER
diff_hunk: @@ -0,0 +1,39 @@
+import torch
+
+
+class ResidualBlock(torch.nn.Module):
path: ssm/model/block/residual_block.py
user: login: FilippoOlivo
body: Maybe another solution could be creating a Block which handles residual, activation and LayerNormalization instead of only handling residual. This makes the code flexible without adding an "trivial" Module. @GiovanniCanali what do you think? 
created_at: 2025-04-17T11:46:45Z
author_association: OWNER
diff_hunk: @@ -0,0 +1,39 @@
+import torch
+
+
+class ResidualBlock(torch.nn.Module):
path: ssm/model/block/residual_block.py
user: login: GiovanniCanali
body: I understand your point. Let's keep it the way you thought it in the first place.
created_at: 2025-04-17T13:31:53Z
author_association: COLLABORATOR
diff_hunk: @@ -11,22 +12,26 @@ class DeltaNetwork(torch.nn.Module):
     input dimension.
     """
 
-    def __init__(self, input_dim, dt):
+    def __init__(self, input_dim, dt_min, dt_max, dt_init_floor):
         """
         Initialization of the Delta Network.
 
         :param int input_dim: The input dimension.
-        :param float dt: The time step for discretization.
+        :param float dt_min: The minimum time step.
+        :param float dt_max: The maximum time step.
+        :param float dt_init_floor: The minimum value for the time step.
         """
         super().__init__()
-
-        # Initialize the parameters
-        self.dt = dt
         self.input_dim = input_dim
-
-        # Define the layers
+        self.linear = torch.nn.Linear(input_dim, 1, bias=False)
         self.activation = torch.nn.Softplus()
-        self.linear = torch.nn.Linear(input_dim, 1)
+
+        dt = torch.exp(
path: ssm/model/block/s6_block.py
user: login: GiovanniCanali
body: Ok, then I am not sure this is how we are supposed to proceed. It's up to you, we can go back to the previous implementation at any time in the future
created_at: 2025-04-17T13:33:58Z
author_association: COLLABORATOR
diff_hunk: @@ -11,22 +12,26 @@ class DeltaNetwork(torch.nn.Module):
     input dimension.
     """
 
-    def __init__(self, input_dim, dt):
+    def __init__(self, input_dim, dt_min, dt_max, dt_init_floor):
         """
         Initialization of the Delta Network.
 
         :param int input_dim: The input dimension.
-        :param float dt: The time step for discretization.
+        :param float dt_min: The minimum time step.
+        :param float dt_max: The maximum time step.
+        :param float dt_init_floor: The minimum value for the time step.
         """
         super().__init__()
-
-        # Initialize the parameters
-        self.dt = dt
         self.input_dim = input_dim
-
-        # Define the layers
+        self.linear = torch.nn.Linear(input_dim, 1, bias=False)
         self.activation = torch.nn.Softplus()
-        self.linear = torch.nn.Linear(input_dim, 1)
+
+        dt = torch.exp(
path: ssm/model/block/s6_block.py
user: login: FilippoOlivo
body: I read more carefully mamba paper and, at page 9 they say 
_In our experiments, the Δ parameter (which can be viewed as a bias term) is initialized to 𝜏−1
Δ (Uniform([0.001,0.1])), following prior work on SSMs (Gu, Johnson, Timalsina, et al. 2023)._
Do we want to apply this precise strategy? Tau is soft plus in this context 

created_at: 2025-04-17T13:41:47Z
author_association: OWNER

PR comments:
user: login: FilippoOlivo
created_at: 2025-04-17T09:01:22Z
body: @GiovanniCanali Thanks for the suggestion! I didn't notice the duplications, I will fix them soon. They I suggest to apply the same logic to S4. What do you think? 

author_association: OWNER
user: login: GiovanniCanali
created_at: 2025-04-17T13:37:43Z
body: > @GiovanniCanali Thanks for the suggestion! I didn't notice the duplications, I will fix them soon. They I suggest to apply the same logic to S4. What do you think?

If you refer to the handling of residual connections, it is ok for me. In that case, please use this very same PR for changes.
If you meant something else, please clarify.
author_association: COLLABORATOR
user: login: FilippoOlivo
created_at: 2025-04-18T07:35:04Z
body: @GiovanniCanali I modify the initialization of dt for all the S4 and S6 blocks and I add residual connection in S4. Ready to merge from my side 
author_association: OWNER