MambaConfig has no dropout field — training on small datasets will overfit

[eeshsaxena]

Enhancement: MambaConfig Has No dropout or weight_decay Fields — Training Is Under-Regularized

mamba/model.py defines MambaConfig as:

@dataclass
class MambaConfig:
    d_model: int = 512
    n_layers: int = 24
    vocab_size: int = 50257
    d_state: int = 16
    d_conv: int = 4
    expand: int = 2
    pad_vocab_size_multiple: int = 8
    norm_eps: float = 1e-5

There is no dropout rate in the config. Looking at mamba_block.py, neither ResidualMambaBlock nor the SSM layers apply any dropout. Without regularization, the model will overfit on small datasets, and users experimenting with training on custom corpora will get worse results than necessary.

Impact

• No dropout means no regularization during training.
• The default n_layers=24 is quite deep; without dropout, training on small-to-medium datasets will overfit.
• There is no way to configure dropout through the standard MambaConfig interface.

Suggested Fix

Add dropout to MambaConfig and apply it in ResidualMambaBlock:

@dataclass
class MambaConfig:
    ...
    dropout: float = 0.0   # set to 0.1 for training on small datasets

class ResidualMambaBlock(nn.Module):
    def __init__(self, ..., dropout=0.0):
        ...
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, x):
        residual = x
        x = self.norm(x)
        x = self.mamba(x)
        x = self.dropout(x)
        return x + residual

This matches how the original Mamba paper and reference implementation handle regularization.