stac
STAC: Spiking Transformer Augmenting Cognition for Conversational AI
Overview
STAC (Spiking Transformer Augmenting Cognition) is a research framework with two distinct approaches:
- STAC V1: Complete end-to-end training pipeline with learnable AdEx neurons (see
stac-v1/) - STAC V2: Experimental conversion framework that transforms pretrained transformer LLMs (DistilGPT-2, SmolLM2-1.7B-Instruct) into Spiking Neural Networks (SNNs) for potential energy savings while retaining multi-turn conversational ability in simulation
⚠️ Important: This repository currently runs software-level SNN simulations only. No metrics have been collected on physical neuromorphic hardware yet. Energy savings figures are theoretical projections based on spike-count analysis, not measured hardware data.
Key Features
✔️ Proof-of-concept ANN→SNN conversion using SpikingJelly
✔️ Multi-turn context retention via a Temporal Spike Processor
✔️ Extensive software tests for position IDs, KV-cache, and spike-rate sanity
➖ Hardware power profiling — planned, not implemented
➖ Full operator coverage & optimisation — work in progress
Quick Start
# 1. Install dependencies
pip install -r requirements.txt
# 2. Convert DistilGPT-2 to SNN (fast)
python run_conversion.py --model_name distilgpt2 --timesteps 8 --simplified
# 3. Test multi-turn conversation
python snn_multi_turn_conversation_test.py --mode snn --turns 3 --timesteps 8
# 4. Run comprehensive validation
python test_conversational_snn.py --model_name distilgpt2 --test_all --timesteps 8
Core Components
STAC V2 (Current)
| Component | Purpose |
|———–|———|
| smollm2_converter.py | Specialized converter with TemporalSpikeProcessor |
| convert.py | Generic ANN→SNN conversion pipeline |
| run_conversion.py | Main CLI entry point for conversions |
| spikingjelly_compat.py | Cross-version compatibility layer |
| test_conversational_snn.py | Comprehensive test suite (1K+ lines) |
| snn_multi_turn_conversation_test.py | Simple conversation smoke test |
STAC V1 (Original Research)
| Component | Purpose |
|———–|———|
| stac-v1/stacv1.ipynb | Complete end-to-end training pipeline with learnable AdEx neurons |
| stac-v1/README.md | V1 documentation and research contributions |
Implementation Status
STAC V2 (Current)
Completed (prototype level)
- ✅ Core conversion flow (GELU→ReLU, quantization, ann2snn)
- ✅ Temporal dynamics & KV-cache handling in PyTorch
- ✅ Spike-count telemetry hooks and unit tests
Pending / In Progress
- ⏳ Hardware benchmarking on Loihi-2 / Akida
- ⏳ Expanded operator support (e.g., rotary embeddings, flash-attention variants)
- ⏳ Integration with SCANUE multi-agent alignment layer
- ⏳ Robust CLI/UX and documentation polish
STAC V1 (Complete)
Completed (research prototype)
- ✅ End-to-end training pipeline with learnable AdEx neurons
- ✅ Hyperdimensional Memory Module (HEMM) integration
- ✅ Surrogate gradient training on WikiText-2
- ✅ L1 spike regularization for energy efficiency
- ✅ Comprehensive validation suite
Documentation
STAC V2 (Current)
- 🔄 Conversion Workflow - Step-by-step conversion guide
- 📚 API Reference - Function and class documentation
- 🖥️ Hardware Requirements - System specifications
STAC V1 (Original Research)
- 📖 STAC V1 Documentation - End-to-end training pipeline documentation
- 🧠 STAC V1 Implementation - Complete Jupyter notebook with learnable AdEx neurons
Testing & Validation
The repository includes extensive testing for multi-turn conversational correctness:
# Test specific components
python test_conversational_snn.py --model_name distilgpt2 --test_position_boundaries
python test_conversational_snn.py --model_name distilgpt2 --test_attention_mask
python test_conversational_snn.py --model_name distilgpt2 --test_multi_turn
python test_conversational_snn.py --model_name distilgpt2 --test_energy
# Run all tests
python test_conversational_snn.py --model_name distilgpt2 --test_all
License
This project is licensed under the MIT License - see the LICENSE file for details.