Skip to content

Sparse Evo-MemoryLM: Detailed Mathematics and Memory Design

Implementation status: Implemented. The reference implementation is EvoSpikeNet-Core/evospikenet/sparse_event_memory.py.
Scope: SparseEventMemoryLM, EventCSRLayer, TiedFactorizedVocabulary, SparseEventMemoryBlock, and EventMemoryExpert.
Out of scope: This is a separate, opt-in architecture. It does not replace the existing dense SpikingEvoTextLM / ChronoSpikeAttention / SpikingFFN path.

Contents

  1. Why this architecture is needed
  2. Implementation boundary
  3. Notation
  4. End-to-end data flow
  5. Factorized tied vocabulary
  6. INT8 CSR synapses
  7. EvoLIF-style event state
  8. Expert routing and low-rank adapters
  9. Loss, Adam, and local plasticity
  10. Actual memory layout and lifetime
  11. Sizing equations and dense-model comparison
  12. Training, inference, and generation
  13. Worked sizing example
  14. Constraints and operating guidance

Why this architecture is needed

Why dense language-model VRAM grows rapidly

A conventional dense Transformer-style block contains Q/K/V/output projections and an FFN. With hidden width \(D\), its dominant weights are approximately

\[ 4D^2 + (D\cdot4D + 4D\cdot D)=12D^2. \]

When FP32 Adam trains every weight, the minimum storage for weights, gradients, first moments, and second moments is approximately

\[ M_{\mathrm{Adam,min}}\approx16P\ \mathrm{bytes}, \]

where \(P\) is the trainable parameter count. This lower bound excludes activations, CUDA workspaces, and input/output logits.

EvoLIF neurons and softmax-free ChronoSpikeAttention do not remove that dominant term while dense Q/K/V/FFN weights and their Adam state remain present.

The separation used by Sparse Evo-MemoryLM

The model splits capacity into three storage classes:

Storage class Role Representation Update method
Primary memory Input and recurrent synapses Fixed-topology CSR with INT8 values Local Hebbian update
Small plastic subsystem Vocabulary basis, routers, adapters, LayerNorm Floating-point nn.Parameter tensors Adam
Dynamic state Membrane, previous spikes, activity records INT16 / temporary floating tensors Created and released per sequence call

The aim is not to make arbitrary capacity free. It is to remove large primary synapse matrices from Adam so that GPU memory does not scale with their gradients and optimizer moments.

Implementation boundary

The following statements describe the current source code, rather than a future hardware design.

  • CSR crow_indices, col_indices, and quantized_values are buffers, not nn.Parameter objects.
  • INT8 quantized_values are stored persistently, but forward() casts them to a floating CSR tensor for torch.sparse.mm.
  • Connectivity is created at initialization. Training does not add, prune, or rewire connections.
  • The router uses a mean-pooled sequence representation and selects Top-k experts. The Top-k index selection is discrete; routing is not differentiable across selection boundaries.
  • EventMemoryExpert iterates causally over tokens. It also collects each output in a Python list and stacks it, so sequence outputs and autograd information still scale with sequence length. The implementation is not constant-activation-memory recurrent execution.
  • generate() re-evaluates the whole prompt for every new token. There is no KV cache or persistent recurrent-state cache.
  • --ssl-task reconstruction, MetaSTDP, and AEG are not used by the sparse training branch.

Notation

Symbol Meaning
\(V\) Vocabulary size
\(D\) d_model, event-state width
\(R\) factor_rank, shared vocabulary factor rank
\(L\) num_transformer_blocks, sparse event block count
\(M\) num_experts per block
\(K\) router_top_k, experts executed per sequence
\(\rho\) connectivity, row-level CSR density
\(F\) CSR fan-in: \(\min(D,\max(1,\lceil\rho D\rceil))\)
\(A\) adapter_rank
\(B\) Batch size
\(S\) Sequence length
\(C\) Candidate count for sampled softmax
\(\theta\) event_threshold
\(\lambda\) event_leak, an integer coefficient with denominator 256
\(\eta_h\) plasticity_learning_rate

End-to-end data flow

flowchart TD
    A[Token IDs] --> B[INT8 codebook lookup]
    B --> C[Shared factor basis]
    C --> D[Sequence router]
    D --> E[Top-k sparse event experts]
    E --> F[INT8 CSR input and recurrent synapses]
    F --> G[INT16 membrane and spike state]
    G --> H[LayerNorm output state]
    H --> I[Sampled vocabulary projection]
    I --> J[Cross-entropy]
    J --> K[Adam on small trainable tensors]
    K --> L[Local Hebbian update of INT8 CSR values]

    classDef fixed fill:#e8f3ff,stroke:#1976d2,color:#102a43
    classDef trainable fill:#e8f5e9,stroke:#2e7d32,color:#102a43
    classDef state fill:#fff3e0,stroke:#ef6c00,color:#102a43
    class B,F fixed
    class C,D,H,I,K trainable
    class G,L state

