Skip to content

[BUG] MambaTab forward never calls its Mamba block #412

Description

@ChrisW09

Describe the bug
MambaTab.forward never calls the Mamba block. self.mamba is constructed in __init__ (deeptab/architectures/mambatab.py:84-87), but the forward pass is initial_layer → unsqueeze → norm_f → embedding_activation → squeeze → tabular_head (mambatab.py:89-120) — self.mamba is never invoked.

Every MambaTabClassifier/Regressor/LSS fit therefore trains a single linear layer + LayerNorm + activation + MLP head. All Mamba parameters sit in the optimizer with zero gradients, wasting memory and compute while the model silently underperforms what it claims to be. (Only BaseModel.encode() for pretraining ever touches self.mamba.)

To Reproduce

import numpy as np, pandas as pd
from deeptab.models import MambaTabRegressor

X = pd.DataFrame({"a": np.random.randn(80), "b": np.random.randn(80)})
y = np.random.randn(80)
m = MambaTabRegressor()
m.fit(X, y, max_epochs=1, accelerator="cpu")

task = m._task_model
grads = {n: p.grad for n, p in task.named_parameters() if "mamba" in n}
print(all(g is None for g in grads.values()))  # True — Mamba got no gradients

Or simply read MambaTab.forward: self.mamba does not appear.

Expected behavior
The contextualization block should be applied between the embedding/norm stage and the head, e.g. x = self.mamba(x) before pooling/squeezing, matching the MambaTab paper and the other Mamba-family architectures in the package.

Screenshots
n/a

Desktop (please complete the following information):

  • OS: macOS (Darwin 25.5.0, arm64)
  • Python version: 3.11.15
  • deeptab Version: 2.0.0 (main @ 4e6a359)

Additional context
Found via static review during a systematic bug-hunt of v2.0.0; the smoke tests pass because the degenerate model still produces finite predictions — nothing checks that the backbone actually runs. A regression test asserting non-None gradients on backbone parameters after one training step would catch this class of bug for every architecture.

Metadata

Metadata

Assignees

No one assigned

    Labels

    bugSomething isn't working

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions