GSMC-Torch: Gated Spiking Memory Cell Plugin

A standalone, production-grade PyTorch plugin library for Gated Spiking Memory Cells (GSMC).

GSMC provides a fundamental solution to the vanishing gradient problem in Spiking Neural Networks (SNNs) and Spiking Transformers via a learnable Constant-Error Carousel (CEC), while maintaining strict binary inter-neuron communication and multiplier-free operations (~41 pJ/step/neuron on 45 nm, 46× cheaper than ANN-LSTM).

---

Key Features

• Drop-in PyTorch Plugin (nn.Module): Seamlessly integrates into any PyTorch model architecture (SNN-Transformers, RNNs, ConvNets, hybrid models).
• Dual Execution Modes:
  • GSMCv2 / GSMCLayer: Vectorized sequence-to-sequence layer for fast BPTT over multi-timestep sequence tensors (batch, time, features).
  • GSMCCell: Low-level single-timestep stateful cell for step-by-step unrolling, streaming inference, or custom Transformer attention blocks.
• Vanishing-Gradient Immunity: Preserves temporal gradients over $T=784$ steps 34+ orders of magnitude above LIF baselines.
• Split-Gamma Reset ($\gamma_r = 0.1$): Eliminates the "reset tax" on temporal gradients while maintaining negative feedback stabilization.
• Exact IFT Spike Gradients (gradient_mode="exact"): Surrogate-free implicit function theorem gradient through the spike-time map dt*/dx = -γ·τ·θ / (x·(x−θ)) — zero gradient for silent neurons, bounded by the split-gamma caps, with a "hybrid" variant that adds a small arctan smearing on silent neurons to prevent dead-neuron collapse during training.
• Hardware Energy Model: Built-in 45nm CMOS energy metrics generator (compute_energy_per_step).

---

Spike Gradient Modes

GSMC supports three gradient_mode choices (default: "surrogate", so existing code is unchanged):

mode	firing neurons	silent neurons	behavior
surrogate	arctan proxy, capped	arctan proxy	robust and dense; no exact spike-time signal
exact	exact IFT flux	zero gradient	exact spike-time gradients; silent neurons stay dead once silent
hybrid	exact IFT flux	tiny arctan smearing	exact firing dynamics + dead-neuron revival

from gsmc_torch import GSMCv2, FusedGSMCv2

# Exact IFT spike gradient (calibrate init firing — silent-at-init exact cells cannot learn)
layer = GSMCv2(input_size=1, hidden_size=128, gradient_mode="exact")
layer.calibrate_init_fire(target=0.1)

# Hybrid: exact firing gradients with silent-neuron revival (recommended over pure exact)
layer = GSMCv2(input_size=1, hidden_size=128, gradient_mode="hybrid", silent_scale=0.02)

# Tune the reset-path gradient independently (e.g. surrogate reset + exact output)
layer = GSMCv2(input_size=1, hidden_size=128, gradient_mode="exact", reset_gradient_mode="surrogate", tau=1.0)

# Exact + "escape-noise" existence regularizer: keeps silent neurons alive (64/64)
# by pinning the membrane near the threshold on both sides.
layer = GSMCv2(input_size=1, hidden_size=128, gradient_mode="exact",
              existence_scale=0.05, existence_repel=0.2)
task_loss = layer(x)
total = task_loss
if layer.existence_loss is not None:
    total = task_loss + layer.existence_loss   # add it to your loss manually
total.backward()

Important caveats for "exact" (no smearing):

• Exact gradients only exist through firing neurons. A cell that fires ~0% at init is a deadlock — call calibrate_init_fire() after construction.
• During training, neurons that fall silent receive zero gradient and stay dead, progressively draining capacity. The existence_scale regularizer (relu(θ−V)² for under-active neurons + existence_repel·relu(V−θ)² for over-active ones) keeps all neurons alive, but on our first-bit recall benchmark it alone does not close the accuracy gap to dense gradients.
• Measured ranking on the first-bit recall task (test acc / firing / alive-of-64): hybrid 100% / 0.226 / 49 = surrogate 100% / 0.185 / 54 > exact + existence 80.5% / 0.285 / 64 > exact pure 52.9% / 0.032 / 12. The exact firing dynamics are best expressed through "hybrid".
• The adaptive threshold is treated as a constant for the exact gradient (theta.detach()); threshold still adapts through the forward reset path.
• All gradient-mode knobs (gradient_mode, reset_gradient_mode, tau, silent_scale, existence_scale, existence_repel) are mutable at runtime on both the layer (GSMCv2/FusedGSMCv2) and the underlying .cell (e.g. layer.gradient_mode = "hybrid").

