Baidu FactorEngram Lifts Hard NIAH-3 by 43 Points over Engram at 340M

FactorEngram: Factorized N-gram Memory with Basis-Level Gating for Language Models

Bowen Yang, Jingbo Zhou, Qinghong Miao, Hua Wu

cs.CL, cs.AI

2026-09-29

FactorEngram stores n-grams as sparse shared-basis coefficients with per-basis gates. At 340M, NIAH-3 jumps from Engram's 15.0 to 58.3; WikiText PPL drops from 23.19 to 21.29.

What problem this solves

A Transformer still has to rebuild phrases like 'the Eiffel Tower' layer by layer, even though those strings carry fairly static lexical and factual content. Lookup-based memory adds a side branch that hashes local token patterns into a table and reads out learned embeddings. Engram is the best-known design in this family; DeepSeek-V4.1Flash already uses it.

Engram still treats each retrieved vector as one block: one hashed slot, one scalar gate. A polysemous pattern such as 'the bank' can only be turned up or down as a whole, so the model cannot keep the finance sense and drop the river sense. Parameter sharing happens only when two n-grams collide in the hash, which has nothing to do with meaning. Both failures come from the same choice: the atom of memory is a full n-gram embedding.

Method

FactorEngram, from Baidu and NTU, stores each local pattern as coefficients over a dictionary of basis vectors that is shared inside a layer. Coefficient tables are addressed by the discrete suffix: unigrams use direct indexing, 2-grams and 3-grams use hashing, with several lookup heads per length. Retrieved slices are concatenated in a fixed order so each coordinate lines up with one dictionary row. An L1 penalty on those coefficients, applied before gating, pushes each pattern onto a small subset of bases.

The same dictionary scores the context. The backbone hidden state is projected, RMSNormed, and dotted with every basis vector; a sigmoid turns the scores into per-basis gates that scale the retrieved coefficients. Reconstruction is a linear combination of the gated coefficients, then a projection back to backbone width, then a short depthwise causal convolution with a residual. The result is added to the residual stream. Attention and the FFN stay untouched.

Default placement is before attention in the middle of the stack: layers 10 and 12 of a 24-layer 340M model. Training is next-token NLL plus λ times the L1 term, with λ = 1e-3. Tables, dictionary, projections, and the convolution are learned jointly with the backbone from scratch.

Results

The 340M backbone sees 30B FineWeb-Edu tokens; the 1B backbone sees 120B. Context is 8K. Extra lookup parameters are 1B at 340M (250K slots for both 2-gram and 3-gram tables) and 2B at 1B (250K / 750K). Baselines are a plain Transformer at the same scale and the authors' Engram reproduction. Downstream average covers PIQA, HellaSwag, WinoGrande, ARC-Easy/Challenge, Social IQA, and BoolQ. NIAH-1/2/3 are three 8K Needle-in-a-Haystack variants of increasing difficulty.

SetupWikiText↓LAMBADA↓Avg acc↑NIAH-1NIAH-2NIAH-3
340M Transformer23.1924.9048.1044.247.722.1
340M Engram22.2923.9449.0669.837.615.0
340M FactorEngram21.2921.6949.8270.679.158.3
1B Transformer15.5711.7355.6147.572.825.7
1B Engram15.6011.8355.8145.875.216.6
1B FactorEngram15.5711.3955.8072.084.544.2

At 340M, WikiText falls from 23.19 to 21.29 and LAMBADA from 24.90 to 21.69 versus the Transformer; average downstream accuracy rises from 48.10 to 49.82. The gap that matters is retrieval: NIAH-2 and NIAH-3 beat Engram by 41.5 and 43.3 points. Engram itself undershoots the Transformer on those two tasks (37.6 vs 47.7, 15.0 vs 22.1).

At 1B the language-modeling and downstream picture flattens. WikiText matches the Transformer at 15.57; average accuracy is 55.80 against Engram's 55.81. Retrieval still moves: NIAH-1 45.8 → 72.0, NIAH-2 75.2 → 84.5, NIAH-3 16.6 → 44.2 versus Engram, a range of 9.3 to 27.6 points.

All ablations use the 340M model. Replacing per-basis gates with a scalar gate after reconstruction drops NIAH-3 from 58.3 to 9.4 and average accuracy from 49.82 to 48.15. Removing unigrams cuts NIAH-2/3 to 44.2 and 23.0. Adding unigrams to Engram helps all three NIAH scores but slightly hurts perplexity and average accuracy. Raising λ from 0 to 1e-3 helps most metrics; 1e-4 wins WikiText at 21.18; 1e-2 hurts everything, with WikiText at 23.89, worse than no sparsity.

A single memory module works best around layer 12. Adding a second module at layer 10 gives the highest NIAH-3 among tested pairs. Inside the layer, before-attention wins every accuracy column and LAMBADA. Before-FFN matches WikiText (21.28 vs 21.29) but NIAH-3 falls to 18.8. Injecting into the SwiGLU up-projection is worst across the board (WikiText 23.34, NIAH-3 12.2).

Why it matters

This is a representation change on the Engram line, not a new scaling recipe. If you already ship n-gram lookup, two pieces are worth copying: shared bases, so related patterns can reuse components, and per-basis gates, so context can keep some components and drop others. Gains concentrate on long-context retrieval. Language modeling and downstream improvements are small, and at 1B they essentially vanish.

The engineering defaults are concrete: cover unigrams through trigrams; insert in middle layers, before attention; keep L1 around 1e-3. The cost is an extra 1B to 2B table parameters. Latency, bandwidth, and memory footprint are not reported. Nothing in the paper speaks to models larger than 1B, or to production-scale Engram stacks.

Limitations

The authors say so: evaluation stops at 340M and 1B. They note that 1B still matters for on-device work, and leave larger scales to future work. Read against Table 1, that caveat is the whole story for perplexity and downstream: at 1B FactorEngram is no longer ahead of Engram. Retrieval is the remaining claim.

Several checks are missing. Engram losing to the Transformer on NIAH-2/3 is never diagnosed, so it is unclear whether the reproduction, hash collisions, or the missing unigram is at fault. NIAH is 8K only. Individual downstream scores are folded into an average. The dictionary is never inspected the way a sparse autoencoder would be, so shared semantic components remain an architectural story, not a measured one. TN-gram also factorizes n-grams and keeps scalar gating; it is cited, not compared. There is no iso-parameter run that puts the same extra weights into the FFN or an MoE.

Terms

Source

What people are saying

Related papers

All paper explainers