Engram

N-grams were introduced by Shannon, and fastText turned them into learned vectors. Hash Embeddings hashed them into tables and Product Key Memory introduced large, sparsely accessed memory within transformers. Engram from DeepSeek is the culmination of all this work.

How it works

The fundamental idea is to give the model a space where it can store information about relationships between local n-grams to avoid spending FLOPs computing these relationships. The idea is that this will allow the model to put more capacity elsewhere. Trading off this memory for FLOPs works as the memory look ups are cheap.

We start by sliding a window (stride 1) over our token sequence and extracting every run of $n$ consecutive tokens, for each n-gram order we want. We then hash these into tables of learned embeddings. We fetch these learned values and add them to the hidden states, or residual streams in the case of mHC.

Looking up our table

The hashing can be implemented as a simple multiplicative XOR hash. We keep separate tables for the bigrams and trigrams and for each of these, we split the hashing over multiple tables to avoid hash collisions. For example, we could have 4 tables per n-gram order, giving us a total of 8 tables. We would produce 8 hashes in total and look up 8 embedding vectors. These are later projected to the model dimension $d_{m o d e l}$ before being added to our hidden states, so their concatenated length, $d_{m e m}$ , can be chosen freely. In practice, each table stores $\frac{d_{m e m}}{8}$ -sized vectors which we concatenate to produce our lookup embedding vector $e_{t} \in R^{d_{m e m}}$ . Storing parts of the embedding in different tables to avoid hash collisions is a well-known method.

Gating lookup information

We'd like to give the model some control over how much of the Engram embedding vector it wants, so we compute a gate using a projection and a dot product using the following method:

\begin{matrix} (1) & \begin{aligned} k_{t} & = W_{k} e_{t} \\ v_{t} & = W_{v} e_{t} \end{aligned} \end{matrix}

\begin{matrix} (2) & α_{t} = σ (\frac{h_{t}^{⊤} k_{t}}{\sqrt{d_{m o d e l}}}) \end{matrix}

Where $W_{k}$ and $W_{v}$ are projection matrices to create key and value vectors from our fetched embedding vector $e_{t}$ . After the projections, we use $k_{t}$ to measure how similar $e_{t}$ is to $h_{t}$ through a dot product, computed in Eq. 2. $σ$ is the sigmoid function forcing the gate to be between 0 and 1. The $t$ subscript here denotes that the vectors belong to token $t$ and $\sqrt{d_{m o d e l}}$ is to control the variance and avoid vanishing gradients from the sigmoid.

Once we have $α_{t}$ , we can use it to gate $v_{t}$ and add it to $h_{t}$ :

\begin{matrix} (3) & h_{t}^{n e w} = h_{t} + α_{t} v_{t} \end{matrix}

Note that $α$ is per token but it is a single scalar per token, so it cannot discriminate which of the $e_{t}$ values should be let in. The DeepSeek Engram paper goes further, computing a depthwise dilated 1D convolution over the token dimension and then applying a SiLU gate, but in their ablations, that made little impact so we'll leave it out for brevity. The gradient reaching the Engram table rows flows through the $v$ and $k$ vectors.

Extending to mHC

With mHC, we have $n$ streams and an $h_{t}$ vector belonging to each stream. We use a separate $W_{k}$ projection for each stream as that determines the gate and produces a separate $α_{t}$ for each stream, but share the $W_{v}$ projection for efficiency reasons. Other components stay the same, allowing the Engram embedding vector $e_{t}$ to be added to each of the streams.

Tokeniser compression

To ensure that n-grams for the same word written in different formats (toddler vs TODDLER) end up in the same lookup, we need to ensure the tokeniser produces the same tokens for them. Since subword tokenisers prioritise lossless compression, they end up assigning different token IDs to the same word, even if one is a capitalised version of the other, which learns separate embeddings for the same word, increasing table and model size. To resolve this concern, DeepSeek applies the following text normalisation rules:

Lowercasing so TODDLER and toddler end up with the same token IDs.
Removing accents so café and cafe get the same token
Unicode normalisation so ＡＢＣ (full-width) and ABC get the same tokens
Collapsing repeated whitespace and removing leading and trailing whitespace.

