In ssm.py the A matrix is initialized as:
A = repeat(torch.arange(1, d_state + 1, dtype=torch.float32), "n -> d n", d=d_model)
A_log = torch.log(A)
This is the diagonal HiPPO initialization from S4, but Mamba (Gu & Dao, 2023) specifically uses a different initialization - A is initialized with random negative values and the B matrix is initialized with a uniform distribution. Using the S4 HiPPO init in a "selective" SSM gives you something closer to S4 behavior than Mamba behavior.
The comment in the code even says "following S4" which confirms this. Should either use the correct Mamba initialization from the paper or note clearly that this deviates from it.
In ssm.py the A matrix is initialized as:
A = repeat(torch.arange(1, d_state + 1, dtype=torch.float32), "n -> d n", d=d_model)
A_log = torch.log(A)
This is the diagonal HiPPO initialization from S4, but Mamba (Gu & Dao, 2023) specifically uses a different initialization - A is initialized with random negative values and the B matrix is initialized with a uniform distribution. Using the S4 HiPPO init in a "selective" SSM gives you something closer to S4 behavior than Mamba behavior.
The comment in the code even says "following S4" which confirms this. Should either use the correct Mamba initialization from the paper or note clearly that this deviates from it.