Skip to content

Dense Chrono-Spike LM training and prompt retrieval

This document summarizes the current implementation of SpikingEvoTextLM, ChronoSpikeAttention, MetaSTDP, AEG, and the prompt-time retrieval path based on the verified code.

1. Core implementation contract

The dense text path is no longer a single simplified block. It is a system with the following integrated contract:

  • token-to-spike conversion via TASEncoderDecoder
  • causal temporal attention via ChronoSpikeAttention
  • configurable output neuron layers (EvoLIF, Izhikevich, LIF)
  • optional AEG and MetaSTDP adaptation
  • prompt-time retrieval via SNNRAGHybrid and dense fallback

The safety defaults in the current code are:

  • EVOSPIKENET_FORCE_HARD_SPIKE_PARITY=true
  • EVOSPIKENET_TRAIN_CONTINUOUS_RELAXATION=false
  • any continuous-relaxation mode requires explicit opt-in

2. Neuron structure and types

LIF

  • LIFNeuronLayer uses integer membrane potential, threshold, and leak states
  • inference is strictly hard-spike
  • training uses evo_lif_spike_with_ste() to preserve a stable identity gradient

EvoLIF

  • the EvoLIF path in SpikingEvoTextLM uses LIFNeuronLayer and caps the scale with resolve_evo_lif_scale_factor()
  • EVOSPIKENET_EVO_LIF_SCALE_FACTOR is constrained to a maximum of 100.0

Izhikevich

  • IzhikevichNeuronLayer keeps a float state for \(v\) and \(u\)
  • it fires when the threshold is crossed and resets after spike
  • izhikevich_spike_with_ste() preserves hard inference while allowing task gradients to flow

Training-side stability helper

  • spike_activation_with_residual_gradient() adds a bounded residual to prevent zero-spike collapse in short temporal windows

3. Current ChronoSpikeAttention model

The causal mask is

\[ M(t, t') = \exp\left(-\max(0, t - t') / \tau\right) \]

with tau learnable when learnable_tau=True and optionally per-head when per_head_tau=True.

The layer accepts neuron_type values LIF, EvoLIF, and Izhikevich, and converts its attention output into a spike train via output_lif before the next layer consumes it.

4. Data flow in SpikingEvoTextLM

The current forward path is:

  1. token encoding to embeddings and spike trains
  2. optional AEG importance gating
  3. transformer block processing
  4. time aggregation of spike activity
  5. output_potential_sum accumulation
  6. readout merge of mean activity + deep output + direct embedding shortcut
  7. final logits from the direct/deep blend

The final readout mixes the two streams as

\[ \text{logits} = \sigma(\alpha) \cdot \text{direct\_logits} + (1 - \sigma(\alpha)) \cdot \text{deep\_logits} \]

with \(\alpha = \text{sigmoid}(\text{readout\_direct\_logit\_mix})\).

5. Training convergence and reward design

The convergence profile is resolved in examples/train_spiking_evospikenet_lm.py by _resolve_convergence_profile().

  • default: stable_baseline
  • stable_converge disables AEG and MetaSTDP to favor stable training
  • aggressive_convergence tunes reward clipping and learning-rate style behavior
  • explicit environment overrides remain the highest-priority control

The reward path supports raw, clipped_raw, and ema_delta modes, with the EMA form being the default stable option:

\[ \text{reward} = -\Delta \text{EMA}(\text{loss}) \]

This dampens noisy loss spikes and keeps the adaptation signal more stable than a plain negative loss.

6. Corpus flow and chunked training

The training corpus pipeline is:

  1. get_training_corpus(args)
  2. _apply_rag_japanese_preprocess() for Japanese normalization/chunking
  3. _iter_preprocessed_hf_japanese_wikipedia_chunks()
  4. _tokenize_corpus_in_chunks()
  5. _build_next_token_dataset()

The design intentionally avoids loading the entire corpus into memory at once. It instead creates bounded text chunks, tokenizes those chunks, and then builds a next-token dataset in a shift-by-one structure.

7. Prompt-time retrieval and fallback logic

The retrieval path in evospikenet/snn_rag.py uses SNNRAGHybrid to encode queries with both spike and dense features.

  • KnowledgeGraphIntegrator builds a lightweight on-device document graph
  • ChronoSpikeAttention contributes temporal correlation scoring
  • low-confidence spike retrieval can trigger dense fallback using vector similarity
  • spike_fallback_threshold governs the fallback gate

This is a hybrid retrieval strategy: spike-based retrieval remains the primary path, while dense retrieval is used as a fallback when the spike signal is too weak or the query is unknown.

8. Summary

The current implementation is best understood as a production-stabilized spiking language model rather than a purely biologically literal architecture. It keeps hard-spike inference as the default path, allows controlled approximations only when explicitly requested, blends direct and deep readout signals for stable learning, and combines prompt-time retrieval with dense fallback for low-confidence queries.