Metadata-Version: 2.4
Name: gsmc-torch
Version: 0.2.0
Summary: Gated Spiking Memory Cell (GSMC) -- Standalone PyTorch plugin for vanishing-gradient-free spiking neural networks
Author: Sumit
License: MIT
Project-URL: Documentation, https://github.com/Griffith-7/GSMC-SNN#readme
Project-URL: Repository, https://github.com/Griffith-7/GSMC-SNN
Project-URL: Issues, https://github.com/Griffith-7/GSMC-SNN/issues
Keywords: spiking-neural-networks,snn,neuromorphic,vanishing-gradients,constant-error-carousel,pytorch,snn-transformer,plugin
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.9
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.0
Provides-Extra: dev
Requires-Dist: pytest>=7.0; extra == "dev"
Requires-Dist: torchvision>=0.15; extra == "dev"
Dynamic: license-file

# GSMC-Torch: Gated Spiking Memory Cell Plugin

[![Python 3.9+](https://img.shields.io/badge/python-3.9%2B-blue.svg)](https://www.python.org/downloads/)
[![PyTorch 2.0+](https://img.shields.io/badge/pytorch-2.0%2B-ee4c2c.svg)](https://pytorch.org/)
[![License: MIT](https://img.shields.io/badge/license-MIT-green.svg)](LICENSE)

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 |

```python
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:

```bash
cd gsmc-plugin
pip install -e .
```

Or install with development dependencies:

```bash
pip install -e ".[dev]"
```

---

## Quickstart

```python
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:

```python
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:

```bash
pytest
```

---

## License

MIT License. See [LICENSE](LICENSE) for details.