Each block receives the same input sequence, executes only its selected experts, then adds their outputs using router gates. The result becomes the input of the next block.

Factorized tied vocabulary

Stored representation

The vocabulary layer has a fixed INT8 codebook

\[ Q\in\mathbb{Z}_8^{V\times R} \]

and an Adam-trained shared basis

\[ E\in\mathbb{R}^{R\times D}. \]

For token \(w\), the input embedding is

\[ \mathbf{x}_w=Q_wE\in\mathbb{R}^{D}. \]

The same \(Q,E\) serve input embedding and output projection, avoiding separate dense \(V\times D\) input and output matrices.

Output projection

For hidden state \(\mathbf{h}\in\mathbb{R}^{D}\), compute latent vocabulary coordinates

\[ \mathbf{z}=\frac{E\mathbf{h}}{\sqrt{R}}\in\mathbb{R}^{R}. \]

For candidate set \(\mathcal{C}\), the logits are

\[ \ell_c=\mathbf{z}^{\mathsf T}Q_c,\qquad c\in\mathcal{C}. \]
  • forward() uses \(\mathcal{C}=\{0,\ldots,V-1\}\) and materializes full-vocabulary logits.
  • sampled_cross_entropy() uses all unique targets plus random unique negatives, limiting \(|\mathcal{C}|=C\).

The training loss is candidate-set cross entropy:

\[ \mathcal{L}_{\mathrm{sampled}}=-\frac{1}{BS}\sum_{b,t} \log\frac{\exp(\ell_{y_{b,t}})}{\sum_{c\in\mathcal{C}}\exp(\ell_c)}. \]

It is not exactly the same objective as a full-vocabulary softmax.

INT8 CSR synapses

Connectivity topology

Each EventCSRLayer(D,D) has \(F\) connections in every output row:

\[ F=\min(D,\max(1,\lceil\rho D\rceil)),\qquad N_{\mathrm{conn}}=DF. \]

For row \(i\), an offset \(o_i\) and a stride \(s_i\) chosen to be coprime with the input width give columns

\[ j_{i,p}=(o_i+s_ip)\bmod D,\qquad p\in\{0,\ldots,F-1\}. \]

This produces a distributed fixed neighborhood without materializing a dense mask. Stored values are

\[ w_{i,p}\in\{-127,\ldots,127\}\subset\mathbb{Z}_8. \]

For input \(\mathbf{x}\), the actual computation casts values to floating point:

\[ [\mathcal{S}(\mathbf{x})]_i=\sum_{p=0}^{F-1}\mathrm{float}(w_{i,p})x_{j_{i,p}}. \]

It is not an all-INT8 neuromorphic multiply-accumulate kernel.

Persistent storage

A CSR layer uses INT64 crow_indices, INT64 col_indices, and INT8 values. Its persistent storage is therefore

\[ M_{\mathrm{CSR,layer}}=8(D+1)+8DF+DF=8(D+1)+9DF\ \mathrm{bytes}. \]

This excludes temporary float CSR values and sparse-kernel workspaces.

EvoLIF-style event state

For expert \(m\) and time \(t\), let the input be \(\mathbf{x}_t\) and the previous spike be \(\mathbf{s}_{t-1}\). The low-rank adapter is

\[ \mathbf{a}_t=(\mathbf{x}_tA_{\downarrow})A_{\uparrow}, \quad A_{\downarrow}\in\mathbb{R}^{D\times A}, \quad A_{\uparrow}\in\mathbb{R}^{A\times D}. \]

The combined input and recurrent synaptic current is

\[ \mathbf{i}_t=\mathcal{S}_{\mathrm{in}}(\mathbf{x}_t) +\mathcal{S}_{\mathrm{rec}}(\mathbf{s}_{t-1})+\mathbf{a}_t. \]

The membrane is an INT16 state updated in fixed point:

\[ \widetilde{\mathbf{i}}_t=\mathrm{round}(\theta\,\mathrm{detach}(\mathbf{i}_t))\in\mathbb{Z}_{32}, \]
\[ \mathbf{v}_t=\mathrm{clip}_{[-32768,32767]} \left(\left\lfloor\frac{\lambda\mathbf{v}_{t-1}}{256}\right\rfloor+ \widetilde{\mathbf{i}}_t\right)\in\mathbb{Z}_{16}, \]
\[ \mathbf{s}_t=\mathbb{1}[\mathbf{v}_t\ge\theta], \qquad\mathbf{v}_t\leftarrow0\quad\text{at fired elements}. \]