See benchmarks/exact_vs_surrogate.py for the full reproducibility script: temporal gradient preservation over T=784, first-bit recall training curves, and fwd+bwd wall-clock parity (~parity across all modes).

---

Installation

Install directly in editable mode:

cd gsmc-plugin
pip install -e .

Or install with development dependencies:

pip install -e ".[dev]"

---

Quickstart

import torch
from gsmc_torch import GSMCv2, GSMCCell

# 1. High-level sequence layer (batch_first=True)
layer = GSMCv2(input_size=1, hidden_size=128)

# Input binary spikes: (batch=32, time=784, features=1)
x = (torch.rand(32, 784, 1) > 0.5).float()

# Forward pass -> returns output binary spikes (32, 784, 128)
spikes = layer(x)
print(f"Output spikes shape: {spikes.shape}")

# 2. Low-level stateful step cell
cell = GSMCCell(input_size=1, hidden_size=128)
state = cell.init_state(batch_size=32)

x_t = (torch.rand(32, 1) > 0.5).float()
s_next, state = cell(x_t, state)
print(f"Step spike shape: {s_next.shape}")

---

Integration into Spiking Transformers

GSMC can be used directly inside Spiking Attention blocks to provide long-horizon temporal memory:

import torch
import torch.nn as nn
from gsmc_torch import GSMCv2

class SpikingAttentionBlock(nn.Module):
    def __init__(self, embed_dim=64, hidden_dim=128):
        super().__init__()
        self.q_proj = nn.Linear(embed_dim, embed_dim)
        self.k_proj = nn.Linear(embed_dim, embed_dim)
        self.v_proj = nn.Linear(embed_dim, embed_dim)
        
        # GSMC Temporal Memory Cell replacing standard attention decay
        self.gsmc_memory = GSMCv2(input_size=embed_dim, hidden_size=hidden_dim, batch_first=True)
        self.out_proj = nn.Linear(hidden_dim, embed_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        attn_features = (self.q_proj(x) * self.k_proj(x)) + self.v_proj(x)
        spiking_attn = (attn_features > 0.0).float()
        memory_spikes = self.gsmc_memory(spiking_attn)
        return self.out_proj(memory_spikes)

---

Mathematical Formulation

Forward Dynamics

All affine operations act on binary vectors $X[t], S[t-1] \in {0, 1}$, eliminating dense multiplications on neuromorphic hardware:

$$\begin{aligned} f[t] &= \sigma(W_f X[t] + U_f S[t-1] + b_f) && \text{(Forget gate: initial } b_f=8.0) \ i[t] &= \sigma(W_i X[t] + U_i S[t-1]) && \text{(Input write gate)} \ o[t] &= \sigma(W_o X[t] + U_o S[t-1] + b_o) && \text{(Output exposure gate: initial } b_o=-2.0) \ g[t] &= \tanh(W_g X[t] + U_g S[t-1]) && \text{(Candidate state)} \ A[t] &= f[t] \odot M[t-1] + i[t] \odot g[t] && \text{(Memory-bus accumulator)} \ V[t] &= o[t] \odot \text{Norm}(A[t]) + W_d X[t] && \text{(Exposed membrane voltage)} \ S[t] &= \Theta(V[t] - \theta_t) && \text{(Spike generation)} \ M[t] &= A[t] - v_{th} \tilde{S}[t] && \text{(Refractory reset via split } \gamma_r) \end{aligned}$$

BPTT Temporal Jacobian

The temporal Jacobian decomposes into:

$$J_t = \frac{\partial M[t]}{\partial M[t-1]} = \mathrm{diag}(f[t]) + \mathcal{B}_t$$

Holding $S$ constant gives $\prod_{t=1}^T \mathrm{diag}(f[t])$, a learnable constant-error carousel immune to exponential decay.

---

Hardware Energy Footprint (45 nm CMOS)

Model	Dense MACs	Energy ($\text{pJ}/\text{step}/\text{neuron}$)	Energy vs ANN-LSTM
VanillaLIF	0	16.9 pJ	112× cheaper
GSMC v2 (gsmc_torch)	0	41.2 pJ	46× cheaper
SpikingLSTM	0	40.8 pJ	46× cheaper
ANN-LSTM	Dense ($32 \times 32$)	1900.8 pJ	1.0× (Baseline)

---

Running Tests

Execute the comprehensive Pytest suite:

pytest

---

License

MIT License. See LICENSE for details.