BitMamba module for performing computations using the BitNet architecture. Args: dim (int): The input dimension (D). depth (int): The depth of the BitNet architecture. dt_rank (Union[int, str], optional): The rank of the time step tensor. Defaults to "auto".
| 544 | |
| 545 | |
| 546 | class BitMamba(nn.Module): |
| 547 | """ |
| 548 | BitMamba module for performing computations using the BitNet architecture. |
| 549 | |
| 550 | Args: |
| 551 | dim (int): The input dimension (D). |
| 552 | depth (int): The depth of the BitNet architecture. |
| 553 | dt_rank (Union[int, str], optional): The rank of the time step tensor. Defaults to "auto". |
| 554 | d_state (int, optional): The dimension of the state tensor (N in paper/comments). Defaults to 16. |
| 555 | expand_factor (int, optional): The expansion factor for the inner dimension (E in paper/comments). Defaults to 2. |
| 556 | d_conv (int, optional): The dimension of the convolutional filters. Defaults to 4. |
| 557 | dt_min (float, optional): The minimum value for the time step. Defaults to 0.001. |
| 558 | dt_max (float, optional): The maximum value for the time step. Defaults to 0.1. |
| 559 | dt_init (str, optional): The initialization method for the time step. Can be "random" or "constant". Defaults to "random". |
| 560 | dt_scale (float, optional): The scaling factor for the time step. Defaults to 1.0. |
| 561 | dt_init_floor (float, optional): The floor value for the initialized time step. Defaults to 1e-4. |
| 562 | bias (bool, optional): Whether to include bias terms. Defaults to False. |
| 563 | conv_bias (bool, optional): Whether to include bias terms in the convolutional layers. Defaults to True. |
| 564 | pscan (bool, optional): Whether to use parallel scan mode or sequential mode when training. Defaults to True. |
| 565 | """ |
| 566 | |
| 567 | def __init__( |
| 568 | self, |
| 569 | dim: int, # D |
| 570 | num_tokens: int, |
| 571 | sequence_length: int, |
| 572 | depth: int, |
| 573 | dt_rank: Union[int, str] = "auto", |
| 574 | d_state: int = 16, # N in paper/comments |
| 575 | expand_factor: int = 2, # E in paper/comments |
| 576 | d_conv: int = 4, |
| 577 | dt_min: float = 0.001, |
| 578 | dt_max: float = 0.1, |
| 579 | dt_init: str = "random", # "random" or "constant" |
| 580 | dt_scale: float = 1.0, |
| 581 | dt_init_floor=1e-4, |
| 582 | bias: bool = False, |
| 583 | conv_bias: bool = True, |
| 584 | pscan: bool = True, # use parallel scan mode or sequential mode when training |
| 585 | return_embeddings: bool = True, |
| 586 | return_tokens: bool = True, |
| 587 | *args, |
| 588 | **kwargs |
| 589 | ): |
| 590 | super().__init__(*args, **kwargs) |
| 591 | self.dim = dim |
| 592 | self.num_token = num_tokens |
| 593 | self.sequence_length = sequence_length |
| 594 | self.depth = depth |
| 595 | self.dt_rank = dt_rank |
| 596 | self.d_state = d_state |
| 597 | self.expand_factor = expand_factor |
| 598 | self.d_conv = d_conv |
| 599 | self.dt_min = dt_min |
| 600 | self.dt_max = dt_max |
| 601 | self.dt_init = dt_init |
| 602 | self.dt_scale = dt_scale |
| 603 | self.dt_init_floor = dt_init_floor |