SpikeMamba presents a novel integration of spiking neural networks (SNNs) with the Mamba state space model architecture, investigating the potential for biologically-inspired temporal dynamics in language modeling. This research explores the computational benefits of combining Leaky Integrate-and-Fire (LIF) neurons with selective state space mechanisms, examining energy efficiency, temporal processing capabilities, and neuromorphic computing applications in large language models.

The SpikeMamba architecture consists of several key components that work in concert to process sequential data through spiking dynamics:

1. SpikingMambaLayer: The fundamental building block that integrates Mamba's selective state space mechanism with spiking neuron dynamics
2. LIFNeuron: Leaky Integrate-and-Fire neurons with adaptive thresholds and refractory periods
3. SpikingMambaBlock: Complete processing blocks with normalization and spike regularization
4. SpikingMambaLM: The full language model with embedding, multiple spiking layers, and output projection

The model supports multiple integration strategies for incorporating spiking dynamics:

• Pre-spiking: Spiking neurons process input before Mamba transformation
• Post-spiking: Spiking neurons process Mamba output
• Pre-post: Bidirectional spiking integration
• Residual: Spiking neurons in residual connections

The core spiking mechanism follows the LIF neuron model with the following differential equation:

τ_m * dV/dt = -V + I_syn + I_ext

Where:

• τ_m is the membrane time constant
• V is the membrane potential
• I_syn is the synaptic current
• I_ext is external input

The Mamba component processes the spiking output through selective state space equations:

h_t = Ā * h_{t-1} + B̄ * x_t
y_t = C * h_t + D * x_t

Where the selection mechanism adapts based on spiking activity patterns.

To enable backpropagation through the non-differentiable spike generation, we employ a fast sigmoid surrogate gradient:

∂S/∂V ≈ σ'(V - θ) = σ(V - θ) * (1 - σ(V - θ))

This allows for gradient-based optimization while maintaining the discrete nature of spike generation.

The model implements learnable, adaptive thresholds that adjust based on recent spiking activity:

θ_adapt(t) = θ_base + α * θ_scale * spike_history(t)

This mechanism enables the model to maintain appropriate firing rates across different input distributions and temporal scales.

Soft gating mechanisms combine continuous Mamba outputs with discrete spike trains:

output = mamba_out * (γ * spikes + (1 - γ) * continuous_signal)

Where γ is a learnable gating parameter that controls the balance between spiking and continuous processing.

Multiple temporal pooling approaches are implemented:

• Adaptive pooling: Based on spike activity levels
• Learnable pooling: Parameterized temporal integration
• None: Direct temporal processing

L2 regularization on membrane potentials encourages sparsity and biologically realistic firing patterns:

L_spike = λ * Σ(V_membrane²)

The integration of spiking dynamics with state space models investigates energy-efficient, event-driven computation in language models. Potential benefits include:

• Reduced power consumption through sparse activation patterns
• Hardware acceleration on neuromorphic processing units
• Biologically plausible temporal dynamics for event-based processing

The combination of spiking dynamics with state space models may provide computational advantages for:

• Long-range dependency modeling through temporal integration
• Pattern recognition in sequential data
• Event-based processing with sparse representations

The discrete nature of spikes and adaptive thresholds may facilitate:

• Mitigation of catastrophic forgetting through sparse representations
• Online learning capabilities with dynamic threshold adaptation
• Task-specific adaptation through spike pattern modulation

from spike_mamba.main import create_spiking_mamba_model, SpikingMambaConfig, MambaConfig

# Basic model configuration
model = create_spiking_mamba_model(
    d_model=512,           # Model dimension
    n_layer=6,             # Number of layers
    vocab_size=1000,       # Vocabulary size
    spike_mode="pre_post", # Integration mode
    threshold=1.0,         # Spike threshold
    tau_mem=20.0,          # Membrane time constant
    adaptive_threshold=True, # Enable adaptive thresholds
    spike_regularization=0.01 # Regularization strength
)

# Advanced configuration with custom parameters
mamba_config = MambaConfig(
    d_model=768,
    n_layer=12,
    vocab_size=50277,
    d_state=16,
    d_conv=4,
    expand=2
)

spiking_config = SpikingMambaConfig(
    mamba_config=mamba_config,
    threshold=1.5,
    tau_mem=25.0,
    tau_syn=5.0,
    reset_mode="subtract",
    adaptive_threshold=True,
    refractory_period=3,
    spike_regularization=0.02,
    spike_integration="pre_post",
    temporal_pooling="adaptive"
)

import torch
import torch.nn.functional as F
from spike_mamba.main import SpikingMambaTrainer

# Initialize trainer
trainer = SpikingMambaTrainer(
    model=model,
    spike_loss_weight=0.01,
    enable_logging=True
)

# Training loop example
def train_step(model, input_ids, targets):
    model.train()
    
    # Forward pass
    output = model(input_ids, return_spike_stats=True)
    
    # Compute loss
    total_loss, loss_dict = trainer.compute_loss(
        output.logits, 
        targets, 
        output.spike_reg_loss
    )
    
    # Backward pass
    total_loss.backward()
    
    return total_loss, loss_dict, output.spike_stats

