$ show deepseek-ai/Engram
This repository contains the official research implementation of the Engram module, a conditional memory mechanism for large language models that augments transformer blocks with static N-gram lookup. The codebase is a single standalone Python demo script (`engram_demo_v1.py`) that illustrates the core data flow: tokenization, compression, hash-based N-gram indexing, multi-head embedding retrieval, gating via learned keys and query dot-product, and a short depthwise convolution with per-group RMSNorm. The architecture is a mock LLM (embedding, transformer blocks with mocked Attention/MoE, linear head) to focus solely on the Engram module. The repo is a research artifact, not a production service; it lacks tests, packaging, and error handling but provides a clear, commented implementation of the novel hash-based conditional memory pattern. The primary value for engineers is the extractable pattern for deterministic N-gram lookup and gating, and the depthwise convolution with group normalization.
The LLM identified 3 files worth highlighting for reuse:
The NgramHashMapping class implements a deterministic hash function for N-gram token sequences using prime numbers and XOR mixing. It builds a per-layer hash mapping from compressed token IDs to embedding indices, with per-head prime moduli. This pattern is reusable for any project needing token-based hash indexing for sparse retrieval, and the design is cleanly separated from the embedding logic.
The ShortConv class implements a depthwise 1D convolution followed by per-group RMSNorm and SiLU activation. It is used in the Engram module for local sequence mixing after gating. The pattern is self-contained and reusable for efficient channel-wise temporal processing in other architectures.
The MultiHeadEmbedding class concatenates multiple embedding tables into a single nn.Embedding by using precomputed offsets per head. This allows efficient multi-head lookup from a single embedding matrix. The pattern is simple but clearly demonstrates how to manage multiple embedding tables with contiguous indexing.