PaperlayerView exact revision on arXiv

Source-linked AI summary

NGM: A Plug-and-Play Training-Free Memory Module for LLMs

Yuwen Qu, Wenhui Dong, Chenyang Si, Caifeng Shan

arXiv:2605.16893v1cs.AI

TL;DR

LLMs lack a native lookup mechanism for local, static patterns, while existing memory approaches require trainable embeddings and additional infrastructure. NGM is a training-free, plug-and-play module that reuses pretrained token embeddings for causal N-gram memory and injects them with cosine gating, improving average Qwen3 performance by +0.5 to +1.2 points across eight benchmarks, especially code and knowledge tasks.

  • Problem

    LLMs lack a native lookup primitive for local lexical and symbolic dependencies, while existing lookup methods require trainable parameters, dedicated training, or specialized infrastructure.

  • Method

    NGM is a training-free, plug-and-play module that constructs causal multi-scale N-gram representations from pretrained token embeddings and injects them through a non-parametric cosine gate.

  • Results

    +0.5 to +1.2 points in average performance across Qwen3 models from 0.6B to 14B, with strongest gains on code generation and knowledge-intensive benchmarks.

  • Takeaways & Limitations

    NGM provides a parameter-free way to add local-memory signals, with benefits clearest when short-range pattern stability matters rather than uniformly across tasks.

  • Takeaways & Limitations

    NGM’s fixed, order-insensitive injection rule may mishandle order-sensitive phrases and suit different tasks or generation styles unequally.

Abstract

from arXiv · show

Recent studies introduce conditional memory modules that decouple knowledge storage from neural computation, enabling more direct knowledge access. Compared to MoE, which relies on dynamic computation paths, explicit lookup provides a more efficient knowledge retrieval mechanism. However, these approaches still depend on learned memory embeddings, requiring additional training and limiting flexibility. To address this, we propose N-gram Memory (NGM), a training-free, plug-and-play module composed of a Causal N-Gram Encoder and a Cosine-Gated Memory Injector. The Causal N-Gram Encoder directly averages the pretrained token embeddings of the backbone model to construct N-gram representations, thereby eliminating the need to train separate N-gram embeddings from scratch. This design requires neither an additional memory table nor a retrieval pipeline. The Cosine-Gated Memory Injector then uses a non-parametric cosine gate with ReLU to modulate the retrieved embeddings into the contextual representations. We evaluate NGM on the Qwen3 series from 0.6B to 14B across eight benchmarks. NGM improves average performance by 0.5 to 1.2 points, with particularly clear gains on code generation and knowledge-intensive tasks (e.g., +3.0 on LiveCodeBench and +3.03 on GPQA for Qwen3-14B). Moreover, NGM also improves performance in multimodal benchmarks (e.g., MMStar +1.53 on Qwen3-VL-2B).

1 Introduction

NGM addresses the cost and inflexibility of learned lookup memory by reusing pretrained token embeddings to provide training-free local-memory signals in frozen LLMs. Across Qwen3 models and multimodal tasks, it improves average benchmark performance, especially for code and knowledge-intensive evaluation.

  • Motivation: Lookup-style memory separates static pattern reuse from dynamic Transformer computation, targeting named entities, identifiers, units, terminology, and formulaic phrases.Prior approaches demonstrate this direction but require learned memory components, additional trainable parameters, and dedicated training.
  • Method: Averaging pretrained token embeddings within a local causal window constructs N-gram features without training a separate N-gram embedding space.This directly targets the need for additional trainable parameters, dedicated training, and specialized memory infrastructure in prior lookup-style pipelines.
  • Representation alignment: 0.312 and 0.137 mean cosine similarities at the default injection layers exceed shuffled controls’ 0.172 and 0.084 and random controls’ 0.014 and 0.008.The comparison is reported for Qwen3-8B across model depth.
  • Method: NGM is a training-free, plug-and-play module that injects local N-gram signals into frozen decoder-only LLMs using the pretrained embedding space as reusable memory.The method treats aggregated nearby-token embeddings as cues that the decoder can reuse rather than reconstruct entirely through computation.
  • Empirical results: +0.5 to +1.2 average points across Qwen3 models from 0.6B to 14B, including +3.0 on LiveCodeBench and +3.03 on GPQA for Qwen3-14B.NGM was evaluated across eight benchmarks covering mathematics, code, knowledge, and alignment.
  • Multimodal results: NGM improves all reported Qwen3-VL-2B multimodal benchmarks when applied only to the language decoder.This extends the method beyond text-only evaluation and indicates generality across multimodal tasks.

2 Related work

Prior work connects local lexical statistics and conditional lookup to scalable memory, while interpretability studies characterize the residual stream as a vocabulary-aligned workspace. These lines of research motivate explicit memory access and embedding-based integration in language models.

  • Conditional memory and embedding scaling: Conditional memory extends the intuition of classical N-gram models [25, 7] and local lexical predictability [32] [2] through lookup-based sparsity, complementing MoE’s conditional computation [36] [16].Classical N-gram models use fixed-order Markov assumptions, while MoE scales capacity through conditional computation.
  • Residual stream alignment: Interpretability work views the Transformer residual stream as a shared linear workspace, with logit-lens and probing studies showing intermediate states remain partly projectable into vocabulary space [11].With tied embeddings, this projection directly aligns with the input embedding layer; the passage also distinguishes models with untied embeddings.

3 Methodology

NGM is a training-free, plug-and-play memory module that derives multi-scale local signals from the backbone’s pretrained token embeddings and injects them into frozen decoder representations. Its Causal N-gram Encoder and Cosine-Gated Memory Injector require no learned memory embeddings, external retrieval, or additional parameters.

  • 3 Methodology: NGM combines a Causal N-gram Encoder with a Cosine-Gated Memory Injector to construct and inject local memory signals into decoder representations.The encoder builds multi-scale N-gram representations from backbone token embeddings, while the injector scores them against decoder hidden states and adds an aggregated residual.
  • 3 Methodology: During inference, all backbone parameters remain frozen, and NGM’s added computation depends only on the current input sequence and a small predefined set of N-gram sizes.The module is therefore plug-and-play and does not require separately trained N-gram embeddings or an additional memory table.
  • Causal N-gram Encoder: The Causal N-gram Encoder averages trailing windows of pretrained token embeddings at multiple scales, using left-padding to preserve length and causality without learned memory entries.This bag-of-embeddings representation is intentionally order-insensitive within each window and provides a parameter-free local summary.
  • Cosine-Gated Memory Injector: The injector uses cosine similarity between decoder hidden states and N-gram vectors, optionally applies ReLU, and writes the gated aggregate through a scaled residual connection.The gate selects memory vectors compatible with the current decoder state without learned projections or external retrieval.
  • Cosine-Gated Memory Injector: With KV caching, NGM constructs representations from the full input prefix and slices the last L positions to align memory with the currently available hidden states.This preserves causal consistency during autoregressive decoding.

4 Experiments

Across five Qwen3 scales and eight text benchmarks, NGM consistently improves average performance without adding trainable parameters, with strongest gains on code and knowledge tasks. It also transfers to multimodal models and is supported by ablations showing the importance of multi-scale construction, ReLU gating, and stack fusion.

  • Experimental setup: NGM is evaluated on GSM8K, MATH500, HumanEval, LiveCodeBench v5, MMLU-Redux, GPQA-Diamond, IFEval, and TruthfulQA across Qwen3 scales of 0.6B, 1.7B, 4B, 8B, and 14B.The setup uses Qwen3 for controlled scaling analysis and follows standard benchmark protocols, with MMLU-Redux evaluated at context length 4096.
  • Main results: +1.2, +0.5, +0.6, +0.8, and +0.7 average-score gains are achieved by NGM from Qwen3-0.6B through 14B, with especially consistent improvements on LiveCodeBench, HumanEval, GSM8K, and GPQA.Table 1 evaluates models with and without NGM under identical decoding settings; NGM adds no trainable parameters.
  • Multimodal transfer: +1.53 on MMStar is the largest gain in a preliminary Qwen3-VL-2B-Instruct evaluation, where NGM improves or matches the baseline on all five multimodal and text benchmarks.NGM modifies only the language decoder; vision tokens are excluded from N-gram pooling, and the visual encoder and fusion modules remain unchanged.
  • Ablations: 72.17 versus 70.38 shows that ReLU gating is important for stable gains, while stack fusion reaches 72.17 versus 71.07 for concat.Multi-scale N = {2, 3} gives the strongest average result, whereas adding n = 4 improves some tasks without improving robustness; the Compressed Tokenizer yields task-specific gains but no average improvement.
  • Mechanistic analysis: 1.27×–2.42× diagonal/off-diagonal ratios show that NGM interactions are predominantly local, supporting token-wise cosine gating without a learned projection.Alignment remains above shuffled and random controls, while diagonal-dominant interactions preserve linear-time cost.
  • Efficiency: 3.4–7.3% prefill and 1.9–2.3% decode overhead occur for 256–1024-token prompts, rising to 16.0% and 9.9% at 2048 tokens.The released implementation recomputes N-gram features over the full prefix; a streaming cache could reduce per-step cost to O(|N|d).

5 Conclusion

NGM is a training-free N-gram memory module that builds causal multi-scale representations from the backbone’s token embeddings and injects them through a non-parametric cosine gate without trainable parameters. Across Qwen3 models from 0.6B to 14B, it improves average performance by +0.5 to +1.2 points, with strongest gains on code generation, while mixed task-wise results expose limits of fixed injection.

  • Conclusion: NGM constructs causal multi-scale N-gram representations from the backbone’s own token embeddings and injects them through a non-parametric cosine gate without trainable parameters.The module is training-free and adds no trainable parameters.
  • Conclusion: +0.5 to +1.2 points in average performance across Qwen3 models from 0.6B to 14B, with the strongest gains on code generation.These results summarize NGM’s overall improvement across the evaluated model scale range.
  • Conclusion: Mixed task-wise results indicate limitations in using a fixed injection rule, motivating future exploration of lightly parameterized alternatives.The conclusion identifies the fixed injection rule as a direction for improvement.

6 Limitations

NGM’s causal N-gram encoder is order-insensitive because it approximates phrases as bags of embeddings, while its heuristic, model-specific injection rule is not suitable for every task.

  • The causal N-gram encoder uses a bag-of-embeddings approximation that is order-insensitive.This can mishandle order-sensitive or noncompositional phrases and inject misleading signals.
  • The cosine gate and model-specific fixed scale λ are heuristic rather than context-adaptive learned components.
  • Mixed task-wise results show that NGM’s training-free injection rule does not suit all tasks or generation.

A NGM implementation

NGM constructs causal n-gram representations from backbone token embeddings, aligns them with cached hidden states, and injects them through cosine-gated residual modulation. The implementation operates on batched sequence tensors and supports configurable n-gram sizes, ReLU gating, and output scaling.

  • A NGM implementation: The implementation aligns the n-gram representations with the L hidden-state positions used by the KV cache.It retains the final L positions, producing an aligned shape of [B, L, N, D].
  • A NGM implementation: NGM averages backbone token embeddings into causal n-gram representations for each configured n-gram size.The resulting tensor has shape [B, T, N, D].
  • A NGM implementation: Cosine gating computes normalized hidden-state-to-n-gram similarities, optionally applies ReLU, and uses them to weight each n-gram representation.The similarities are computed with an einsum over the hidden and n-gram dimensions.
  • A NGM implementation: The module injects the weighted n-gram signal through a residual connection scaled by the configurable output_scale.The output is hidden_states + output_scale * out.

A.1 Model-specific NGM settings

NGM experiments use standardized N = {2, 3} and ReLU gating, with model-specific insertion layers, backbone depths, and output scales defined for converted Qwen3-NGM checkpoints. Layer placements are practical heuristics informed by Engram’s layer-sensitivity analysis [8], not universally optimal settings.

  • A.1 Model-specific NGM settings: All reported models use N = {2, 3} with ReLU gating, while Table 7 specifies model-specific insertion layers, backbone depths, and output scales.Inserted decoder layers are reported using 1-based numbering.
  • A.1 Model-specific NGM settings: For Qwen3 models up to 8B, NGM uses the 15th decoder layer; for the 14B model, the second insertion moves to the 20th layer.These choices reference Engram’s layer-sensitivity analysis [8], Qwen3 backbone depth, and experiment defaults.
  • A.1 Model-specific NGM settings: The layer choices balance early memory injection for local pattern reconstruction against later injection after the hidden state becomes suitable for gating.The placements should be interpreted as practical defaults rather than universally optimal choices.

B Case studies: NGM vs. base model on GSM8K

On GSM8K, Qwen3-8B-NGM solves three representative problems that the base Qwen3-8B model misses, addressing arithmetic, semantic, and reading-comprehension failures. The examples show NGM preserving correct intermediate calculations and prompt-consistent quantity and cost interpretations.

  • Arithmetic error correction: In the gumball problem, NGM computes 4 × 12 = 48 and 48 + 22 = 70, while the base model incorrectly computes 4 × 12 = 28 and answers 50.NGM maintains correct arithmetic in the short-range computation chain.
  • Semantic misinterpretation: In the tadpole problem, NGM correctly computes 11 + 6 − 2 = 15, whereas the base model ignores the six newly visible tadpoles and answers 9.NGM accumulates all quantity changes mentioned in the prompt.
  • Reading comprehension error: In the fair-cost problem, NGM correctly interprets $20.25 as the total ticket expenditure, while the base model treats it as the per-ticket price and inflates downstream totals.NGM preserves local textual coherence across the reasoning chain.
Loading 2605.16893v1…