astrohebbian 0.1.0__tar.gz

astrohebbian-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

astrohebbian-0.1.0/PKG-INFO

@@ -0,0 +1,150 @@
+Metadata-Version: 2.4
+Name: astrohebbian
+Version: 0.1.0
+Summary: Standalone Astrocyte-Hebbian spiking linear Transformer components
+Author: Sumith
+License: MIT License
+        
+        Copyright (c) 2026
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all
+        copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+        
+Project-URL: Homepage, https://github.com/Griffith-7/Snn-Transformer
+Project-URL: Source, https://github.com/Griffith-7/Snn-Transformer
+Project-URL: Issue Tracker, https://github.com/Griffith-7/Snn-Transformer/issues
+Keywords: spiking-neural-networks,transformer,attention,hebbian,snns,pytorch
+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.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.0
+Provides-Extra: test
+Requires-Dist: pytest>=7; extra == "test"
+Provides-Extra: benchmark
+Requires-Dist: torchvision>=0.15; extra == "benchmark"
+Requires-Dist: matplotlib>=3.7; extra == "benchmark"
+Provides-Extra: quality
+Requires-Dist: ruff>=0.6; extra == "quality"
+Dynamic: license-file
+
+# Astrocyte-Hebbian Spiking Transformer Plugin
+
+A standalone PyTorch plugin containing Astrocyte-Hebbian spiking linear-attention components. This project is intentionally independent from Exact-SNN; no Exact-SNN code, imports, or dependencies are included.
+
+## Scope
+
+This is a focused model plugin, not a general SNN framework. It provides:
+
+- `AstrocyteHebbianAttention`: multi-head linear attention using binary Q/K/V activations.
+- `AstrocyteHebbianBlock`: pre-norm Transformer-style block with a spiking FFN.
+- `AstrocyteHebbianClassifier`: ready-to-train sequence classifier.
+- `spike_fn`: binary Heaviside forward pass with surrogate gradients.
+
+The implementation avoids an `N x N` attention matrix by computing the `K^T V` trace. It still uses ordinary dense PyTorch tensors for projections, normalization, residual paths, and training.
+
+## What this is (and is not)
+
+- **Binary activation spikes**: the inter-layer signals (Q, K, V, and the FFN hidden activation) are binary Heaviside spikes (0/1).
+- **Surrogate gradients**: training uses a fast-sigmoid surrogate gradient through the spike threshold; it is not an exact spike-time gradient library.
+- **Dense PyTorch execution**: forward/backward run on ordinary dense GPU tensors. This is a CPU/GPU software package, not an event-driven neuromorphic-hardware implementation, and reported times/memory are wall-clock/FLOP measurements, not hardware energy.
+- **Full-sequence psMNIST mode**: the core `AstrocyteHebbianClassifier` (and the frozen baseline below) is full-sequence attention over N=784 pixels.
+- **Separate causal LM experimental mode**: a distinct, experimental causal path (`CausalAstrocyteLanguageModel`) is provided for small language-model proof-of-concept work only.
+
+This is a **focused model plugin, not a complete SNN framework**.
+
+## Install
+
+```bash
+pip install -e .
+```
+
+For development and tests:
+
+```bash
+pip install -e .[test]
+pytest -q
+```
+
+For the optional psMNIST benchmark:
+
+```bash
+pip install -e .[benchmark]
+python astrohebbian/benchmark.py
+```
+
+For a controlled three-seed summary:
+
+```bash
+python benchmarks/multi_seed.py --seeds 1 2 3 --output results/multi_seed.json
+```
+
+For the small causal language-model proof of concept:
+
+```bash
+python benchmarks/lm_prototype.py \
+    --data /path/to/pretraining_code.jsonl \
+    --max-bytes 10000000 --steps 50 --output results/lm_prototype.json
+```
+
+The causal LM is a separate experimental path. The full-sequence psMNIST model
+and its baseline remain unchanged.
+
+## Example
+
+```python
+import torch
+from astrohebbian import AstrocyteHebbianClassifier
+
+model = AstrocyteHebbianClassifier(
+    input_dim=1,
+    d_model=128,
+    seq_len=784,
+    num_heads=4,
+    v_levels=1,
+)
+
+pixels = torch.randn(8, 784, 1)
+logits = model(pixels)
+print(logits.shape)  # torch.Size([8, 10])
+```
+
+## Results
+
+Frozen three-seed psMNIST baseline (N=784, 60k train / 10k test, 6 epochs, RTX 3050, batch 64):
+
+| Model | 3-seed mean test acc | Peak VRAM |
+| --- | ---: | ---: |
+| AstroHebbian Pure SNN | **86.82% ± 2.48%** | ~1004 MB |
+| Transformer (dense O(N²)) | 77.80% ± 3.06% | ~1565 MB |
+
+The SNN beats the dense baseline by **+9.03 pts accuracy** at **~36% lower peak VRAM**.
+Mean runtime is not a headline figure: seed 2 was a large hardware/runtime outlier,
+so only accuracy and memory are claimed as reliable. Full per-seed data:
+`results/multi_seed.json`; the N-scaling memory crossover is in `results/n_scaling.png`.
+
+## Project status
+
+This is the clean standalone starting point for production hardening. The current attention is full-sequence rather than causal or streaming. Results are tracked in `docs/benchmark_baseline.md`.

astrohebbian-0.1.0/README.md

