Complex-valued training about mamba HOT 5 OPEN

I also wondered about this matter (similarly raised in annotated-S4 and S5) and found a clue that "the sum in (7) will implicitly include the conjugate pairs of A, B, C and therefore resolve to twice the real part of the original sum." in S4D paper.
Also, check out the mathematical properties of complex conjugates.

from mamba.

I also wanna train a complex-valued Mamba model. And according to the state-space model paper I have read before and the Mamba repository, I changed the initialization of A into complex-valued which was proposed in S4D.

real-valued initialization :

complex-valued initialization:

        # # S4D real initialization
        # A = repeat(
        #     torch.arange(1, self.d_state + 1, dtype=torch.float32, device=device),
        #    "n -> d n",
        #    d=self.d_inner,
        # ).contiguous()
        # A_log = torch.log(A)  # Keep A_log in fp32
        # self.A_log = nn.Parameter(A_log)
        # self.A_log._no_weight_decay = True        

        # S4D complex initialization
        log_A_real = torch.log(0.5 * torch.ones(self.d_inner, self.d_state))
        A_imag = math.pi * repeat(torch.arange(self.d_state), 'n -> h n', h=self.d_inner)
        
        self.log_A_real = nn.Parameter(log_A_real)
        self.log_A_real._no_weight_decay = True 
        
        self.A_imag = nn.Parameter(A_imag)
        self.A_imag._no_weight_decay = True 

        # real-valued
        # A = -torch.exp(self.A_log.float())  # (d_inner, d_state)

        # complex-valued        
        A = -torch.exp(self.log_A_real) + 1j * self.A_imag

By doing this, the state space is computed in the complex-value domain instead of the real-value domain. However, the state-space model still perform the mapping from the real-sequence u to the real-value sequence y. The inputs of the Mamba still be real-valued. This paper may be helpful.

Hope this can help you.

from mamba.

@Tworan Did you encounter this error when modifying it like the above?

RuntimeError: B must have shape (batch_size, n_groups, dstate, !is_complex ? seqlen : seqlen * 2)

from mamba.

@hungdche Same problem here. There is a mismatch between the .cpp function interface and complex input. Have you solved the problem yet? Thanks!

from mamba.

@Lily-Le I have not. Unsure how to proceed

from mamba.