💻 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.

Training status: This Engram module is heavily undertrained. Treat the weights and reported results as preliminary research, not a mature or fully trained release.

---

📥 Direct Downloads & Quick Access

---

📌 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:

---

📊 Benchmark Results (HumanEval 50-Problem Suite)

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

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}
}