@@ -0,0 +1,98 @@
+# Astrocyte-Hebbian Spiking Transformer Plugin
+
+A standalone PyTorch plugin containing Astrocyte-Hebbian spiking linear-attention components. This project is intentionally independent from Exact-SNN; no Exact-SNN code, imports, or dependencies are included.
+
+## Scope
+
+This is a focused model plugin, not a general SNN framework. It provides:
+
+- `AstrocyteHebbianAttention`: multi-head linear attention using binary Q/K/V activations.
+- `AstrocyteHebbianBlock`: pre-norm Transformer-style block with a spiking FFN.
+- `AstrocyteHebbianClassifier`: ready-to-train sequence classifier.
+- `spike_fn`: binary Heaviside forward pass with surrogate gradients.
+
+The implementation avoids an `N x N` attention matrix by computing the `K^T V` trace. It still uses ordinary dense PyTorch tensors for projections, normalization, residual paths, and training.
+
+## What this is (and is not)
+
+- **Binary activation spikes**: the inter-layer signals (Q, K, V, and the FFN hidden activation) are binary Heaviside spikes (0/1).
+- **Surrogate gradients**: training uses a fast-sigmoid surrogate gradient through the spike threshold; it is not an exact spike-time gradient library.
+- **Dense PyTorch execution**: forward/backward run on ordinary dense GPU tensors. This is a CPU/GPU software package, not an event-driven neuromorphic-hardware implementation, and reported times/memory are wall-clock/FLOP measurements, not hardware energy.
+- **Full-sequence psMNIST mode**: the core `AstrocyteHebbianClassifier` (and the frozen baseline below) is full-sequence attention over N=784 pixels.
+- **Separate causal LM experimental mode**: a distinct, experimental causal path (`CausalAstrocyteLanguageModel`) is provided for small language-model proof-of-concept work only.
+
+This is a **focused model plugin, not a complete SNN framework**.
+
+## Install
+
+```bash
+pip install -e .
+```
+
+For development and tests:
+
+```bash
+pip install -e .[test]
+pytest -q
+```
+
+For the optional psMNIST benchmark:
+
+```bash
+pip install -e .[benchmark]
+python astrohebbian/benchmark.py
+```
+
+For a controlled three-seed summary:
+
+```bash
+python benchmarks/multi_seed.py --seeds 1 2 3 --output results/multi_seed.json
+```
+
+For the small causal language-model proof of concept:
+
+```bash
+python benchmarks/lm_prototype.py \
+    --data /path/to/pretraining_code.jsonl \
+    --max-bytes 10000000 --steps 50 --output results/lm_prototype.json
+```
+
+The causal LM is a separate experimental path. The full-sequence psMNIST model
+and its baseline remain unchanged.
+
+## Example
+
+```python
+import torch
+from astrohebbian import AstrocyteHebbianClassifier
+
+model = AstrocyteHebbianClassifier(
+    input_dim=1,
+    d_model=128,
+    seq_len=784,
+    num_heads=4,
+    v_levels=1,
+)
+
+pixels = torch.randn(8, 784, 1)
+logits = model(pixels)
+print(logits.shape)  # torch.Size([8, 10])
+```
+
+## Results
+
+Frozen three-seed psMNIST baseline (N=784, 60k train / 10k test, 6 epochs, RTX 3050, batch 64):
+
+| Model | 3-seed mean test acc | Peak VRAM |
+| --- | ---: | ---: |
+| AstroHebbian Pure SNN | **86.82% ± 2.48%** | ~1004 MB |
+| Transformer (dense O(N²)) | 77.80% ± 3.06% | ~1565 MB |
+
+The SNN beats the dense baseline by **+9.03 pts accuracy** at **~36% lower peak VRAM**.
+Mean runtime is not a headline figure: seed 2 was a large hardware/runtime outlier,
+so only accuracy and memory are claimed as reliable. Full per-seed data:
+`results/multi_seed.json`; the N-scaling memory crossover is in `results/n_scaling.png`.
+
+## Project status
+
+This is the clean standalone starting point for production hardening. The current attention is full-sequence rather than causal or streaming. Results are tracked in `docs/benchmark_baseline.md`.

astrohebbian-0.1.0/astrohebbian/__init__.py

@@ -0,0 +1,23 @@
+"""Standalone Astrocyte-Hebbian spiking Transformer components."""
+
+from .causal import CausalAstrocyteHebbianAttention
+from .language_model import CausalAstrocyteLanguageModel
+from .model import (
+    AstrocyteHebbianAttention,
+    AstrocyteHebbianBlock,
+    AstrocyteHebbianClassifier,
+    SpikingFFN,
+    SurrogateHeaviside,
+    spike_fn,
+)
+
+__all__ = [
+    "AstrocyteHebbianAttention",
+    "AstrocyteHebbianBlock",
+    "AstrocyteHebbianClassifier",
+    "SpikingFFN",
+    "SurrogateHeaviside",
+    "spike_fn",
+    "CausalAstrocyteHebbianAttention",
+    "CausalAstrocyteLanguageModel",
+]

astrohebbian-0.1.0/astrohebbian/benchmark.py

