Package: hsss
HSSSImplementation of a Hierarchical Mamba as described in the paper: "Hierarchical State Space Models for Continuous Sequence-to-Sequence Modeling" but instead of using traditional SSMs were using Mambas. Basically the flow is single input -> low level mambas -> concat -> high level ssm -> multiple outputs.Paper linkI believe in this architecture alot as it segments local and global learning.installpip install hsssusageimporttorchfromhsssimportLowLevelMamba,HSSS# Reandom tensorx=torch.randn(1,10,8)# Low level modelmamba=LowLevelMamba(dim=8,# dimension of inputdepth=6,# depth of inputdt_rank=4,# rank of inputd_state=4,# state of inputexpand_factor=4,# expansion factor of inputd_conv=6,# convolution dimension of inputdt_min=0.001,# minimum time step of inputdt_max=0.1,# maximum time step of inputdt_init="random",# initialization method of inputdt_scale=1.0,# scaling factor of inputbias=False,# whether to use bias in inputconv_bias=True,# whether to use bias in convolution of inputpscan=True,# whether to use parallel scan in input)# Low level model 2mamba2=LowLevelMamba(dim=8,# dimension of inputdepth=6,# depth of inputdt_rank=4,# rank of inputd_state=4,# state of inputexpand_factor=4,# expansion factor of inputd_conv=6,# convolution dimension of inputdt_min=0.001,# minimum time step of inputdt_max=0.1,# maximum time step of inputdt_init="random",# initialization method of inputdt_scale=1.0,# scaling factor of inputbias=False,# whether to use bias in inputconv_bias=True,# whether to use bias in convolution of inputpscan=True,# whether to use parallel scan in input)# Low level mamba 3mamba3=LowLevelMamba(dim=8,# dimension of inputdepth=6,# depth of inputdt_rank=4,# rank of inputd_state=4,# state of inputexpand_factor=4,# expansion factor of inputd_conv=6,# convolution dimension of inputdt_min=0.001,# minimum time step of inputdt_max=0.1,# maximum time step of inputdt_init="random",# initialization method of inputdt_scale=1.0,# scaling factor of inputbias=False,# whether to use bias in inputconv_bias=True,# whether to use bias in convolution of inputpscan=True,# whether to use parallel scan in input)# HSSShsss=HSSS(layers=[mamba,mamba2,mamba3],dim=12,# dimension of modeldepth=3,# depth of modeldt_rank=2,# rank of modeld_state=2,# state of modelexpand_factor=2,# expansion factor of modeld_conv=3,# convolution dimension of modeldt_min=0.001,# minimum time step of modeldt_max=0.1,# maximum time step of modeldt_init="random",# initialization method of modeldt_scale=1.0,# scaling factor of modelbias=False,# whether to use bias in modelconv_bias=True,# whether to use bias in convolution of modelpscan=True,# whether to use parallel scan in modelproj_layer=True,)# Forward passout=hsss(x)print(out)Citation@misc{bhirangi2024hierarchical,title={Hierarchical State Space Models for Continuous Sequence-to-Sequence Modeling},author={Raunaq Bhirangi and Chenyu Wang and Venkatesh Pattabiraman and Carmel Majidi and Abhinav Gupta and Tess Hellebrekers and Lerrel Pinto},year={2024},eprint={2402.10211},archivePrefix={arXiv},primaryClass={cs.LG}}LicenseMITTodoImplement the chunking of the tokens by spliting it up the sequence dimensionMake the fusion projection layer dynamic and not use just a linear, ffn, or cross attention or even an output head.