# Example training iteration
input_ids = torch.randint(0, 1000, (batch_size, seq_len))
targets = input_ids[:, 1:]  # Next token prediction
targets = torch.cat([targets, input_ids[:, :1]], dim=1)  # Shift for causal LM

loss, loss_dict, spike_stats = train_step(model, input_ids, targets)
print(f"Total loss: {loss.item():.4f}")
print(f"Spike rate: {spike_stats.spike_rate:.4f}")

• Gradient Flow: Surrogate gradients enable end-to-end training through discrete spike generation
• Spike Regularization: Balances task performance with biological realism through membrane potential regularization
• State Management: Careful handling of temporal states across layers for proper spike dynamics
• Memory Efficiency: Sparse activations may reduce memory requirements compared to dense models

1. Computational Overhead: Spiking dynamics introduce additional computational complexity compared to standard transformers
2. Training Stability: Surrogate gradients may introduce training instabilities, particularly with high spike rates
3. Hyperparameter Sensitivity: Multiple spiking parameters (thresholds, time constants, refractory periods) require careful tuning
4. Evaluation Metrics: Standard NLP metrics may not adequately capture spiking-specific benefits such as energy efficiency
5. Memory Requirements: State management across layers increases memory overhead during training

1. Hardware Implementation: Investigation of neuromorphic hardware compatibility and acceleration
2. Energy Efficiency: Quantification of power consumption benefits through sparse activation patterns
3. Biological Plausibility: Comparison with biological neural networks and validation of temporal dynamics
4. Task-Specific Optimization: Adaptation for specific NLP tasks and evaluation of performance trade-offs
5. Scaling Properties: Investigation of model behavior at larger scales and longer sequences

pip install mamba-ssm torch loguru

import torch
from spike_mamba.main import create_spiking_mamba_model

# Create model
model = create_spiking_mamba_model(
    d_model=512,
    n_layer=6,
    vocab_size=1000,
    spike_mode="pre_post"
).to('cuda')

# Forward pass
input_ids = torch.randint(0, 1000, (2, 64)).to('cuda')
output = model(input_ids, return_spike_stats=True)
print(f"Spike rate: {output.spike_stats.spike_rate:.4f}")

from spike_mamba.main import SpikingMambaGenerator

# Initialize generator
generator = SpikingMambaGenerator(
    model=model,
    tokenizer=tokenizer,  # Your tokenizer
    enable_logging=True
)

# Generate text with spike monitoring
result = generator.generate(
    input_ids=input_ids,
    max_length=100,
    temperature=0.8,
    top_k=50,
    top_p=0.9,
    return_spike_stats=True
)

# Analyze spike patterns
for i, spike_stats in enumerate(result['spike_stats_history']):
    print(f"Step {i}: Spike rate = {spike_stats.spike_rate:.4f}")
    print(f"  Membrane potential: {spike_stats.avg_membrane_potential:.4f}")

from spike_mamba.main import SpikingMambaConfig, MambaConfig, LIFNeuron

# Create custom LIF neuron
lif_config = SpikingMambaConfig(
    mamba_config=MambaConfig(d_model=256),
    threshold=1.2,
    tau_mem=15.0,
    tau_syn=3.0,
    adaptive_threshold=True,
    refractory_period=2,
    spike_regularization=0.005
)

lif_neuron = LIFNeuron(lif_config, d_model=256)

# Test LIF neuron
x = torch.randn(1, 10, 256)
spikes, state = lif_neuron(x)
print(f"Spike output shape: {spikes.shape}")
print(f"Spike rate: {torch.mean(spikes).item():.4f}")

# Enable detailed logging
import logging
logging.basicConfig(level=logging.DEBUG)

# Create model with logging
model = create_spiking_mamba_model(
    d_model=256,
    n_layer=4,
    vocab_size=1000,
    enable_logging=True,
    log_spike_stats=True
)

# Forward pass with detailed output
output = model(input_ids, return_spike_stats=True)

# Access detailed spike statistics
spike_stats = output.spike_stats
print(f"Total spikes: {spike_stats.total_spikes}")
print(f"Layer spike rates: {spike_stats.layer_spike_rates}")
print(f"Average membrane potential: {spike_stats.avg_membrane_potential:.4f}")
print(f"Max membrane potential: {spike_stats.max_membrane_potential:.4f}")

This is an active research project exploring the intersection of neuromorphic computing and large language models. We welcome contributions from researchers interested in:

• Spiking neural networks and temporal dynamics
• State space models and sequence modeling
• Neuromorphic computing and hardware acceleration
• Language modeling and natural language processing
• Biologically inspired artificial intelligence

• Algorithm Development: Novel spiking mechanisms and integration strategies
• Hardware Implementation: Neuromorphic chip compatibility and optimization
• Theoretical Analysis: Mathematical foundations and convergence properties
• Empirical Evaluation: Benchmarking and performance analysis
• Biological Validation: Comparison with biological neural networks

Join our research community focused on advancing neuromorphic language models and biologically inspired AI architectures.

Note: This is a work-in-progress research project. The architecture and algorithms are under active development and may change significantly as we explore the potential of spiking neural networks in language modeling.