darioooooo0o/spark-4b-engram

💻 DeepSeek Engram (Coding-Specialized) on Spark-X2.5-4B

Requests, questions or suggestions? Message me on X: https://x.com/imdariotoo

⚠️ EXPERIMENTAL RESEARCH ARTIFACT: This repository contains an experimental, code-specialized implementation of Conditional Memory via Scalable Lookup (Engram) (arXiv:2601.07372) trained on top of a 100% frozen Spark-X2.5-4B model.

---

📥 Direct Downloads & Quick Access

File	Format	Size	Description	Direct Download Link
Engram 20% Trained Weights	.safetensors	3.1 GB	Full 826M-param Coding Engram module	Download Weights (Direct)
Canonical Tokenizer Map	.pt	1.1 MB	Precomputed canonical token mapping	Download Map (Direct)
Full File Tree	—	—	Architecture files & benchmark data	Browse All Files & Versions

---

📌 Overview & Multi-Language Scope

This model explores how O(1) conditional memory lookup can augment a frozen language model for multi-language code generation and algorithmic reasoning:

• Polyglot Code Memory Bank: The 262k-slot hash tables act as an external memory bank storing multi-token syntax patterns, recursion templates, and standard library idioms across Python, C++, Java, JavaScript, C#, SQL, Bash, and Rust.
• 100% Frozen Backbone (Zero Drift): The base 4.11B transformer weights remain completely frozen. No weights in Spark-4B were modified, fine-tuned, or adapted with LoRA.
• Preserved General Reasoning: Because the base model remains untouched, the model experiences zero catastrophic forgetting of its original conversational or non-coding skills.

---

📚 Training Corpus (25,000,000 Multi-Language Tokens)

The Engram module was trained across 25M tokens sourced from diverse programming datasets:

1. nickrosh/Evol-Instruct-Code-80k-v1 (~12M tokens, Multilingual):
   • Complex algorithmic problems, data structures, and competitive programming across Python, C++, Java, JavaScript, C#, Bash, PHP, and SQL.
2. sahil2801/CodeAlpaca-20k + iamtarun/python_code_instructions (~3M tokens):
   • Multi-language task completions, idiomatic one-liners, standard library manipulations, and docstring-to-code implementations.
3. codeparrot/codeparrot-clean (~10M tokens):
   • Clean real-world repository code, modular package architectures, class hierarchies, and production syntax.

---

🏗️ Architecture Specifications

Following DeepSeek's paper, the memory module is structured as follows:

Component	Specification	Description
N-gram Orders	(2, 3, 4)	Multi-scale n-gram context modeling for code tokens
Heads per Order	8 heads (24 heads total per layer)	Multi-head bitwise XOR hashing
Slots per Head	262,144 prime slots	Coprime prime moduli per head eliminating cross-head collisions
Embedding Dimension	64	Compact, high-density idiom memory representation
Target Layers	Layer 2 and Layer 18	Dual-layer topology (shallow + mid) matching empirical sweep peaks
Tokenizer Compression	131,072 -> 100,096 keys	NFKC + NFD + Accent Strip + Lowercase canonicalization (-23.6%)
Context Gating	Signed Square-Root Gate	sigmoid(sign(S) * sqrt(|S|)) with FP32 RMSNorm
Temporal Convolution	Causal ShortConv (kernel=4, dilation=2)	Strict causal left-padding with SiLU activation and skip connection
Memory Capacity	826.24M Parameters	Sized at 20.09% of the Spark-4B base model

---

📊 Benchmark Results (HumanEval 50-Problem Suite)

Evaluation conducted greedily (temperature=0.0, max_new_tokens=150) against standard unit tests (check(entry_point)).

Model Configuration	Base Backbone	Engram Parameters	HumanEval (50 Problems)	Pass Rate	Net Gain
Stock Spark-1.7B	1.7B	0M (Stock)	11 / 50	22.0%	Baseline
Spark-1.7B + 20% Engram (4B-Adapted)	1.7B	~820M (Zero-shot sliced from 4B)	0 / 50	0.0%	-22.0%
Stock Spark-4B	4.1B	0M (Stock)	15 / 50	30.0%	Baseline
Spark-4B + 20% Engram (Native)	4.1B	826.2M (Trained on 25M code tokens)	36 / 50	72.0%	+42.0% 🚀