@@ -0,0 +1,227 @@
+"""Reproducible psMNIST benchmark for Astrocyte-Hebbian attention."""
+
+import argparse
+import json
+import math
+import random
+import sys
+import time
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader, Subset
+from torchvision import datasets, transforms
+
+ROOT = Path(__file__).resolve().parent.parent
+DATA_DIR = ROOT / "data"
+CKPT_DIR = ROOT / "checkpoints"
+
+sys.path.insert(0, str(ROOT))
+from astrohebbian.model import AstrocyteHebbianClassifier  # noqa: E402
+
+
+class TransformerClassifier(nn.Module):
+    def __init__(self, d_in, d_model, seq_len, num_classes, num_heads):
+        super().__init__()
+        self.proj_in = nn.Linear(d_in, d_model)
+        self.pos_encoder = nn.Parameter(torch.randn(1, seq_len, d_model) * 0.02)
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=d_model, nhead=num_heads, batch_first=True
+        )
+        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=1)
+        self.classifier = nn.Linear(d_model, num_classes)
+
+    def forward(self, x):
+        x = self.proj_in(x) + self.pos_encoder[:, : x.shape[1], :]
+        return self.classifier(self.transformer(x)[:, -1, :])
+
+
+def set_seed(seed):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+        torch.backends.cudnn.benchmark = False
+    torch.backends.cudnn.deterministic = True
+
+
+def warmup_cosine(step, warmup_steps, total_steps):
+    if step < warmup_steps:
+        return step / max(1, warmup_steps)
+    progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)
+    return 0.5 * (1.0 + math.cos(math.pi * progress))
+
+
+@torch.no_grad()
+def evaluate(model, test_loader, device, seq_len):
+    model.eval()
+    correct = 0
+    total = 0
+    for data, target in test_loader:
+        data = data.view(data.size(0), seq_len, 1).to(device)
+        target = target.to(device)
+        correct += model(data).argmax(dim=1).eq(target).sum().item()
+        total += data.size(0)
+    return 100.0 * correct / total
+
+
+def train_model(model, name, train_loader, test_loader, device, epochs, seq_len):
+    print(f"\n--- Training {name} ---", flush=True)
+    model = model.to(device)
+    optimizer = optim.Adam(model.parameters(), lr=1e-3)
+    criterion = nn.CrossEntropyLoss()
+    steps_per_epoch = len(train_loader)
+    scheduler = optim.lr_scheduler.LambdaLR(
+        optimizer,
+        lambda step: warmup_cosine(step, steps_per_epoch, steps_per_epoch * epochs),
+    )
+    if device.type == "cuda":
+        torch.cuda.reset_peak_memory_stats(device)
+    start_time = time.perf_counter()
+    for epoch in range(epochs):
+        model.train()
+        correct = 0
+        total = 0
+        for data, target in train_loader:
+            data = data.view(data.size(0), seq_len, 1).to(device, non_blocking=True)
+            target = target.to(device, non_blocking=True)
+            optimizer.zero_grad(set_to_none=True)
+            output = model(data)
+            loss = criterion(output, target)
+            loss.backward()
+            nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+            optimizer.step()
+            scheduler.step()
+            correct += output.argmax(dim=1).eq(target).sum().item()
+            total += data.size(0)
+        print(
+            f"Epoch {epoch + 1}/{epochs} | train acc: {100.0 * correct / total:.2f}%"
+            f" | lr: {scheduler.get_last_lr()[0]:.6f}",
+            flush=True,
+        )
+
+    elapsed = time.perf_counter() - start_time
+    peak_memory = (
+        torch.cuda.max_memory_allocated(device) / (1024 * 1024)
+        if device.type == "cuda"
+        else 0.0
+    )
+    test_accuracy = evaluate(model, test_loader, device, seq_len)
+    result = {
+        "test_accuracy": test_accuracy,
+        "train_time_seconds": elapsed,
+        "peak_vram_mb": peak_memory,
+    }
+    print(
+        f"[{name}] Time: {elapsed:.2f}s | VRAM: {peak_memory:.2f} MB"
+        f" | TEST acc: {test_accuracy:.2f}%",
+        flush=True,
+    )
+    return result
+
+
+def build_parser():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--device", choices=("auto", "cpu", "cuda"), default="auto")
+    parser.add_argument("--epochs", type=int, default=6)
+    parser.add_argument("--batch-size", type=int, default=64)
+    parser.add_argument("--test-batch-size", type=int, default=512)
+    parser.add_argument("--d-model", type=int, default=128)
+    parser.add_argument("--num-heads", type=int, default=4)
+    parser.add_argument("--num-layers", type=int, default=1)
+    parser.add_argument("--v-levels", type=int, default=1)
+    parser.add_argument("--train-size", type=int, default=None)
+    parser.add_argument("--test-size", type=int, default=None)
+    parser.add_argument("--output", type=Path, default=None, help="Write results to JSON")
+    parser.add_argument(
+        "--save-checkpoints", action="store_true", help="Save model state dictionaries"
+    )
+    return parser
+
+
+def select_device(requested):
+    if requested == "cuda" and not torch.cuda.is_available():
+        raise RuntimeError("CUDA was requested but is not available")
+    if requested == "auto":
+        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    return torch.device(requested)
+
+
+def main(argv=None):
+    args = build_parser().parse_args(argv)
+    if args.epochs <= 0 or args.batch_size <= 0 or args.test_batch_size <= 0:
+        raise ValueError("epochs and batch sizes must be positive")
+    set_seed(args.seed)
+    device = select_device(args.device)
+    seq_len = 784
+    transform = transforms.Compose(
+        [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
+    )
+    train_dataset = datasets.MNIST(str(DATA_DIR), train=True, download=True, transform=transform)
+    test_dataset = datasets.MNIST(str(DATA_DIR), train=False, download=True, transform=transform)
+    if args.train_size is not None:
+        train_dataset = Subset(train_dataset, range(min(args.train_size, len(train_dataset))))
+    if args.test_size is not None:
+        test_dataset = Subset(test_dataset, range(min(args.test_size, len(test_dataset))))
+    generator = torch.Generator().manual_seed(args.seed)
+    train_loader = DataLoader(
+        train_dataset, batch_size=args.batch_size, shuffle=True,
+        generator=generator, pin_memory=device.type == "cuda"
+    )
+    test_loader = DataLoader(test_dataset, batch_size=args.test_batch_size, shuffle=False)
+
+    print(f"BENCHMARK | seed: {args.seed} | device: {device}", flush=True)
+    print(
+        f"psMNIST (N={seq_len}) | train: {len(train_dataset)} | test: {len(test_dataset)}"
+        f" | {args.epochs} epochs | warmup+cosine | clip 1.0",
+        flush=True,
+    )
+    results = {
+        "config": {
+            "seed": args.seed, "device": str(device), "epochs": args.epochs,
+            "batch_size": args.batch_size, "test_batch_size": args.test_batch_size,
+            "d_model": args.d_model, "num_heads": args.num_heads,
+            "num_layers": args.num_layers, "v_levels": args.v_levels,
+            "train_size": len(train_dataset), "test_size": len(test_dataset),
+        },
+        "models": {},
+    }
+
+    snn = AstrocyteHebbianClassifier(
+        d_model=args.d_model, num_heads=args.num_heads,
+        num_layers=args.num_layers, seq_len=seq_len, v_levels=args.v_levels
+    )
+    results["models"]["AstroHebbian Pure SNN"] = train_model(
+        snn, "AstroHebbian Pure SNN", train_loader, test_loader, device, args.epochs, seq_len
+    )
+    if args.save_checkpoints:
+        CKPT_DIR.mkdir(exist_ok=True)
+        torch.save(snn.state_dict(), CKPT_DIR / "astrohebbian_pure_snn.pt")
+    del snn
+    if device.type == "cuda":
+        torch.cuda.empty_cache()
+
+    transformer = TransformerClassifier(1, args.d_model, seq_len, 10, args.num_heads)
+    results["models"]["Transformer (dense)"] = train_model(
+        transformer, "Transformer baseline", train_loader, test_loader, device, args.epochs, seq_len
+    )
+    if args.save_checkpoints:
+        torch.save(transformer.state_dict(), CKPT_DIR / "transformer_baseline.pt")
+
+    if args.output is not None:
+        args.output.parent.mkdir(parents=True, exist_ok=True)
+        args.output.write_text(json.dumps(results, indent=2) + "\n", encoding="utf-8")
+        print(f"Wrote results to {args.output}", flush=True)
+    return results
+
+
+if __name__ == "__main__":
+    main()
+
+if __name__ == "__main__":
+    main()

astrohebbian-0.1.0/astrohebbian/causal.py

@@ -0,0 +1,99 @@
+"""Causal, state-updating Astrocyte-Hebbian attention."""
+
+import torch
+import torch.nn as nn
+
+from .model import spike_fn
+
+
+class CausalAstrocyteHebbianAttention(nn.Module):
+    """Causal linear attention with a recurrent decayed Hebbian state.
+
+    Unlike the full-sequence attention module, each output only reads the
+    prefix ending at its current token. The state update is explicit so the
+    module can later be adapted to streaming execution.
+    """
+
+    def __init__(self, d_model=128, num_heads=4, v_levels=1):
+        super().__init__()
+        if d_model <= 0 or num_heads <= 0 or d_model % num_heads != 0:
+            raise ValueError("d_model must be positive and divisible by num_heads")
+        if v_levels != 1:
+            raise ValueError("the causal prototype currently supports v_levels=1 only")
+        self.d_model = d_model
+        self.num_heads = num_heads
+        self.head_dim = d_model // num_heads
+        self.q_proj = nn.Linear(d_model, d_model)
+        self.k_proj = nn.Linear(d_model, d_model)
+        self.v_proj = nn.Linear(d_model, d_model)
+        self.o_proj = nn.Linear(d_model, d_model)
+        self.decay_logit = nn.Parameter(torch.empty(d_model).uniform_(5.0, 8.0))
+        self.eps = 1e-6
+
+    def forward(self, x, implementation="recurrent"):
+        if x.ndim != 3 or x.shape[-1] != self.d_model:
+            raise ValueError(f"x must have shape (batch, sequence, {self.d_model})")
+        if x.shape[1] == 0:
+            raise ValueError("sequence length must be positive")
+        if implementation not in ("recurrent", "parallel"):
+            raise ValueError("implementation must be 'recurrent' or 'parallel'")
+        batch_size, sequence_length, _ = x.shape
+        query = spike_fn(self.q_proj(x)).view(
+            batch_size, sequence_length, self.num_heads, self.head_dim
+        ).transpose(1, 2)
+        key = spike_fn(self.k_proj(x)).view(
+            batch_size, sequence_length, self.num_heads, self.head_dim
+        ).transpose(1, 2)
+        value = spike_fn((torch.tanh(self.v_proj(x)) + 1.0) / 2.0 - 0.5).view(
+            batch_size, sequence_length, self.num_heads, self.head_dim
+        ).transpose(1, 2)
+        decay = torch.sigmoid(self.decay_logit).view(
+            1, self.num_heads, self.head_dim
+        ).to(x.dtype)
+        if implementation == "parallel":
+            positions = torch.arange(
+                sequence_length, device=x.device, dtype=x.dtype
+            ).view(1, 1, sequence_length, 1)
+            inverse_decay = decay.reciprocal()
+            powers = inverse_decay.unsqueeze(-2).pow(positions)
+            pair = key.unsqueeze(-1) * value.unsqueeze(-2)
+            trace = (
+                torch.cumsum(pair * powers.unsqueeze(-1), dim=2)
+                * decay.unsqueeze(-2).pow(positions).unsqueeze(-1)
+            )
+            mass = (
+                torch.cumsum(key * powers.squeeze(-1), dim=2)
+                * decay.unsqueeze(-2).pow(positions)
+            )
+            numerator = torch.matmul(query.unsqueeze(-2), trace).squeeze(-2)
+            denominator = (
+                (query * mass).sum(dim=-1, keepdim=True).clamp_min(self.eps)
+            )
+            output = numerator / denominator
+            output = output.transpose(1, 2).contiguous().view(
+                batch_size, sequence_length, self.d_model
+            )
+            return self.o_proj(output)
+
+        trace = torch.zeros(
+            batch_size, self.num_heads, self.head_dim, self.head_dim,
+            device=x.device, dtype=x.dtype,
+        )
+        mass = torch.zeros(
+            batch_size, self.num_heads, self.head_dim,
+            device=x.device, dtype=x.dtype,
+        )
+        outputs = []
+        for position in range(sequence_length):
+            key_position = key[:, :, position].unsqueeze(-1)
+            value_position = value[:, :, position].unsqueeze(-2)
+            trace = trace * decay.unsqueeze(-1) + key_position * value_position
+            mass = mass * decay + key[:, :, position]
+            numerator = torch.matmul(query[:, :, position].unsqueeze(-2), trace).squeeze(-2)
+            denominator = (
+                query[:, :, position] * mass
+            ).sum(dim=-1, keepdim=True).clamp_min(self.eps)
+            outputs.append(numerator / denominator)
+        output = torch.stack(outputs, dim=2)
+        output = output.transpose(1, 2).contiguous().view(batch_size, sequence_length, self.d_model)
+        return self.o_proj(output)

astrohebbian-0.1.0/astrohebbian/language_model.py

@@ -0,0 +1,40 @@
+"""Small causal language-model prototype using the causal attention module."""
+
+import torch
+import torch.nn as nn
+
+from .causal import CausalAstrocyteHebbianAttention
+from .model import SpikingFFN
+
+
+class CausalAstrocyteLanguageModel(nn.Module):
+    """Byte-level causal LM for small proof-of-concept experiments."""
+
+    def __init__(
+        self, vocab_size=256, d_model=64, seq_len=128, num_heads=4,
+        attention_implementation="parallel",
+    ):
+        super().__init__()
+        if vocab_size <= 0 or seq_len <= 0:
+            raise ValueError("vocab_size and seq_len must be positive")
+        self.seq_len = seq_len
+        if attention_implementation not in ("recurrent", "parallel"):
+            raise ValueError("attention_implementation must be recurrent or parallel")
+        self.attention_implementation = attention_implementation
+        self.embedding = nn.Embedding(vocab_size, d_model)
+        self.pos_encoder = nn.Parameter(torch.randn(1, seq_len, d_model) * 0.02)
+        self.norm1 = nn.LayerNorm(d_model)
+        self.attention = CausalAstrocyteHebbianAttention(d_model, num_heads)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.ffn = SpikingFFN(d_model, expansion=2)
+        self.lm_head = nn.Linear(d_model, vocab_size)
+
+    def forward(self, token_ids):
+        if token_ids.ndim != 2 or token_ids.shape[1] > self.seq_len:
+            raise ValueError("token_ids must have shape (batch, sequence <= seq_len)")
+        hidden = self.embedding(token_ids) + self.pos_encoder[:, : token_ids.shape[1], :]
+        hidden = hidden + self.attention(
+            self.norm1(hidden), implementation=self.attention_implementation
+        )
+        hidden = hidden + self.ffn(self.norm2(hidden))
+        return self.lm_head(hidden)

astrohebbian-0.1.0/astrohebbian/model.py

@@ -0,0 +1,174 @@
+"""Astrocyte-Hebbian spiking linear Transformer components.
+
+This package is intentionally standalone. It depends on PyTorch only and does
+not import or depend on Exact-SNN.
+"""
+
+import torch
+import torch.nn as nn
+
+
+class SurrogateHeaviside(torch.autograd.Function):
+    """Binary threshold in the forward pass with a fast-sigmoid backward pass."""
+
+    @staticmethod
+    def forward(ctx, input):
+        ctx.save_for_backward(input)
+        return (input > 0).to(input.dtype)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        (input,) = ctx.saved_tensors
+        alpha = 10.0
+        sigmoid = torch.sigmoid(alpha * input)
+        return grad_output * sigmoid * (1 - sigmoid) * alpha
+
+
+spike_fn = SurrogateHeaviside.apply
+
+
+class AstrocyteHebbianAttention(nn.Module):
+    """Multi-head spiking linear attention with learnable channel decay."""
+
+    def __init__(self, d_model=128, num_heads=4, v_levels=1):
+        super().__init__()
+        if d_model <= 0 or num_heads <= 0 or d_model % num_heads != 0:
+            raise ValueError("d_model must be positive and divisible by num_heads")
+        if v_levels <= 0:
+            raise ValueError("v_levels must be positive")
+
+        self.d_model = d_model
+        self.num_heads = num_heads
+        self.head_dim = d_model // num_heads
+        self.value_levels = int(v_levels)
+        self.q_proj = nn.Linear(d_model, d_model)
+        self.k_proj = nn.Linear(d_model, d_model)
+        self.v_proj = nn.Linear(d_model, d_model)
+        self.o_proj = nn.Linear(d_model, d_model)
+        self.register_buffer(
+            "thresholds",
+            torch.arange(1, self.value_levels + 1, dtype=torch.float32)
+            / (self.value_levels + 1),
+        )
+        self.decay_logit = nn.Parameter(torch.empty(d_model).uniform_(5.0, 8.0))
+        self.eps = 1e-6
+
+    def forward(self, x):
+        if x.ndim != 3:
+            raise ValueError("x must have shape (batch, sequence, d_model)")
+        batch_size, sequence_length, d_model = x.shape
+        if sequence_length == 0:
+            raise ValueError("sequence length must be positive")
+        if d_model != self.d_model:
+            raise ValueError(f"last dimension must be {self.d_model}")
+
+        query = spike_fn(self.q_proj(x))
+        key = spike_fn(self.k_proj(x))
+        value_unit = (torch.tanh(self.v_proj(x)) + 1.0) / 2.0
+        if self.value_levels == 1:
+            value = spike_fn(value_unit - self.thresholds[0])
+        else:
+            value_spikes = spike_fn(
+                value_unit.unsqueeze(-1) - self.thresholds.view(1, 1, 1, -1)
+            )
+            value = value_spikes.mean(dim=-1)
+
+        positions = torch.arange(
+            sequence_length - 1,
+            -1,
+            -1,
+            device=x.device,
+            dtype=x.dtype,
+        )
+        decay = torch.sigmoid(self.decay_logit).to(x.dtype)
+        weights = decay.unsqueeze(0).pow(positions.unsqueeze(1))
+        weighted_key = key * weights.unsqueeze(0)
+
+        query = query.view(
+            batch_size, sequence_length, self.num_heads, self.head_dim
+        ).transpose(1, 2)
+        weighted_key = weighted_key.view(
+            batch_size, sequence_length, self.num_heads, self.head_dim
+        ).transpose(1, 2)
+        value = value.view(
+            batch_size, sequence_length, self.num_heads, self.head_dim
+        ).transpose(1, 2)
+
+        key_value_trace = torch.matmul(weighted_key.transpose(-2, -1), value)
+        key_mass = weighted_key.sum(dim=-2, keepdim=True).transpose(-2, -1)
+        numerator = torch.matmul(query, key_value_trace)
+        denominator = torch.matmul(query, key_mass) + self.eps
+        output = numerator / denominator
+
+        output = output.transpose(1, 2).contiguous().view(batch_size, sequence_length, d_model)
+        return self.o_proj(output)
+
+
+class SpikingFFN(nn.Module):
+    """Feed-forward network with a binary hidden activation."""
+
+    def __init__(self, d_model=128, expansion=4):
+        super().__init__()
+        if d_model <= 0 or expansion <= 0:
+            raise ValueError("d_model and expansion must be positive")
+        self.fc1 = nn.Linear(d_model, d_model * expansion)
+        self.fc2 = nn.Linear(d_model * expansion, d_model)
+
+    def forward(self, x):
+        return self.fc2(spike_fn(self.fc1(x)))
+
+
+class AstrocyteHebbianBlock(nn.Module):
+    """Pre-norm Transformer block using Astrocyte-Hebbian attention."""
+
+    def __init__(self, d_model=128, num_heads=4, expansion=4, v_levels=1):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(d_model)
+        self.attention = AstrocyteHebbianAttention(d_model, num_heads, v_levels)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.ffn = SpikingFFN(d_model, expansion)
+
+    def forward(self, x):
+        x = x + self.attention(self.norm1(x))
+        return x + self.ffn(self.norm2(x))
+
+
+class AstrocyteHebbianClassifier(nn.Module):
+    """Sequence classifier built from standalone Astrocyte-Hebbian blocks."""
+
+    def __init__(
+        self,
+        input_dim=1,
+        d_model=128,
+        seq_len=784,
+        num_classes=10,
+        num_layers=1,
+        num_heads=4,
+        v_levels=1,
+    ):
+        super().__init__()
+        if seq_len <= 0 or num_classes <= 0 or num_layers <= 0:
+            raise ValueError("seq_len, num_classes, and num_layers must be positive")
+        if input_dim <= 0 or d_model <= 0 or num_heads <= 0:
+            raise ValueError("input_dim, d_model, and num_heads must be positive")
+        self.embedding = nn.Linear(input_dim, d_model)
+        self.pos_encoder = nn.Parameter(torch.randn(1, seq_len, d_model) * 0.02)
+        self.blocks = nn.ModuleList(
+            [
+                AstrocyteHebbianBlock(
+                    d_model, num_heads, v_levels=v_levels
+                )
+                for _ in range(num_layers)
+            ]
+        )
+        self.classifier = nn.Linear(d_model, num_classes)
+
+    def forward(self, x):
+        if x.ndim != 3 or x.shape[-1] != self.embedding.in_features:
+            raise ValueError("x must have shape (batch, sequence, input_dim)")
+        if x.shape[1] > self.pos_encoder.shape[1]:
+            raise ValueError("sequence length exceeds configured seq_len")
+        x = self.embedding(x) + self.pos_encoder[:, : x.shape[1], :]
+        for block in self.blocks:
+            x = block(x)
+        return self.classifier(x[:, -1, :])

astrohebbian-0.1.0/astrohebbian.egg-info/PKG-INFO

@@ -0,0 +1,150 @@
+Metadata-Version: 2.4
+Name: astrohebbian
+Version: 0.1.0
+Summary: Standalone Astrocyte-Hebbian spiking linear Transformer components
+Author: Sumith
+License: MIT License
+        
+        Copyright (c) 2026
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all
+        copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+        
+Project-URL: Homepage, https://github.com/Griffith-7/Snn-Transformer
+Project-URL: Source, https://github.com/Griffith-7/Snn-Transformer
+Project-URL: Issue Tracker, https://github.com/Griffith-7/Snn-Transformer/issues
+Keywords: spiking-neural-networks,transformer,attention,hebbian,snns,pytorch
+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.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.0
+Provides-Extra: test
+Requires-Dist: pytest>=7; extra == "test"
+Provides-Extra: benchmark
+Requires-Dist: torchvision>=0.15; extra == "benchmark"
+Requires-Dist: matplotlib>=3.7; extra == "benchmark"
+Provides-Extra: quality
+Requires-Dist: ruff>=0.6; extra == "quality"
+Dynamic: license-file
+
+# Astrocyte-Hebbian Spiking Transformer Plugin
+
+A standalone PyTorch plugin containing Astrocyte-Hebbian spiking linear-attention components. This project is intentionally independent from Exact-SNN; no Exact-SNN code, imports, or dependencies are included.
+
+## Scope
+
+This is a focused model plugin, not a general SNN framework. It provides:
+
+- `AstrocyteHebbianAttention`: multi-head linear attention using binary Q/K/V activations.
+- `AstrocyteHebbianBlock`: pre-norm Transformer-style block with a spiking FFN.
+- `AstrocyteHebbianClassifier`: ready-to-train sequence classifier.
+- `spike_fn`: binary Heaviside forward pass with surrogate gradients.
+
+The implementation avoids an `N x N` attention matrix by computing the `K^T V` trace. It still uses ordinary dense PyTorch tensors for projections, normalization, residual paths, and training.
+
+## What this is (and is not)
+
+- **Binary activation spikes**: the inter-layer signals (Q, K, V, and the FFN hidden activation) are binary Heaviside spikes (0/1).
+- **Surrogate gradients**: training uses a fast-sigmoid surrogate gradient through the spike threshold; it is not an exact spike-time gradient library.
+- **Dense PyTorch execution**: forward/backward run on ordinary dense GPU tensors. This is a CPU/GPU software package, not an event-driven neuromorphic-hardware implementation, and reported times/memory are wall-clock/FLOP measurements, not hardware energy.
+- **Full-sequence psMNIST mode**: the core `AstrocyteHebbianClassifier` (and the frozen baseline below) is full-sequence attention over N=784 pixels.
+- **Separate causal LM experimental mode**: a distinct, experimental causal path (`CausalAstrocyteLanguageModel`) is provided for small language-model proof-of-concept work only.
+
+This is a **focused model plugin, not a complete SNN framework**.
+
+## Install
+
+```bash
+pip install -e .
+```
+
+For development and tests:
+
+```bash
+pip install -e .[test]
+pytest -q
+```
+
+For the optional psMNIST benchmark:
+
+```bash
+pip install -e .[benchmark]
+python astrohebbian/benchmark.py
+```
+
+For a controlled three-seed summary:
+
+```bash
+python benchmarks/multi_seed.py --seeds 1 2 3 --output results/multi_seed.json
+```
+
+For the small causal language-model proof of concept:
+
+```bash
+python benchmarks/lm_prototype.py \
+    --data /path/to/pretraining_code.jsonl \
+    --max-bytes 10000000 --steps 50 --output results/lm_prototype.json
+```
+
+The causal LM is a separate experimental path. The full-sequence psMNIST model
+and its baseline remain unchanged.
+
+## Example
+
+```python
+import torch
+from astrohebbian import AstrocyteHebbianClassifier
+
+model = AstrocyteHebbianClassifier(
+    input_dim=1,
+    d_model=128,
+    seq_len=784,
+    num_heads=4,
+    v_levels=1,
+)
+
+pixels = torch.randn(8, 784, 1)
+logits = model(pixels)
+print(logits.shape)  # torch.Size([8, 10])
+```
+
+## Results
+
+Frozen three-seed psMNIST baseline (N=784, 60k train / 10k test, 6 epochs, RTX 3050, batch 64):
+
+| Model | 3-seed mean test acc | Peak VRAM |
+| --- | ---: | ---: |
+| AstroHebbian Pure SNN | **86.82% ± 2.48%** | ~1004 MB |
+| Transformer (dense O(N²)) | 77.80% ± 3.06% | ~1565 MB |
+
+The SNN beats the dense baseline by **+9.03 pts accuracy** at **~36% lower peak VRAM**.
+Mean runtime is not a headline figure: seed 2 was a large hardware/runtime outlier,
+so only accuracy and memory are claimed as reliable. Full per-seed data:
+`results/multi_seed.json`; the N-scaling memory crossover is in `results/n_scaling.png`.
+
+## Project status
+
+This is the clean standalone starting point for production hardening. The current attention is full-sequence rather than causal or streaming. Results are tracked in `docs/benchmark_baseline.md`.

astrohebbian-0.1.0/astrohebbian.egg-info/SOURCES.txt

@@ -0,0 +1,14 @@
+LICENSE
+README.md
+pyproject.toml
+astrohebbian/__init__.py
+astrohebbian/benchmark.py
+astrohebbian/causal.py
+astrohebbian/language_model.py
+astrohebbian/model.py
+astrohebbian.egg-info/PKG-INFO
+astrohebbian.egg-info/SOURCES.txt
+astrohebbian.egg-info/dependency_links.txt
+astrohebbian.egg-info/requires.txt
+astrohebbian.egg-info/top_level.txt
+tests/test_plugin.py

astrohebbian-0.1.0/astrohebbian.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

astrohebbian-0.1.0/astrohebbian.egg-info/requires.txt

@@ -0,0 +1,11 @@
+torch>=2.0
+
+[benchmark]
+torchvision>=0.15
+matplotlib>=3.7
+
+[quality]
+ruff>=0.6
+
+[test]
+pytest>=7

astrohebbian-0.1.0/astrohebbian.egg-info/top_level.txt

@@ -0,0 +1 @@
+astrohebbian

astrohebbian-0.1.0/pyproject.toml

@@ -0,0 +1,41 @@
+[build-system]
+requires = ["setuptools>=68"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "astrohebbian"
+version = "0.1.0"
+description = "Standalone Astrocyte-Hebbian spiking linear Transformer components"
+readme = "README.md"
+requires-python = ">=3.10"
+license = { file = "LICENSE" }
+authors = [{ name = "Sumith" }]
+keywords = ["spiking-neural-networks", "transformer", "attention", "hebbian", "snns", "pytorch"]
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "Intended Audience :: Science/Research",
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+dependencies = ["torch>=2.0"]
+
+[project.urls]
+Homepage = "https://github.com/Griffith-7/Snn-Transformer"
+Source = "https://github.com/Griffith-7/Snn-Transformer"
+"Issue Tracker" = "https://github.com/Griffith-7/Snn-Transformer/issues"
+
+[project.optional-dependencies]
+test = ["pytest>=7"]
+benchmark = ["torchvision>=0.15", "matplotlib>=3.7"]
+quality = ["ruff>=0.6"]
+
+[tool.setuptools.packages.find]
+include = ["astrohebbian*"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+pythonpath = ["."]

astrohebbian-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

astrohebbian-0.1.0/tests/test_plugin.py

@@ -0,0 +1,269 @@
+"""
+Tests for the Astrocyte-Hebbian spiking linear Transformer plugin.
+
+This suite mirrors the reference tests from
+`Astrocyte-Hebbian-Spiking-Linear-Transformer-Block/tests/test_astrohebbian.py`
+but is adapted to this plugin's API:
+
+* `AstrocyteHebbianBlock` exposes `self.attention` (not `self.attn`).
+* `AstrocyteHebbianAttention` exposes `self.num_heads` / `self.head_dim`
+  (not `self.h` / `self.dh`) and keeps the spike thresholds in the
+  `self.thresholds` buffer.
+
+Run:  pytest tests/ -v
+"""
+import torch
+import torch.nn as nn
+
+from astrohebbian import (
+    AstrocyteHebbianAttention,
+    AstrocyteHebbianBlock,
+    AstrocyteHebbianClassifier,
+    CausalAstrocyteHebbianAttention,
+    CausalAstrocyteLanguageModel,
+    SpikingFFN,
+    spike_fn,
+)
+
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+
+def test_attention_rejects_empty_or_mismatched_sequences():
+    attention = AstrocyteHebbianAttention(d_model=16, num_heads=4)
+    for inputs in (torch.randn(2, 0, 16), torch.randn(2, 4, 8)):
+        try:
+            attention(inputs)
+        except ValueError:
+            pass
+        else:
+            raise AssertionError("invalid attention input was accepted")
+
+
+def test_classifier_state_round_trip_is_deterministic():
+    torch.manual_seed(12)
+    model = AstrocyteHebbianClassifier(
+        input_dim=1, d_model=16, seq_len=8, num_classes=3, num_heads=4
+    ).eval()
+    inputs = torch.randn(2, 8, 1)
+    expected = model(inputs)
+    restored = AstrocyteHebbianClassifier(
+        input_dim=1, d_model=16, seq_len=8, num_classes=3, num_heads=4
+    ).eval()
+    restored.load_state_dict(model.state_dict())
+    assert torch.equal(expected, restored(inputs))
+
+# ---------------------------------------------------------------------------
+# Surrogate-heavy-side spike function
+# ---------------------------------------------------------------------------
+
+class TestSpikeFn:
+    def test_forward_is_binary(self):
+        x = torch.randn(64, 32, device=DEVICE) * 3.0
+        s = spike_fn(x)
+        assert set(s.unique().tolist()) <= {0.0, 1.0}, "spike_fn must output 0/1"
+        assert (s == (x > 0).float()).all()
+
+    def test_backward_is_finite_and_nonzero(self):
+        x = torch.randn(64, 32, device=DEVICE, requires_grad=True)
+        s = spike_fn(x)
+        s.sum().backward()
+        assert torch.isfinite(x.grad).all(), "surrogate gradient must be finite"
+        assert x.grad.abs().sum() > 0, "surrogate gradient must not vanish everywhere"
+
+
+# ---------------------------------------------------------------------------
+# Attention mechanism (the core novel contribution)
+# ---------------------------------------------------------------------------
+
+class TestAttention:
+    def test_output_shape(self):
+        attn = AstrocyteHebbianAttention(d_model=128, num_heads=4).to(DEVICE)
+        x = torch.randn(2, 784, 128, device=DEVICE)
+        out = attn(x)
+        assert out.shape == x.shape
+
+    def test_no_n_squared_attention_matrix(self):
+        """The whole point: never materialize an N x N attention/score matrix.
+
+        We assert this empirically by checking that for seq_len N, the only
+        large intermediate tensors have size O(N*d) or O(d*d), never O(N*N).
+        We instrument by running with an enormous N and confirming no O(N^2)
+        allocation dominates memory (would blow up VRAM).
+        """
+        attn = AstrocyteHebbianAttention(d_model=64, num_heads=2).to(DEVICE)
+        # N = 16k tokens: an O(N^2) matrix would need 16k^2 * 4B = 1 GB just
+        # for one float32 attention map. O(N*d) stays at 16k*64*4B = 4 MB.
+        x = torch.randn(1, 4000, 64, device=DEVICE)
+        # If anything materialized a (4000, 4000) tensor, peak memory would
+        # jump. We just assert the forward runs without OOM and is finite.
+        out = attn(x)
+        assert out.shape == x.shape
+        assert torch.isfinite(out).all()
+
+    def test_q_k_binary_v_binary_at_levels_one(self):
+        """At v_levels=1, V is also binary: the strictest all-spiking regime."""
+        attn = AstrocyteHebbianAttention(d_model=128, num_heads=4, v_levels=1).to(DEVICE)
+        x = torch.randn(3, 64, 128, device=DEVICE)
+        # Reconstruct V from v_proj and the plugin's `self.thresholds` buffer.
+        with torch.no_grad():
+            pre_q = spike_fn(attn.q_proj(x))
+            pre_k = spike_fn(attn.k_proj(x))
+            u01 = (torch.tanh(attn.v_proj(x)) + 1.0) / 2.0
+            v = spike_fn(u01 - attn.thresholds.item())
+        assert set(pre_q.unique().tolist()) <= {0.0, 1.0}
+        assert set(pre_k.unique().tolist()) <= {0.0, 1.0}
+        assert set(v.unique().tolist()) <= {0.0, 1.0}
+
+    def test_gradient_flows_to_all_learnable_params(self):
+        attn = AstrocyteHebbianAttention(d_model=64, num_heads=2).to(DEVICE)
+        x = torch.randn(4, 64, 64, device=DEVICE, requires_grad=True)
+        out = attn(x)
+        out.sum().backward()
+        for name, p in attn.named_parameters():
+            assert p.grad is not None, f"no gradient for {name}"
+            assert torch.isfinite(p.grad).all(), f"non-finite gradient for {name}"
+
+    def test_separate_keys_change_output(self):
+        """Sanity: different inputs produce different outputs (no collapse)."""
+        attn = AstrocyteHebbianAttention(d_model=64, num_heads=2)
+        a = torch.randn(2, 32, 64)
+        b = torch.randn(2, 32, 64)
+        out_a = attn(a)
+        out_b = attn(b)
+        assert (out_a - out_b).abs().max() > 1e-6
+
+    def test_v_levels_multi_level(self):
+        """v_levels>1 quantizes V to L levels (still finite, well-behaved)."""
+        for L in (2, 4):
+            attn = AstrocyteHebbianAttention(d_model=64, num_heads=2, v_levels=L).to(DEVICE)
+            x = torch.randn(2, 32, 64, device=DEVICE)
+            out = attn(x)
+            assert out.shape == x.shape
+            assert torch.isfinite(out).all()
+
+
+# ---------------------------------------------------------------------------
+# Block / FFN / Classifier
+# ---------------------------------------------------------------------------
+
+class TestBlock:
+    def test_block_shape_and_finite(self):
+        block = AstrocyteHebbianBlock(d_model=128, num_heads=4).to(DEVICE)
+        x = torch.randn(2, 64, 128, device=DEVICE)
+        out = block(x)
+        assert out.shape == x.shape
+        assert torch.isfinite(out).all()
+
+    def test_ffn_hidden_is_binary(self):
+        ffn = SpikingFFN(d_model=64, expansion=4).to(DEVICE)
+        x = torch.randn(2, 32, 64, device=DEVICE)
+        with torch.no_grad():
+            hidden = spike_fn(ffn.fc1(x))
+        assert set(hidden.unique().tolist()) <= {0.0, 1.0}
+
+    def test_stackable_blocks(self):
+        """num_layers>1 should run cleanly (the advertised deep variant)."""
+        model = AstrocyteHebbianClassifier(
+            d_model=64, num_heads=2, num_layers=2, v_levels=1
+        ).to(DEVICE)
+        x = torch.randn(2, 784, 1, device=DEVICE)
+        logits = model(x)
+        assert logits.shape == (2, 10)
+        assert torch.isfinite(logits).all()
+
+
+def test_causal_attention_does_not_read_future_tokens():
+    torch.manual_seed(4)
+    attention = CausalAstrocyteHebbianAttention(d_model=16, num_heads=4).eval()
+    inputs = torch.randn(1, 8, 16)
+    changed = inputs.clone()
+    changed[:, 7, :] += 50
+    first = attention(inputs)[:, :7]
+    changed_first = attention(changed)[:, :7]
+    assert torch.equal(first, changed_first)
+
+
+def test_parallel_causal_attention_matches_recurrent_path():
+    torch.manual_seed(6)
+    attention = CausalAstrocyteHebbianAttention(d_model=16, num_heads=4).eval()
+    inputs = torch.randn(2, 8, 16)
+    recurrent = attention(inputs, implementation="recurrent")
+    parallel = attention(inputs, implementation="parallel")
+    assert torch.allclose(recurrent, parallel, atol=1e-5, rtol=1e-5)
+
+
+def test_causal_language_model_trains_on_tiny_batch():
+    torch.manual_seed(5)
+    model = CausalAstrocyteLanguageModel(vocab_size=32, d_model=16, seq_len=8)
+    inputs = torch.randint(0, 32, (2, 8))
+    targets = torch.roll(inputs, shifts=-1, dims=1)
+    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
+    losses = []
+    for _ in range(3):
+        optimizer.zero_grad(set_to_none=True)
+        loss = torch.nn.functional.cross_entropy(
+            model(inputs).reshape(-1, 32), targets.reshape(-1)
+        )
+        loss.backward()
+        optimizer.step()
+        losses.append(loss.item())
+    assert losses[-1] < losses[0]
+
+
+class TestClassifier:
+    def test_classifier_shapes(self):
+        model = AstrocyteHebbianClassifier(d_model=128, num_heads=4, num_layers=1).to(DEVICE)
+        x = torch.randn(5, 784, 1, device=DEVICE)
+        logits = model(x)
+        assert logits.shape == (5, 10)
+
+    def test_training_step_reduces_loss_and_finite_grads(self):
+        """End-to-end: a few optimizer steps should reduce the loss without NaNs."""
+        model = AstrocyteHebbianClassifier(
+            d_model=64, num_heads=2, num_layers=1, v_levels=1
+        ).to(DEVICE)
+        opt = torch.optim.Adam(model.parameters(), lr=1e-3)
+        crit = nn.CrossEntropyLoss()
+        x = torch.randn(8, 784, 1, device=DEVICE)
+        y = torch.randint(0, 10, (8,), device=DEVICE)
+
+        losses = []
+        for _ in range(3):
+            opt.zero_grad()
+            out = model(x)
+            loss = crit(out, y)
+            loss.backward()
+            for p in model.parameters():
+                assert p.grad is None or torch.isfinite(p.grad).all()
+            opt.step()
+            losses.append(loss.item())
+        assert losses[-1] < losses[0], f"loss did not decrease: {losses}"
+
+
+# ---------------------------------------------------------------------------
+# Hebbian trace / key-mass normalization semantics
+# ---------------------------------------------------------------------------
+
+class TestMechanism:
+    def test_key_mass_normalization_prevents_blowup(self):
+        """denom = Q . sum_K + eps keeps the output bounded even when the
+        raw numerator is large."""
+        attn = AstrocyteHebbianAttention(d_model=64, num_heads=2).to(DEVICE)
+        # Push k and v to be large/constant
+        B, N, d = 2, 64, 64
+        x = torch.ones(B, N, d, device=DEVICE)
+        out = attn(x)
+        assert torch.isfinite(out).all()
+        assert out.abs().max() < 100.0, f"key-mass norm let output blow up: {out.abs().max()}"
+
+    def test_complexity_claim_reading(self):
+        """The trace contract K^T V is (d,d), never (N,N) -- structural check."""
+        import inspect
+        src = inspect.getsource(AstrocyteHebbianAttention.forward)
+        # The plugin computes the trace via `weighted_key.transpose(-2, -1)`
+        # times `value` -- a (d,d) product per head, never an (N,N) item.
+        assert "matmul(weighted_key.transpose(-2, -1), value)" in src, "must compute K^T V"
+        # Ensure there is no softmax / score matrix anywhere in the module
+        assert "softmax" not in src
+        assert "self_attn" not in src.lower() or "scaled_dot" not in src.lower()