In practice, we need a mapping table which maps original token IDs to compressed IDs. We build this mapping table once at the start of training by going over our vocab, normalising text, storing their original IDs and creating compressed IDs for them. Once the mapping is created, during inference and training, we simply do an array lookup, using the original IDs to retrieve the compressed IDs, fast and cheap.

A naive PyTorch implementation

Below is a naive PyTorch implementation, not including the hashing and compression functions as those are pretty standard. A few notes to help understand the implementation:

On line 22, self.get_embeddings(n_grams) does the lookup using the hashes to get K × n_orders embedding vectors, which are concatenated together to form a vector of size $d_{m e m}$ . K is the number of tables per n-gram order and n_orders is the number of n-gram orders. If we did bigram and trigram with four tables each, we'd look up 8 vectors of size $\frac{d_{m e m}}{8}$ each and concatenate them along the last dimension.
On line 26 you can see we have a shared $W_{v}$ projection of the embedding vector $e$ , but on line 38, each branch has its own $W_{k}$ projection.
We apply RMSNorm to $h$ and $k$ before we compute the similarity to make them scale-independent.
Line 48 passes the similarity through a signed square root, $sign (z) \sqrt{| z |}$ , before the sigmoid, which keeps large similarities from saturating the gate.
Standard implementations also apply a dilated depthwise 1D convolution; we have left it out for simplicity.
Line 60 returns the gated $v$ to be added to our $h$


class Engram(nn.Module):
    def __init__(self, cfg):
        super().__init__()

        self.d_model = cfg["d_model"]
        self.n = cfg["n"]

        # Memory projections: [b, t, d_mem] -> [b, t, d_model]
        self.W_V = nn.Linear(cfg["d_mem"], cfg["d_model"], bias=True)
        self.W_K = nn.ModuleList([
            nn.Linear(cfg["d_mem"], cfg["d_model"], bias=True)
            for _ in range(self.n)
        ])

        self.norm_h = nn.RMSNorm(cfg["d_model"])
        self.norm_k = nn.RMSNorm(cfg["d_model"])

    def engram_context_aware_gating(self, n_grams, x_res):
        # Look up hashed N-gram embeddings and concatenate them
        # [b, t, n_orders, K] -> [b, t, d_mem]
        e = self.get_embeddings(n_grams)

        # Project memory into residual-stream dimension
        # [b, t, d_mem] -> [b, t, d_model]
        v = self.W_V(e)

        # Apply independent gates to each mHC residual stream
        # n * [b, t, d_model] -> [b, t, n, d_model]
        return torch.stack([
            self.context_gate(x_res[:, :, m, :], e, v, m)
            for m in range(self.n)
        ], dim=2)

    def context_gate(self, h, e, v, branch):
        # Project memory to a branch-specific key
        # [b, t, d_mem] -> [b, t, d_model]
        k = self.W_K[branch](e)

        # Compute normalized query-key similarity
        # [b, t, d_model] -> [b, t, 1]
        z = (
            self.norm_h(h) * self.norm_k(k)
        ).sum(-1, keepdim=True) / math.sqrt(self.d_model)

        # Apply signed square-root scaling
        # [b, t, 1] -> [b, t, 1]
        z = z.abs().clamp_min(1e-6).sqrt() * z.sign()

        # Convert similarity to a scalar gate per token
        # [b, t, 1] -> [b, t, 1]
        gate = torch.sigmoid(z)

        # Gate memory values across all embedding dimensions
        # [b, t, 1] * [b, t, d_model] -> [b, t, d_model]
        v_gated = gate * v

        # Return gated memory
        # [b, t, d_model] -> [b, t, d_model]
        return v_gated

Results

The plot below shows the held-out cross-entropy loss of a vanilla 124M model vs the same model with mHC vs mHC plus Engram trained for 2.5B tokens on the FineWeb-Edu dataset. We can see a ~2.5% improvement in validation loss when combining Engram and mHC. For more details on how Engram behaves, please see the Nenya technical report.

Figure 1: A comparison of vanilla 124M transformer vs vanilla + mHC vs vanilla + mHC + Engram.