Astrocyte-Hebbian Spiking Transformer Plugin

A standalone PyTorch plugin containing Astrocyte-Hebbian spiking linear-attention components with support for both Surrogate Gradients and Exact SNN (Implicit Function Theorem / IFT) Gradients.

Scope

This is a focused spiking model plugin providing:

• AstrocyteHebbianAttention: multi-head linear attention using binary Q/K/V activations with configurable gradient mode ("surrogate" or "exact").
• AstrocyteHebbianBlock: pre-norm Transformer-style block with a spiking FFN.
• AstrocyteHebbianClassifier: ready-to-train sequence classifier.
• CausalAstrocyteLanguageModel: byte-level causal language model with streaming recurrence.
• spike_fn: binary Heaviside threshold supporting fast-sigmoid surrogate gradients and exact closed-form IFT gradients.

The implementation avoids an N x N attention matrix by computing the K^T V trace. It uses ordinary dense PyTorch tensors for projections, normalization, residual paths, and training.

Gradient Modes: Exact Mode (Default) vs. Surrogate Mode

• Exact / Reciprocal Mode (gradient_mode="exact", Default): Uses a bounded inverse / reciprocal surrogate derivative $g'(x) = \frac{1}{|x| + \epsilon}$ for $|x| \le 1.0$, evaluating threshold boundary sensitivity directly out of the box. (Note: Full Time-to-First-Spike trajectory IFT solvers are available via exact_snn integration).
• Surrogate Mode (gradient_mode="surrogate"): Uses a fast-sigmoid surrogate derivative approximation $\sigma'(x) \cdot \alpha$ through firing thresholds during backpropagation.

Install

pip install -e .

For development and tests:

pip install -e .[test]
pytest -q

For the optional psMNIST benchmark:

pip install -e .[benchmark]
python astrohebbian/benchmark.py

For comparing Surrogate SNN, Exact SNN, and Dense Transformer on a causal language model task:

python benchmarks/lm_prototype.py \
    --data data/pretraining_code.jsonl \
    --max-bytes 5000000 --steps 30 --output results/lm_comparison.json

For sequence classification comparison:

python benchmarks/surrogate_vs_exact.py --epochs 5 --output results/surrogate_vs_exact.json

Example

import torch
from astrohebbian import AstrocyteHebbianClassifier, CausalAstrocyteLanguageModel

# Sequence Classifier with Reciprocal Spike Gradients
classifier = AstrocyteHebbianClassifier(
    input_dim=1,
    d_model=128,
    seq_len=784,
    num_heads=4,
    v_levels=1,
    gradient_mode="exact",  # Options: "surrogate" | "exact"
)

pixels = torch.randn(8, 784, 1)
logits = classifier(pixels)
print(logits.shape)  # torch.Size([8, 10])

# Causal Language Model with Reciprocal Spike Gradients
lm = CausalAstrocyteLanguageModel(
    vocab_size=256,
    d_model=64,
    seq_len=128,
    num_heads=4,
    gradient_mode="exact",
)

tokens = torch.randint(0, 256, (4, 128))
lm_logits = lm(tokens)
print(lm_logits.shape)  # torch.Size([4, 128, 256])

Benchmark Results

Causal Language Model (pretraining_code.jsonl, 5MB text sample, 30 steps, CUDA)

Key Observations:

• Lower Training Loss at Step 30: AstroHebbian SNN in exact/reciprocal mode reached a lower training loss (3.624) than surrogate mode (3.826) after 30 steps.
• VRAM Savings: Both SNN variants saved ~18% Peak VRAM (~29.2–29.7 MB) relative to the dense Transformer (~36.4 MB).
• Speed: Reciprocal spike gradients completed 30 steps in 1.99s (vs 2.38s for surrogate gradients).

Project status

This standalone plugin is an experimental research package for prototyping and benchmarking spiking linear transformers. Results are tracked in docs/benchmark_baseline.md and results/lm_comparison.json. Further multi-epoch training and validation are recommended before production deployment.