Key Findings & Research Insights:

1. Massive Algorithmic Unlock (+42.0% absolute on 4B):
   • On the 50-problem HumanEval suite, the coding Engram on Spark-4B more than doubles problem-solving capability (30.0% -> 72.0%), solving 21 problems that stock 4B failed.
   • The frozen 4B backbone provides rich semantic representations (hidden states), allowing the signed sqrt gate to accurately retrieve memorized idioms and code patterns directly into the residual stream.
2. Representation Space Coupling (Asymmetric Transfer):
   • While the N-gram hash tables (tables.weight, 805M params) are vocabulary-dependent and architecture-agnostic, the dense projections (W_q, W_v) are tightly coupled to the specific model's hidden representation geometry.
   • Slicing 4B's projections (2560 -> 2048) to fit 1.7B disrupts the gate dot-products, proving that while memory tables can be shared, projection layers must be native or fine-tuned to each backbone's activation space.

---

🧩 Quantization & Deployment Compatibility

• Weight-Only Quantization (GPTQ / AWQ / INT4 / INT8):
  • Fully Compatible: You can quantize the frozen base 4B transformer weights (e.g. GPTQ/AWQ to reduce backbone memory from ~8 GB to ~2.5 GB).
  • Because inter-layer activations (hidden states) remain in 16-bit floating point (bfloat16/float16), the Engram module (enable_engram()) plugs directly into the quantized backbone without numerical degradation.
  • The Engram weights (~3.1 GB safetensors) remain unquantized in bfloat16 / float32.
• GGUF / llama.cpp:
  • Not yet supported in standard llama.cpp: GGUF requires custom C++/GGML tensor graph kernels for Engram's multi-head XOR hashing, prime moduli grid lookup, and Causal ShortConv. Use Hugging Face / PyTorch for Engram inference.

---

🚀 Quickstart & Inference

1. Requirements

pip install torch transformers safetensors

2. Loading the Model with Engram

import torch
from safetensors.torch import load_file
from configuration_spark import Spark2_5Config
from modeling_spark import Spark2_5ForCausalLM
from transformers import AutoTokenizer

device = "cuda" if torch.cuda.is_available() else ("xpu" if hasattr(torch, "xpu") and torch.xpu.is_available() else "cpu")
dtype = torch.bfloat16

# 1. Load Base Model
model_id = "XHToken/Spark-X2.5-4B"
tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
model = Spark2_5ForCausalLM.from_pretrained(model_id, torch_dtype=dtype).to(device)

# 2. Inject Engram Memory at Layers (2, 18)
model.enable_engram(
    target_layers=(2, 18),
    mem_dim=64,
    num_heads=8,
    slots_per_head=262144,
    orders=(2, 3, 4),
    kernel_size=4,
    dilation=2,
    device=device,
    dtype=dtype,
)

# 3. Load Trained Engram Weights
weights = load_file("engram_layer_weights_spark4b_20pct_25m.safetensors")
for k in weights:
    weights[k] = weights[k].to(device)
model.load_state_dict(weights, strict=False)
model.eval()

# 4. Generate Python Code
prompt = "def greatest_common_divisor(a: int, b: int) -> int:\n    \"\"\" Return a greatest common divisor of two integers a and b \"\"\"\n"
inputs = tokenizer(prompt, return_tensors="pt").to(device)

with torch.no_grad():
    outputs = model.generate(**inputs, max_new_tokens=100, temperature=0.0)

print(tokenizer.decode(outputs[0], skip_special_tokens=True))

---

📖 Citation & References

@article{deepseek2026engram,
  title={Conditional Memory via Scalable Lookup: A New Axis of Sparsity for Large Language Models},
  author={DeepSeek-AI and Peking University},
  journal={arXiv preprint arXiv:2601.07372},
  year={2026}
}