Lattice_Spike

LatticeSpike (reference implementation)

For the versioned v5.3 development protocol, named comparison suites, locked model selection, and paired statistics, see docs/paper_ready_evaluation.md.

PyTorch implementation of LatticeSpike, the architecture in lattice_snn_hallucination.tex: online LLM hallucination detection by combining a Galois-connection Box-embedding lattice with a box-volume-modulated recurrent spiking neural network.

Install

pip install -r requirements.txt

(Only torch, numpy, scikit-learn are required. transformers is optional, for real-LLM signal extraction.)

For development and the full unit suite:

pip install -r requirements-dev.txt
PYTHONPATH=. python -m pytest -q

Run the end-to-end demo

python train.py                 # unsupervised train + evaluate on synthetic signals
python train.py --surrogate sigmoid --contaminate 0.15 --trim 0.2

The demo trains unsupervised on (mostly) grounded synthetic sequences and reports token-level AUPRC, lead time, and per-channel gate attribution.

The consolidated real-data diagnostic results and the validated detector changes are documented in docs/latticespike_experiment_results.md.

Module map (→ paper section)

The six input channels

RSNN input dim = 2d + k + 5.

Evaluating baselines and visualizing results

Run the baseline detectors (EPR and Lookback Lens) on a JSONL of prompt/answer pairs:

python scripts/eval_baselines_dgx.py \
  --model Qwen/Qwen2.5-0.5B-Instruct \
  --input examples/tiny_token_eval.jsonl \
  --output output/predictions.jsonl

Then convert the predictions to an interactive HTML visualization:

python scripts/visualize.py output/predictions.jsonl
# writes output/predictions_viz.html

Optional flags:

# filter to a single method
python scripts/visualize.py output/predictions.jsonl --method epr

# limit to the first N cards
python scripts/visualize.py output/predictions.jsonl --limit 20

# specify output path
python scripts/visualize.py output/predictions.jsonl --output output/my_viz.html

The visualizer accepts both raw eval_baselines_dgx.py output and pre-converted visualizer records — format is detected automatically from the first row.

Training on real LLM signals from raw data

train.py --data-source llm runs the full pipeline in one command: raw labelled JSONL → per-token LLM signal extraction (llm_signals.py) → LatticeSpike training → output/rsnn_output.pt. (The synthetic-generator pipeline is --data-source synthetic, the default — see the module docstring in train.py for its own flags and AUPRC/lead-time eval.)

# signal extraction isn't cached, so start small and iterate on --limit
python train.py --data-source llm --limit 20 --epochs 1 --device cpu

# then scale up
python train.py --data-source llm \
  --llm-data data/auto_labelled_llama8b.jsonl \
  --llm-model Qwen/Qwen2.5-0.5B-Instruct \
  --epochs 5 --device cuda    # or --device cpu

What each stage does:

1. Load raw data — data/auto_labelled_llama8b.jsonl, one record per line with (among other fields) generated_answer (the text to extract signals from — override with --text-field), token_labels (per-token G/U/H, auto-labelled by an NLI judge with a heuristic fallback — see auto_label_method per record), and t_star_index (hallucination onset token, or None).
2. Extract LLM signals — for each record, llm_signals.extract_signals runs one forward pass of --llm-model (default Qwen/Qwen2.5-0.5B-Instruct) with hidden states + attentions, and produces per-token: h_box (mid-layer hidden state), delta (per-layer KL between consecutive-token next-token distributions, via the logit lens), attn_ent (mean attention entropy), erank (effective rank of stacked hiddens), and topk_probs/topk_ids (top---topk next-token candidates). This is the expensive step (one LLM forward pass per record).
3. Train — batches are padded to a common length (collate_batch) and fed through the full LatticeSpike model (box encoder + RSNN) for --epochs passes with CombinedLoss.
4. Save — a final no-grad eval pass over every record collects the RSNNOutput (spikes, u, u_pred, gates, rho, anomaly, burst) plus sequence metadata, and saves it to <--out-dir>/rsnn_output.pt (default output/rsnn_output.pt), alongside rsnn_checkpoint.pt — a training checkpoint (model + optimizer state + args + final avg loss, load with torch.load(path, weights_only=False)["model"] for just the weights).

Run python train.py --help for the full flag list. Most flags are shared with the synthetic pipeline but take different defaults under --data-source llm (see LLM_MODE_DEFAULTS in train.py, also printed in --help) unless you override them explicitly.

Using a real LLM

from transformers import AutoModelForCausalLM, AutoTokenizer
from Lattice_Spike import LatticeSpike
from Lattice_Spike.llm_signals import extract_signals

tok = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B")
lm  = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-7B", output_hidden_states=True)

signals = extract_signals(lm, tok, "The capital of France is", topk=64)
meta = signals.pop("meta")

model = LatticeSpike(d_model=meta["d_model"], vocab_size=lm.config.vocab_size,
                     num_layers=meta["L"])
out = model(signals)          # out.rho is the per-token burst score

Note: box_token_emb is a box-space token table. For real use, tie it to (a learned projection of) the LLM embedding table so the concretization residual is computed in a consistent space; in the demo it is structured synthetically.