The output event uses a straight-through surrogate:

\[ \widetilde{\mathbf{s}}_t= \mathbf{s}_t+\sigma(\mathbf{i}_t)-\mathrm{stopgrad}(\sigma(\mathbf{i}_t)). \]

The forward value is a hard spike, while the backward pass can propagate through \(\sigma(\mathbf{i}_t)\) to adapters and inputs. The INT8 CSR values are buffers and do not receive Adam gradients.

The output state is

\[ \mathbf{h}_t=\mathrm{LayerNorm}(\mathbf{x}_t+\mathrm{Dropout}(\widetilde{\mathbf{s}}_t)). \]

Expert routing and low-rank adapters

For \(X\in\mathbb{R}^{B\times S\times D}\), the sequence summary is

\[ \bar{\mathbf{x}}_b=\frac1S\sum_{t=1}^{S}\mathbf{x}_{b,t}. \]

The router computes

\[ \mathbf{r}_b=W_r\bar{\mathbf{x}}_b+\mathbf{b}_r\in\mathbb{R}^{M}, \]

and topk supplies selected expert indices \(I_b\).

  • At \(K=1\), \(g_{b,1}=\sigma(r_{b,I_b})\). This avoids the identically-one result of a one-element softmax, which would remove router gradient.
  • At \(K>1\), \(g_{b,k}=\mathrm{softmax}(r_{b,I_{b,k}})\) over selected logits.

The block output is

\[ Y_b=\sum_{k=1}^{K}g_{b,k}f_{I_{b,k}}(X_b). \]

Top-k reduces executed expert work and temporary state approximately in proportion to \(K\). It does not page inactive expert CSR buffers or adapters out of device memory; all expert model buffers remain resident after model.to(device).

Loss, Adam, and local plasticity

Adam-owned tensors

Adam updates only:

  • the shared basis \(E\);
  • router \(W_r,\mathbf{b}_r\);
  • expert adapters \(A_{\downarrow},A_{\uparrow}\);
  • LayerNorm scale and bias.

CSR indices, INT8 synapse values, the vocabulary codebook, and membrane state are not nn.Parameter objects.

Local Hebbian update

Each forward call records sequence/batch means of input, recurrent input, and spikes. Conceptually, for a connection from input \(j\) to output \(i\):

\[ \Delta w_{i,j}=\mathrm{round}(127\eta_h\,\bar{x}_j\bar{s}_i), \]
\[ w_{i,j}\leftarrow\mathrm{clip}_{[-127,127]}(w_{i,j}+\Delta w_{i,j}). \]

The implementation evaluates this only at existing CSR connections through index_select. It updates both input and recurrent CSR layers for every activity record, then clears records.

This is not exact gradient descent of cross entropy for primary synapses. Their teacher signal is local correlation, so convergence, accuracy, and forgetting must be measured empirically.

Actual memory layout and lifetime

Persistent model-resident objects

Object dtype Approximate size Lifetime
CSR values INT8 \(2LMDF\) bytes Model lifetime
CSR column indices INT64 \(16LMDF\) bytes Model lifetime
CSR row pointers INT64 \(16LM(D+1)\) bytes Model lifetime
Vocabulary codebook INT8 \(VR\) bytes Model lifetime
Shared basis default fp32 \(4RD\) bytes Model lifetime; Adam-owned
Routers default fp32 \(4L(DM+M)\) bytes Model lifetime; Adam-owned
Adapters default fp32 \(8LMDA\) bytes Model lifetime; Adam-owned
LayerNorm default fp32 \(8LMD\) bytes Model lifetime; Adam-owned

memory_report() returns trainable parameter count, fixed CSR connection count and persistent byte size, and vocabulary-codebook bytes. It does not state trainable parameter bytes because callers can change their dtype.

Temporary objects during a forward call

Object Representative shape Main dtype Notes
Token code lookup \((B,S,R)\) INT8 index_select output
Code cast / embedding \((B,S,R)\) / \((B,S,D)\) default float Follows basis dtype
Membrane \((B,D)\) INT16 Newly allocated in each expert forward
Previous spikes \((B,D)\) Input dtype Floating point in current implementation
Synaptic current / event \((B,D)\) Float Computed at every token
Output list and stack \(S\) tensors of \((B,D)\) Float + autograd Scales with sequence length
Float CSR values \(DF\) Float Made by quantized_values.to(dtype)
Plasticity record three \((D)\) vectors fp32 One per executed-expert forward
Sampled logits \((BS,C)\) Float Scales with candidates, not full vocabulary

Thus, stored synapses are INT8 but current computation is not wholly sparse or integer: CSR values are converted to float and spikes are dense \((B,D)\) tensors.

Sizing equations and dense-model comparison

Sparse Evo-MemoryLM

