Engram Attention

Engram: Conditional Memory via Scalable Lookup Paper: arXiv:2601.07372 (DeepSeek AI, 2026)

Engram augments a transformer with an N-gram conditional memory module that performs O(1) deterministic lookup into a set of learnable embedding tables. Each position retrieves a bigram/trigram (up to max_ngram_size) fingerprint, fuses the retrieved vectors with the hidden state through a learned gate, and returns a corrected hidden state.

This file adapts the official demo (deepseek-ai/Engram) to the Frankenstein architecture:

  • Works with standard (B, T, D) tensors (no hyper-connection dimension).

  • Uses raw token IDs instead of a normalizing tokenizer wrapper.

  • Plugs in as a standard layer_type in HybridLayer.

class src.model.attention.engram.EngramLayer(*args: Any, **kwargs: Any)[source]

Bases: Module

Engram: Conditional Memory via Scalable Lookup (arXiv:2601.07372).

Given hidden states x of shape (B, T, hidden_size) and the original token IDs (B, T), the module:

  1. Hashes token N-grams into embedding-table indices.

  2. Looks up per-head N-gram embeddings and concatenates them.

  3. Passes the concatenated embedding through a causal depthwise convolution.

  4. Computes a scalar gate: σ(dot(norm(key), norm(query)) / √D).

  5. Returns gate ⊙ value_proj(engram_emb) projected to hidden_size.

The layer is registered as "engram_attn" in HybridLayer and can appear at any position in layer_pattern. When no input_ids are available the module returns a zero tensor (graceful degradation).

Config knobs (all optional, sensible defaults):

  • engram_max_ngram_size (int, default 3): highest N-gram order (2 and 3).

  • engram_n_heads_per_ngram (int, default 4): hash heads per N-gram order.

  • engram_embed_dim_per_head (int, default 32): embedding dim per head.

  • engram_kernel_size (int, default 4): ShortConv kernel width.

  • engram_seed (int, default 42): RNG seed for hash multipliers.

forward(x: torch.Tensor, input_ids: torch.Tensor | None = None, logical_layer_idx: int | None = None) → torch.Tensor[source]
Parameters:

  • x – (B, T, hidden_size) – incoming hidden states

  • input_ids – (B, T) int64 – token indices (required for lookup)

  • logical_layer_idx – ignored, kept for HybridLayer interface parity

Returns:

(B, T, hidden_size) – Engram-corrected hidden states

class src.model.attention.engram.MultiHeadEmbedding(*args: Any, **kwargs: Any)[source]

Bases: Module

Packs num_heads separate embedding tables of sizes N_0, N_1, … into a single nn.Embedding by adding per-head offsets to the lookup indices.

forward(indices: torch.Tensor) → torch.Tensor[source]
Parameters:

indices – (B, T, num_heads) int64 – raw per-head hash indices

Returns:

(B, T, num_heads, embed_dim)

class src.model.attention.engram.NgramHasher(max_ngram_size: int, n_heads_per_ngram: int, base_vocab_size: int, seed: int = 42)[source]

Bases: object

Computes deterministic N-gram hashes from raw token IDs.

For each N-gram size n in [2, max_ngram_size] and each head j, the hash is:

hash_j(i) = XOR(token[i-k] * multiplier[k] for k in 0..n-1) % prime_j

Multipliers are seeded odd random integers drawn at construction time. Each (n, j) pair uses a distinct prime modulus so hash collisions across heads are statistically independent.

hash(input_ids: numpy.ndarray) → numpy.ndarray[source]
Parameters:

input_ids – int64 array of shape (B, T)

Returns:

int64 array of shape (B, T, total_heads) where total_heads = (max_ngram_size - 1) * n_heads_per_ngram

Return type:

hashes

class src.model.attention.engram.ShortConv(*args: Any, **kwargs: Any)[source]

Bases: Module

Causal depthwise Conv1d that mixes the retrieved N-gram embeddings across a small local window before gating. The causal padding ensures position i only sees N-gram information from positions ≤ i.

forward(x: torch.Tensor) → torch.Tensor[source]

Args / returns: (B, T, C)