STAC: Spiking Transformer Augmenting Cognition for Conversational AI

Overview

STAC (Spiking Transformer Augmenting Cognition) is a research framework that explores two complementary approaches to spiking neural network (SNN) language modeling:

• STAC V1: A complete end-to-end training pipeline built around learnable Adaptive Exponential (AdEx) neurons. See stac_v1/.
• STAC V2: An experimental conversion framework that transforms pretrained transformer language models (DistilGPT-2, SmolLM2-1.7B-Instruct) into SNNs, targeting 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. Energy figures reported here are theoretical projections derived from spike-count analysis, not measured hardware data.

Key Features

• Proof-of-concept ANN-to-SNN conversion built on SpikingJelly.
• Multi-turn context retention via a Temporal Spike Processor.
• Test coverage for position IDs, KV-cache behavior, and spike-rate sanity checks.
• Hardware power profiling: planned, not yet implemented.
• Full operator coverage and optimization: work in progress.

Quick Start

# 1. Install dependencies
pip install -r requirements.txt

# 2. Convert DistilGPT-2 to an SNN
python run_conversion.py --model_name distilgpt2 --timesteps 8 --simplified

# 3. Run a multi-turn conversation smoke test
python snn_multi_turn_conversation_test.py --mode snn --turns 3 --timesteps 8

# 4. Run the comprehensive validation suite
python test_conversational_snn.py --model_name distilgpt2 --test_all --timesteps 8

Core Components

STAC V2

Component	Purpose
smollm2_converter.py	Specialized converter with TemporalSpikeProcessor.
convert.py	Generic ANN-to-SNN conversion pipeline.
run_conversion.py	Main CLI entry point for conversions.
spikingjelly_compat.py	Cross-version compatibility layer for SpikingJelly.
test_conversational_snn.py	Comprehensive test suite.
snn_multi_turn_conversation_test.py	Lightweight multi-turn smoke test.

STAC V1

Component	Purpose
stac-v1/stacv1.ipynb	End-to-end training pipeline with learnable AdEx neurons.
stac-v1/README.md	V1 documentation and research notes.
stac_v1/ + run_stac_v1.py	Repo-native runnable V1 pipeline demonstrating hybrid fine-tuning (frozen GPT-2 with a trained spiking and memory head).

Implementation Status

STAC V2

Completed (prototype level)

• Core conversion flow: GELU-to-ReLU substitution, quantization, and ann2snn.
• Temporal dynamics and KV-cache handling in PyTorch.
• Spike-count telemetry hooks and accompanying unit tests.
• Loihi export gating (requires EXPORT_LOIHI=1 and lava.lib.dl.slayer; otherwise the pipeline remains simulation-only and Loihi tests are skipped).

Pending or in progress

• Hardware benchmarking on Loihi-2 and Akida.
• Expanded operator support (rotary embeddings, flash-attention variants, etc.).
• Integration with the SCANUE multi-agent alignment layer.
• CLI, UX, and documentation polish.

STAC V1

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.
• Validation suite covering the full pipeline.

Documentation

STAC V2

• Conversion Workflow: step-by-step conversion guide.
• API Reference: function and class documentation.
• Hardware Requirements: system specifications.

STAC V1

• STAC V1 Documentation: end-to-end training pipeline documentation.
• STAC V1 Implementation: Jupyter notebook with learnable AdEx neurons.

Testing and 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 the full suite
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.