Update core/mamba_block.py

core/mamba_block.py

@@ -7,7 +7,7 @@ import math
 class MambaBlock(nn.Module):
     """
     Production-ready Mamba block for graph processing
-    Device-safe implementation
+    Device-safe implementation with optimizations
     """
     def __init__(self, d_model, d_state=16, d_conv=4, expand=2, dt_rank="auto", bias=False):
         super().__init__()
@@ -50,6 +50,21 @@ class MambaBlock(nn.Module):
         # Activation
         self.act = nn.SiLU()
         
+        # Initialize parameters
+        self._init_parameters()
+        
+    def _init_parameters(self):
+        """Initialize parameters with proper scaling"""
+        # Initialize dt projection specially
+        dt_init_std = self.dt_rank**-0.5 * self.d_state
+        with torch.no_grad():
+            self.dt_proj.bias.uniform_(-dt_init_std, dt_init_std)
+            
+        # Initialize other projections
+        nn.init.xavier_uniform_(self.in_proj.weight)
+        nn.init.xavier_uniform_(self.x_proj.weight)
+        nn.init.xavier_uniform_(self.out_proj.weight)
+        
     def forward(self, x):
         """
         x: (batch, length, d_model)