diff --git a/models/S4.py b/models/S4.py
new file mode 100644
index 0000000..b693190
--- /dev/null
+++ b/models/S4.py
@@ -0,0 +1,304 @@
+"""
+Implementation of the S4 model taken from https://github.com/state-spaces/s4
+"""
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange, repeat
+
+
+class DropoutNd(nn.Module):
+    def __init__(self, p: float = 0.5, tie=True, transposed=True):
+        """
+        tie: tie dropout mask across sequence lengths (Dropout1d/2d/3d)
+        """
+        super().__init__()
+        if p < 0 or p >= 1:
+            raise ValueError(
+                "dropout probability has to be in [0, 1), " "but got {}".format(p)
+            )
+        self.p = p
+        self.tie = tie
+        self.transposed = transposed
+        self.binomial = torch.distributions.binomial.Binomial(probs=1 - self.p)
+
+    def forward(self, X):
+        """X: (batch, dim, lengths...)."""
+        if self.training:
+            if not self.transposed:
+                X = rearrange(X, "b ... d -> b d ...")
+            # binomial = torch.distributions.binomial.Binomial(probs=1-self.p) # This is incredibly slow because of CPU -> GPU copying
+            mask_shape = X.shape[:2] + (1,) * (X.ndim - 2) if self.tie else X.shape
+            # mask = self.binomial.sample(mask_shape)
+            mask = torch.rand(*mask_shape, device=X.device) < 1.0 - self.p
+            X = X * mask * (1.0 / (1 - self.p))
+            if not self.transposed:
+                X = rearrange(X, "b d ... -> b ... d")
+            return X
+        return X
+
+
+class S4DKernel(nn.Module):
+    """Generate convolution kernel from diagonal SSM parameters."""
+
+    def __init__(self, d_model, N=64, dt_min=0.001, dt_max=0.1, lr=None):
+        super().__init__()
+        # Generate dt
+        H = d_model
+        log_dt = torch.rand(H) * (math.log(dt_max) - math.log(dt_min)) + math.log(
+            dt_min
+        )
+
+        C = torch.randn(H, N // 2, dtype=torch.cfloat)
+        self.C = nn.Parameter(torch.view_as_real(C))
+        self.register("log_dt", log_dt, lr)
+
+        log_A_real = torch.log(0.5 * torch.ones(H, N // 2))
+        A_imag = math.pi * repeat(torch.arange(N // 2), "n -> h n", h=H)
+        self.register("log_A_real", log_A_real, lr)
+        self.register("A_imag", A_imag, lr)
+
+    def forward(self, L):
+        """
+        returns: (..., c, L) where c is number of channels (default 1)
+        """
+
+        # Materialize parameters
+        dt = torch.exp(self.log_dt)  # (H)
+        C = torch.view_as_complex(self.C)  # (H N)
+        A = -torch.exp(self.log_A_real) + 1j * self.A_imag  # (H N)
+
+        # Vandermonde multiplication
+        dtA = A * dt.unsqueeze(-1)  # (H N)
+        K = dtA.unsqueeze(-1) * torch.arange(L, device=A.device)  # (H N L)
+        C = C * (torch.exp(dtA) - 1.0) / A
+        K = 2 * torch.einsum("hn, hnl -> hl", C, torch.exp(K)).real
+
+        return K
+
+    def register(self, name, tensor, lr=None):
+        """Register a tensor with a configurable learning rate and 0 weight decay"""
+
+        if lr == 0.0:
+            self.register_buffer(name, tensor)
+        else:
+            self.register_parameter(name, nn.Parameter(tensor))
+
+            optim = {"weight_decay": 0.0}
+            if lr is not None:
+                optim["lr"] = lr
+            setattr(getattr(self, name), "_optim", optim)
+
+
+class S4D(nn.Module):
+    def __init__(
+        self, d_model, d_state=64, dropout=0.0, transposed=True, **kernel_args
+    ):
+        super().__init__()
+
+        self.h = d_model
+        self.n = d_state
+        self.d_output = self.h
+        self.transposed = transposed
+
+        self.D = nn.Parameter(torch.randn(self.h))
+
+        # SSM Kernel
+        self.kernel = S4DKernel(self.h, N=self.n, **kernel_args)
+
+        # Pointwise
+        self.activation = nn.GELU()
+        # dropout_fn = nn.Dropout2d # NOTE: bugged in PyTorch 1.11
+        dropout_fn = DropoutNd
+        self.dropout = dropout_fn(dropout) if dropout > 0.0 else nn.Identity()
+
+        # position-wise output transform to mix features
+        self.output_linear = nn.Sequential(
+            nn.Conv1d(self.h, 2 * self.h, kernel_size=1),
+            nn.GLU(dim=-2),
+        )
+
+    def forward(self, u, **kwargs):  # absorbs return_output and transformer src mask
+        """Input and output shape (B, H, L)"""
+        if not self.transposed:
+            u = u.transpose(-1, -2)
+        L = u.size(-1)
+
+        # Compute SSM Kernel
+        k = self.kernel(L=L)  # (H L)
+
+        # Convolution
+        k_f = torch.fft.rfft(k, n=2 * L)  # (H L)
+        u_f = torch.fft.rfft(u, n=2 * L)  # (B H L)
+        y = torch.fft.irfft(u_f * k_f, n=2 * L)[..., :L]  # (B H L)
+
+        # Compute D term in state space equation - essentially a skip connection
+        y = y + u * self.D.unsqueeze(-1)
+
+        y = self.dropout(self.activation(y))
+        y = self.output_linear(y)
+        if not self.transposed:
+            y = y.transpose(-1, -2)
+        return y
+
+
+class S4Block(nn.Module):
+    """
+    A single S4 block that applies:
+      1. S4D module
+      2. (Optionally) a linear layer + GLU activation,
+      3. Residual connection
+      4. Layer Normalization
+      5. Dropout
+
+    Args:
+        model_dim (int): Dimensionality of the model (d_model).
+        dropout_rate (float): Probability of an element to be zeroed in Dropout.
+        use_glu (bool): Whether to apply a Linear -> GLU stage after the residual.
+    """
+
+    def __init__(
+        self, model_dim: int, dropout_rate: float = 0.1, use_glu: bool = False
+    ):
+        super().__init__()
+        self.s4d = S4D(d_model=model_dim, transposed=False)
+        self.norm = nn.LayerNorm(model_dim)
+        self.drop = nn.Dropout(p=dropout_rate)
+
+        self.use_glu = use_glu
+        if self.use_glu:
+            self.post_linear = nn.Linear(model_dim, 2 * model_dim)
+        else:
+            self.post_linear = None
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass of the S4Block.
+
+        Args:
+            x (torch.Tensor): Input tensor of shape (batch_size, seq_len, model_dim).
+
+        Returns:
+            torch.Tensor: Output tensor of the same shape (batch_size, seq_len, model_dim).
+        """
+
+        # S4 module
+        y = self.s4d(x)
+        y = y + x
+
+        # Optional: Linear -> GLU
+        if self.use_glu:
+            # shape: (batch_size, seq_len, 2 * model_dim)
+            y_glu = self.post_linear(y)
+            # shape: (batch_size, seq_len, model_dim)
+            y_glu = F.glu(y_glu, dim=-1)
+            y = y + y_glu
+
+        # Layer Normalization
+        y = self.norm(y)
+
+        # Dropout
+        y = self.drop(y)
+
+        return y
+
+
+class StackedS4(nn.Module):
+    """
+    A stack of multiple S4Blocks, preceded by an embedding layer
+    and followed by a linear projection.
+
+    Args:
+        num_blocks (int): Number of S4Blocks to stack.
+        model_dim (int): Dimensionality of embeddings and S4 blocks.
+        data_dim (int): Size of the vocabulary (if input is token IDs).
+        label_dim (int): Output dimensionality (e.g., number of classes).
+        dropout_rate (float): Dropout probability for each S4Block.
+        use_glu (bool): If True, each block will include a Linear->GLU stage
+                        that preserves model_dim.
+        second_embedding (bool): If True, the model will expect two input
+                                    token IDs and use two separate embeddings.
+    """
+
+    def __init__(
+        self,
+        num_blocks: int,
+        model_dim: int,
+        data_dim: int,
+        label_dim: int,
+        dropout_rate: float = 0.1,
+        use_glu: bool = False,
+        second_embedding: bool = False,
+    ):
+        super().__init__()
+
+        self.second_embedding = second_embedding
+        embedding_dim = model_dim // 2 if second_embedding else model_dim
+        self.embedding = nn.Embedding(data_dim, embedding_dim)
+        if second_embedding:
+            self.embedding2 = nn.Embedding(data_dim, embedding_dim)
+
+        # Create multiple S4Blocks
+        self.blocks = nn.ModuleList(
+            [
+                S4Block(model_dim=model_dim, dropout_rate=dropout_rate, use_glu=use_glu)
+                for _ in range(num_blocks)
+            ]
+        )
+
+        # The final linear projection remains (model_dim -> label_dim)
+        self.linear = nn.Linear(model_dim, label_dim)
+
+    def mask_grads(self):
+        """
+        This method is included for consistency with other models.
+        """
+        pass
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass of StackedS4.
+
+        Args:
+            x (torch.Tensor): Input tensor of shape (batch_size, seq_len)
+                              containing integer token IDs (if used with nn.Embedding).
+
+        Returns:
+            torch.Tensor: Output tensor of shape (batch_size, seq_len, label_dim).
+                          If a single-vector output is desired (e.g. for classification),
+                          additional pooling or indexing may be required
+                          before the final linear layer or after its output.
+        """
+        # Embedding: (batch_size, seq_len, model_dim)
+        if not self.second_embedding:
+            x = self.embedding(x)
+        else:
+            x = torch.cat(
+                [self.embedding(x[:, :, 0]), self.embedding2(x[:, :, 1])], dim=-1
+            )
+
+        # Pass through each S4Block
+        for block in self.blocks:
+            x = block(x)
+
+        # Final projection: (batch_size, seq_len, label_dim)
+        return self.linear(x)
+
+    def step(self, x: torch.Tensor) -> torch.Tensor:
+        # Embedding for the current step
+        if self.second_embedding:
+            x = torch.cat(
+                [self.embedding(x[:, 0].long()), self.embedding2(x[:, 1].long())],
+                dim=-1,
+            )
+        else:
+            x = self.embedding(x)
+
+        for block in self.blocks:
+            x = block.step(x)
+
+        # Final projection for the step
+        return self.linear(x)
diff --git a/models/s6.py b/models/s6.py
deleted file mode 100644
index d9a2275..0000000
--- a/models/s6.py
+++ /dev/null
@@ -1,209 +0,0 @@
-import math
-
-import torch
-import torch.nn as nn
-from einops import rearrange, repeat
-from mamba_ssm.ops.selective_scan_interface import selective_scan_fn
-
-
-class MambaRecurrence(nn.Module):
-    """
-    Implements the Mamba recurrence layer for sequence modeling.
-
-    Args:
-        d_model (int): Dimension of the model (number of features).
-        d_state (int): Dimension of the state space. Defaults to 16.
-        dt_rank (int or str): Rank for the time-step parameterization. Defaults to
-            "auto".
-        dt_min (float): Minimum value for time-steps. Defaults to 0.001.
-        dt_max (float): Maximum value for time-steps. Defaults to 0.1.
-        dt_init (str): Initialization method for dt. Options are "constant" or "random".
-            Defaults to "random".
-        dt_scale (float): Scale factor for dt initialization. Defaults to 1.0.
-        dt_init_floor (float): Floor value for initializing time-steps. Defaults to
-            1e-4.
-        device (torch.device, optional): Device to run the computations on. If None,
-            uses CUDA if available.
-
-    Forward Args:
-        hidden_states (Tensor): Input tensor of shape (batch_size, sequence_length,
-            d_model).
-
-    Returns:
-        Tensor: Output tensor of the same shape as input.
-    """
-
-    def __init__(
-        self,
-        d_model,
-        d_state=16,
-        dt_rank="auto",
-        dt_min=0.001,
-        dt_max=0.1,
-        dt_init="random",
-        dt_scale=1.0,
-        dt_init_floor=1e-4,
-        device=None,
-    ):
-        super().__init__()
-
-        if device is None:
-            if torch.cuda.is_available():
-                self.device = torch.device("cuda")
-            else:
-                self.device = torch.device("cpu")
-        else:
-            self.device = device
-        self.d_model = d_model
-        self.d_state = d_state
-        self.dt_rank = math.ceil(self.d_model / 16) if dt_rank == "auto" else dt_rank
-
-        self.x_proj = nn.Linear(
-            self.d_model,
-            self.dt_rank + self.d_state * 2,
-            bias=False,
-            device=self.device,
-        )
-        self.dt_proj = nn.Linear(
-            self.dt_rank, self.d_model, bias=True, device=self.device
-        )
-
-        # Initialize special dt projection to preserve variance at initialization
-        dt_init_std = self.dt_rank**-0.5 * dt_scale
-        if dt_init == "constant":
-            nn.init.constant_(self.dt_proj.weight, dt_init_std)
-        elif dt_init == "random":
-            nn.init.uniform_(self.dt_proj.weight, -dt_init_std, dt_init_std)
-        else:
-            raise NotImplementedError
-
-        # Initialize dt bias so that F.softplus(dt_bias) is between dt_min and dt_max
-        dt = torch.exp(
-            torch.rand(self.d_model, device=self.device)
-            * (math.log(dt_max) - math.log(dt_min))
-            + math.log(dt_min)
-        ).clamp(min=dt_init_floor)
-        # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759
-        inv_dt = dt + torch.log(-torch.expm1(-dt))
-        with torch.no_grad():
-            self.dt_proj.bias.copy_(inv_dt)
-        # Our initialization would set all Linear.bias to zero, need to mark this
-        # one as _no_reinit
-        self.dt_proj.bias._no_reinit = True
-
-        # S4D real initialization
-        A = repeat(
-            torch.arange(1, self.d_state + 1, dtype=torch.float32, device=self.device),
-            "n -> d n",
-            d=self.d_model,
-        ).contiguous()
-        A_log = torch.log(A)  # Keep A_log in fp32
-        self.A_log = nn.Parameter(A_log)
-
-        # D "skip" parameter
-        self.D = nn.Parameter(
-            torch.ones(self.d_model, device=self.device)
-        )  # Keep in fp32
-        self.D._no_weight_decay = True
-
-    def forward(self, hidden_states):
-        """
-        hidden_states: (B, L, D)
-        Returns: same shape as hidden_states
-        """
-        batch, seqlen, dim = hidden_states.shape
-        A = -torch.exp(self.A_log.float())  # (d_inner, d_state)
-        x = rearrange(hidden_states, "b l d -> b d l")
-        x_dbl = self.x_proj(rearrange(x, "b d l -> (b l) d"))  # (bl d)
-        dt, B, C = torch.split(
-            x_dbl, [self.dt_rank, self.d_state, self.d_state], dim=-1
-        )
-        dt = self.dt_proj.weight @ dt.t()
-        dt = rearrange(dt, "d (b l) -> b d l", l=seqlen)
-        B = rearrange(B, "(b l) dstate -> b dstate l", l=seqlen).contiguous()
-        C = rearrange(C, "(b l) dstate -> b dstate l", l=seqlen).contiguous()
-        y = selective_scan_fn(
-            x,
-            dt,
-            A,
-            B,
-            C,
-            self.D.float(),
-            z=None,
-            delta_bias=self.dt_proj.bias.float(),
-            delta_softplus=True,
-            return_last_state=False,
-        )
-        return rearrange(y, "b d l -> b l d")
-
-
-class S6(nn.Module):
-    def __init__(
-        self,
-        num_blocks: int,
-        data_dim: int,
-        model_dim: int,
-        label_dim: int,
-        dropout_rate: float = 0.1,
-        second_embedding: bool = False,
-        use_glu: bool = False,
-        d_state: int = 16,
-        dt_rank: str | int = "auto",
-    ):
-        """
-        d_state: state dimension of the SSM (default 16).
-        dt_rank: rank for dt parameterization ("auto" uses ceil(model_dim/16)).
-        """
-        super().__init__()
-        self.second_embedding = second_embedding
-
-        emb_dim = model_dim // 2 if second_embedding else model_dim
-        self.embedding = nn.Embedding(data_dim, emb_dim)
-        if second_embedding:
-            self.embedding2 = nn.Embedding(data_dim, emb_dim)
-
-        self.blocks = nn.ModuleList(
-            [
-                MambaRecurrence(model_dim, d_state=d_state, dt_rank=dt_rank)
-                for _ in range(num_blocks)
-            ]
-        )
-        self.norms = nn.ModuleList([nn.LayerNorm(model_dim) for _ in range(num_blocks)])
-
-        self.dropout = nn.Dropout(dropout_rate)
-        self.linear = nn.Linear(model_dim, label_dim)
-
-        self.use_glu = use_glu
-        if use_glu:
-            self.glu_projs = nn.ModuleList(
-                [nn.Linear(model_dim, 2 * model_dim) for _ in range(num_blocks)]
-            )
-        else:
-            self.glu_projs = nn.ModuleList(
-                [nn.Linear(model_dim, model_dim) for _ in range(num_blocks)]
-            )
-        self.act = nn.GLU()
-
-    def mask_grads(self):
-        pass
-
-    def _embed(self, x: torch.Tensor) -> torch.Tensor:
-        if not self.second_embedding:
-            # x: (B, L)
-            return self.embedding(x)  # -> (B, L, model_dim)
-        else:
-            # x: (B, L, 2)
-            return torch.cat(
-                [self.embedding(x[:, :, 0]), self.embedding2(x[:, :, 1])], dim=-1
-            )  # -> (B, L, model_dim)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = self._embed(x)  # (B, L, D=model_dim)
-        for i, (blk, ln) in enumerate(zip(self.blocks, self.norms)):
-            residual = x
-            x = blk(x)
-            h = self.glu_projs[i](x)
-            h = self.act(h) if self.use_glu else torch.tanh(h)
-            x = residual + x + h
-            x = self.dropout(ln(x))
-        return self.linear(x)  # (B, L, label_dim)
diff --git a/train.py b/train.py
index 3561e50..58ec0c5 100644
--- a/train.py
+++ b/train.py
@@ -303,8 +303,6 @@ def run_experiment(config):
     slstm_at = config.get("slstm_at", [1])
     vf_A_norm_lambda = config.get("vf_A_norm_lambda", 0.001)
     rank = config.get("rank", 0)
-    d_state = config.get("d_state", 16)
-    dt_rank = config.get("dt_rank", "auto")
 
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
@@ -405,10 +403,10 @@ def train_dataloader_multilength():
             use_glu=use_glu,
             second_embedding=second_embedding,
         )
-    elif model_name == "S6":
-        from models.s6 import S6
+    elif model_name == "S4":
+        from models.S4 import StackedS4
 
-        model = S6(
+        model = StackedS4(
             num_blocks=num_blocks,
             model_dim=model_dim,
             data_dim=data_dim,
@@ -416,10 +414,7 @@ def train_dataloader_multilength():
             dropout_rate=dropout_rate,
             use_glu=use_glu,
             second_embedding=second_embedding,
-            d_state=d_state,
-            dt_rank=dt_rank,
         )
-
     elif model_name in ["deltanet", "gateddeltanet", "rwkv7", "rwkv6", "deltaproduct"]:
         from models.fla import StackedBlock