The number of primary synaptic connections across all blocks and experts is

\[ N_{\mathrm{syn}}=2LMDF. \]

The factor 2 represents one input CSR and one recurrent CSR layer per expert. Persistent CSR storage is

\[ M_{\mathrm{syn,persistent}}=2LM\left(8(D+1)+9DF\right)\ \mathrm{bytes}. \]

The dominant trainable parameter count is

\[ P_{\mathrm{trainable}} =RD+L\left[(DM+M)+M(2DA+2D)\right]. \]

This corresponds to the current bias-free adapters plus LayerNorm weight/bias. A rough FP32 Adam lower bound is \(16P_{\mathrm{trainable}}\) bytes.

Comparison

Property Dense ChronoSpike + FFN dominant term Sparse Evo-MemoryLM dominant term
Block weights \(12D^2\) CSR links \(2MDF\) + adapters \(2MDA\)
Primary synapse values Usually fp32 parameters INT8 buffers
Primary-synapse Adam Required Not required
Vocabulary Typically \(O(VD)\) input/output INT8 \(VR\) + shared basis \(RD\)
Training output \(O(BSV)\) for full CE \(O(BSC)\) with sampled CE

When \(F\ll D\) and \(A,R\ll D\), the dense \(D^2\) trainable term and its optimizer state are substantially reduced. Do not compare only one INT8 value byte: CSR uses INT64 indices and can be index-storage dominated.

Training, inference, and generation

Training sequence

sequenceDiagram
    participant T as Trainer
    participant V as TiedFactorizedVocabulary
    participant B as SparseEventMemoryBlock
    participant X as EventMemoryExpert
    participant O as Adam
    participant H as Local Hebbian updater

    T->>V: embed(input token IDs)
    V-->>B: hidden states
    loop each block
        B->>B: mean pool and Top-k route
        B->>X: selected sequences only
        loop each token position
            X->>X: CSR input + CSR recurrence + adapter
            X->>X: INT16 membrane / hard spike / surrogate
        end
        X-->>B: normalized sequence states
    end
    B-->>T: final hidden states
    T->>V: sampled candidate projection and CE
    T->>O: backward + optimizer.step()
    T->>H: apply_local_plasticity()

Call apply_local_plasticity() after optimizer.step(). If a failure, NaN detection, or skipped optimizer step occurs, call clear_local_plasticity() so later batches do not update from stale activity records.

Inference and generation

forward(token_ids) returns full-vocabulary logits. generate() samples the next token using temperature and optional top-k, but passes the full current prompt to forward() at every generation step. KV caching, candidate-only decoding, and persistent membrane-state caching are not implemented.

Worked sizing example

The following implemented memory_report() configuration has been exercised:

\[ V=32768,\quad D=2048,\quad L=12,\quad R=128,\quad\rho=0.005,\quad M=2,\quad K=1,\quad A=16. \]

Here \(F=\lceil0.005\times2048\rceil=11\), so fixed primary CSR links total

\[ N_{\mathrm{syn}}=2\times12\times2\times2048\times11=1,081,344. \]

The measured static report is:

Measure Value
Adam-owned parameters 1,982,488
Fixed INT8 CSR connections 1,081,344
Persistent CSR storage about 10.03 MiB
INT8 vocabulary codebook 4.00 MiB

This is static model accounting only. Candidate logits, autograd graphs, float CSR conversion, optimizer tensors, CUDA allocator overhead, and tokenizer memory are additional. Always measure torch.cuda.max_memory_allocated() for the real batch size and sequence length, including on an 8 GB GPU.

Constraints and operating guidance

Constraints

  1. PyTorch CSR is beta: available GPU kernels, autograd behavior, and performance depend on the PyTorch/CUDA version.
  2. The local rule is not a global-error gradient: cross entropy does not directly optimize primary synapses.
  3. No expert load balancing: the current router has no load-balancing loss, capacity factor, or expert offload.
  4. No generation cache: long outputs repeatedly evaluate their prompt prefix.
  5. Full forward() is vocabulary-sized: it explicitly materializes \((B,S,V)\) logits during inference.
  6. INT8 storage is not INT8 execution: CSR values are cast to float inside forward.
  • Use sampled_cross_entropy() for training and tune --sampled-negatives with measurement.
  • Start with small \(D,L,M\) and short \(S\); inspect loss, firing rate, LocalSynapseUpdates, and peak VRAM.
  • Before increasing --sparse-connectivity, measure both INT64 index storage and sparse-kernel throughput.
  • When resuming, verify architecture: sparse_event_memory and \(D,R,L,M,K,A,\rho\) together with vocabulary size.
  • For Docker SDK artifacts, retain the model, configuration, memory report, and tokenizer archive under llm_type="SparseEventMemoryLM".