| 195 | |
| 196 | class MambaBlock(nn.Module): |
| 197 | def __init__(self, config: MambaConfig): |
| 198 | super().__init__() |
| 199 | |
| 200 | self.config = config |
| 201 | |
| 202 | # projects block input from D to 2*ED (two branches) |
| 203 | self.in_proj = BitLinear(config.dim, 2 * config.d_inner, bias=config.bias) |
| 204 | |
| 205 | self.conv1d = nn.Conv1d( |
| 206 | in_channels=config.d_inner, |
| 207 | out_channels=config.d_inner, |
| 208 | kernel_size=config.d_conv, |
| 209 | bias=config.conv_bias, |
| 210 | groups=config.d_inner, |
| 211 | padding=config.d_conv - 1, |
| 212 | ) |
| 213 | |
| 214 | # projects x to input-dependent Δ, B, C |
| 215 | self.x_proj = BitLinear( |
| 216 | config.d_inner, |
| 217 | config.dt_rank + 2 * config.d_state, |
| 218 | bias=False, |
| 219 | ) |
| 220 | |
| 221 | # projects Δ from dt_rank to d_inner |
| 222 | self.dt_proj = BitLinear(config.dt_rank, config.d_inner, bias=True) |
| 223 | |
| 224 | # dt initialization |
| 225 | # dt weights |
| 226 | dt_init_std = config.dt_rank**-0.5 * config.dt_scale |
| 227 | if config.dt_init == "constant": |
| 228 | nn.init.constant_(self.dt_proj.weight, dt_init_std) |
| 229 | elif config.dt_init == "random": |
| 230 | nn.init.uniform_(self.dt_proj.weight, -dt_init_std, dt_init_std) |
| 231 | else: |
| 232 | raise NotImplementedError |
| 233 | |
| 234 | # dt bias |
| 235 | dt = torch.exp( |
| 236 | torch.rand(config.d_inner) |
| 237 | * (math.log(config.dt_max) - math.log(config.dt_min)) |
| 238 | + math.log(config.dt_min) |
| 239 | ).clamp(min=config.dt_init_floor) |
| 240 | inv_dt = dt + torch.log( |
| 241 | -torch.expm1(-dt) |
| 242 | ) # inverse of softplus: https://github.com/pytorch/pytorch/issues/72759 |
| 243 | with torch.no_grad(): |
| 244 | self.dt_proj.bias.copy_(inv_dt) |
| 245 | # self.dt_proj.bias._no_reinit = True # initialization would set all Linear.bias to zero, need to mark this one as _no_reinit |
| 246 | # todo : explain why removed |
| 247 | |
| 248 | # S4D real initialization |
| 249 | A = torch.arange(1, config.d_state + 1, dtype=torch.float32).repeat( |
| 250 | config.d_inner, 1 |
| 251 | ) |
| 252 | self.A_log = nn.Parameter( |
| 253 | torch.log(A) |
| 254 | ) # why store A in log ? to keep A < 0 (cf -torch.exp(...)) ? for gradient stability ? |