PaperlayerView exact revision on arXiv

Source-linked AI summary

Tensorizing Engram: Sharing Latents Across N-Gram Embeddings is Beneficial in LLMs

Wuyang Zhou, Yuxuan Gu, Giorgos Iacovides, Yuning Qiu, Qibin Zhao, Danilo Mandic

arXiv:2606.08347v1cs.CLcs.LG

TL;DR

Token-level embeddings leave recurring multi-token patterns to be built implicitly, while independent hash tables cannot share latent structure across nested n-grams. TN-gram represents n-gram embeddings with shared CP factors, and achieves comparable or better performance with fewer parameters than hashing-based solutions at the same order.

  • Problem

    Token-level embeddings leave recurring n-gram patterns to be constructed implicitly across Transformer layers, while independent hash tables prevent nested contexts from sharing latent structure.

    Source: §1 · p. 1, §3 · p. 4

  • Method

    TN-gram tensorizes n-gram spaces with shared CP factors across orders, using nested suffix structure to represent different-order embeddings.

    Source: §1 · p. 2, §4.2 · p. 4

  • Results

    TN-gram achieves comparable or better performance with fewer parameters than hashing-based solutions such as Engram at the same n-gram order.

    Source: §1 · p. 2

  • Takeaways & Limitations

    TN-gram is a compact explicit n-gram memory module for incorporating local multi-token context into Transformer blocks.

    Source: §6 · p. 8

    The current implementation is not fully optimized for wall-clock speed, and the best tensor rank may vary with model scale, vocabulary size, maximum n-gram order, and compute constraints.

    Source: §6 · p. 8

Abstract

from arXiv · show

Modern language models represent text using discrete token-level embeddings, which forces recurring multi-token patterns to be learned implicitly across Transformer layers. Both Over-tokenized Transformers and Engram attempt to address this limitation by explicitly incorporating multi-token (n-gram) memories. However, they rely on separate hash tables for each n-gram order, which introduces hash collisions and prevents nested n-grams from sharing the underlying latent structures. To address these issues, we propose Tensorized Engram (TN-gram), a compact memory module that represents tensorized n-gram embeddings through shared factors in the Canonical Polyadic (CP) form. TN-gram learns shared token-position factors together with order-absorption vectors to encode the embeddings of different n-gram order. Comprehensive experiments demonstrate that TN-gram matches or even outperforms Engram-style n-gram modules while requiring much fewer parameters.

1. Introduction

Explicit n-gram memories address recurring local patterns that token-level embeddings leave implicit, but order-specific hash tables collide and miss nested sharing. TN-gram replaces them with shared CP factors and reports comparable-or-better performance with fewer parameters than hashing-based solutions.

1. Introduction

  • Over-tokenized Transformer expands the vocabulary with higher-order tokens, while Engram injects n-gram information through external memory layers (Liu et al., 2026; Cheng et al., 2026).These approaches have demonstrated that scaling explicit n-gram memories can improve language modeling without much more FLOPs.Source: §1 · p. 1
  • Direct n-gram tables grow exponentially with order, while hash tables make the space tractable but can collide and leave higher-order embeddings untrained.For large n, training may not encounter every possible token combination.Source: §1 · p. 1
  • Separate hash tables for different n-gram orders overlook the intrinsic structure that each (n −1)-gram is contained within higher-order n-grams, suggesting shared latent components across their embeddings.Source: §1 · pp. 1–2
  • TN-gram represents n-gram spaces with shared CP components, reducing parameter growth from O(V n × d) to O(nV R + dR), linear in n for fixed rank R.The shared structure is intended to support training embeddings for higher-order n-grams even when they are unseen in training.Source: §1 · p. 2
  • TN-gram achieves comparable or better performance with fewer parameters than hashing-based solutions such as Engram at the same n-gram order.The comparison is reported for the same n-gram order.Source: §1 · p. 2

2. Related Works

Related work scales language-model representations beyond token embeddings through higher-order tokens, external n-gram memories, and tensor-network parameterizations. These approaches include tensorized compression, low-rank fine-tuning, and structured residual connections.

2. Related Works

  • Over-tokenized Transformer adds higher-order n-gram tokens, while Engram retrieves hashed suffix n-grams into Key/Value representations and gates them using the current hidden state.Separate order-specific tables can collide and prevent related contexts from sharing parameters.Source: §2 · p. 3
  • Tensor-network methods compress neural components using Tensor-Train, Tucker, or Kronecker structure, and also support low-rank parameter-efficient fine-tuning.The cited examples include TensorLLM, TensorGPT, LoRETTA, TeRA, and KromHC.Source: §2 · p. 3

3. Preliminaries

The preliminaries introduce tensor notation and CP decomposition, then describe Engram’s hashed n-gram retrieval and context-aware gating. Engram’s order-specific tables prevent latent sharing between nested contexts.

3. Preliminaries

  • A tensor is a multidimensional array, and a mode-n product multiplies it by a matrix along its nth mode.The product replaces that mode’s dimension with the matrix’s row dimension.Source: §3 · p. 3
  • CP decomposition approximates a tensor as a sum of R rank-one components, using factor matrices for its modes.Its parameter count is reduced from the product of mode dimensions times R to a sum of mode dimensions times R.Source: §3 · p. 3
  • Engram hashes local suffix contexts for orders 2 through N into order-specific tables, concatenates retrieved embeddings across orders and heads, then projects them into Key and Value vectors.A context-aware gate uses the current hidden state to filter the retrieved Value before it is added to the residual stream.Source: §3 · pp. 3–4

4. Methodology

TN-gram replaces separate hashed n-gram tables with a shared tensorized CP representation, using nested suffix contexts to share latent structure across orders. Its formulation derives lower-order embeddings from shared factors and integrates them into Transformer memory projections and gating.

4.1. Suffix n-gram Extraction

  • Suffix contexts ending at each token index n-gram memories for orders 2 through N, with left padding when they extend before the sequence.Unigram information is left to the standard token embedding.Source: §4.1 · p. 4

4.2. Tensorizing n-gram Embeddings in TN-gram

  • TN-gram tensorizes each order-n embedding table, sharing CP factors across orders because lower-order suffixes are nested within higher-order contexts.The token-position modes represent the context tokens, while the final mode indexes the latent dimension.Source: §4.2 · p. 4

4.3. TN-gram Formulation

  • The N-gram tensor uses shared token-position factor matrices, a shared model-dimension factor, and a super-diagonal CP core.The factor matrices are ordered from oldest to newest token.Source: §4.3 · p. 4
  • Lower-order embeddings contract the oldest token modes using learned order-absorption vectors, while RMSNorm and a learned scalar stabilize and weight each order's token-space vector.The resulting vector can be mapped to model dimension with F; the method absorbs F into later Key/Value projections instead of materializing that multiplication.Source: §4.3 · pp. 4–5
  • Validation loss gains largely saturate at 5-gram, while parameter counts grow roughly linearly with order and remain lower for TN-gram than Engram at CP rank 1200.Improvements beyond 5-gram are marginal and can sometimes reverse.Source: Figure 2 · p. 5
  • TN-gram concatenates token-space vectors and absorbs the model-dimension factor into Key/Value projections before applying Engram-style context-aware gating.The current hidden state and Key vector determine a scalar gate; the gated Value passes through a depthwise causal convolution and is added to the residual stream.Source: §4.3 · p. 5

4.4. Sharing Latents across n-gram Embeddings

  • When the relevant token-position factor has full column rank, a shorter-context extraction vector is a linear combination of the corresponding longer-context vector.The stated condition is that A_(N−n+1) has full column rank.Source: §4.4 · p. 6

5. Experiments

Across two model depths, TN-gram matches or exceeds Engram on reported benchmarks with fewer added parameters; experiments also examine n-gram order, capacity scaling, vocabulary size, and component ablations.

5. Experiments

  • Validation loss generally improves through N = 5, then gains are marginal or sometimes reverse; at N = 2, TN-gram is slightly worse than Engram.The authors use N = 5 for the main comparisons in Tables 1–3.Source: §5 · p. 6
  • TN-gram has the best CORE at both depths with fewer added parameters; 9-layer validation BPB falls from 1.251 to 1.208, while 18-layer CORE rises from 0.107 to 0.120 and BPB is 1.071 vs. 1.070.Both models use a 5-gram configuration.Source: §5 · p. 6
  • TN-gram leads multiple-choice averages at both depths (33.48 to 35.44 at 9 layers; 36.85 to 37.82 at 18), and also leads language-modeling and BIG-bench averages.Source: §5 · p. 7
  • More Engram hashing heads yield only modest validation-loss gains, whereas increasing TN-gram rank consistently improves loss by enriching shared latent components.Source: Figure 3 · p. 7, §5 · p. 7

5.1. Ablation Study

  • At vocabulary size 8192, TN-gram loss reductions (4.27% training, 4.44% validation) exceed Engram’s (3.99%, 4.17%).Source: §5.1 · pp. 7–8
  • Removing any stabilizing component raises validation loss; removing learnable scales hurts more than removing order-absorption vectors alone, and removing both is worst.The ablated components are RMS normalization, learnable per-order scaling, and order-absorption vectors.Source: §5.1 · p. 8, Figure 4 · p. 8

6. Conclusion

TN-gram uses shared CP factors to represent n-gram contexts, letting different orders share latent structure. Its current implementation trades parameter efficiency for unoptimized wall-clock speed, and its best tensor rank may depend on the setting.

6. Conclusion

  • TN-gram replaces separate hash tables with shared CP factors across tensorized n-gram spaces.Common token-position factors and learnable order-absorption vectors let lower- and higher-order contexts share latent structure.Source: §6 · p. 8
  • This shared-factor design removes hash collisions and supports sharing latent structure across n-gram orders.The conclusion identifies common token-position factors and order-absorption vectors as the shared components.Source: §6 · p. 8
  • TN-gram is more parameter-efficient, but its current implementation is not fully optimized for wall-clock speed.It performs factor gathering, Hadamard products, normalization, and scaling beyond hash-based Engram.Source: §6 · p. 8
  • The experiments match tensor rank to Engram’s parameter budget, although the best rank may vary with model scale, vocabulary, maximum n-gram order, and compute constraints.Adaptive rank-selection strategies remain an open direction.Source: §6 · p. 8

Impact Statement

The paper presents parameter efficiency in explicit n-gram memories as a potential route toward more efficient and sustainable language models.

Impact Statement

  • TN-gram may improve parameter efficiency for explicit n-gram memories.The impact statement presents this as a potential societal benefit.Source: Impact Statement · p. 8
  • More efficient and sustainable language models are a potential downstream impact.The statement frames this as a possibility, not a demonstrated outcome.Source: Impact Statement · p. 8
  • Compact modeling of recurring multi-token patterns may reduce parameter overhead in neural network layers.Source: Impact Statement · p. 8

A. Proof of Theorem 1

The proof studies how dropping the oldest token relates shorter and higher-order n-gram extraction vectors, using shared suffix factors and a full-column-rank condition.

A. Proof of Theorem 1

  • Theorem A.1 considers the shorter context formed by dropping the oldest token under a full-column-rank condition on A_N-n+1.The theorem concerns the relationship between that shorter context’s token-space extraction vector and the higher-order context.Source: A. Proof of Theorem 1 · p. 12
  • The proof’s shared suffix vector collects the factors common to the n-gram orders n and n−1.It is formed from order vectors and token-position factors for the shared suffix.Source: A. Proof of Theorem 1 · p. 12
  • Full column rank makes the rows of A span R^R, allowing coefficients that represent w and depend only on w and A.Source: A. Proof of Theorem 1 · p. 12

B. Downstream Tasks

The downstream evaluation uses a broad task suite and CORE to assess generalization, then compares final-step performance across n-gram orders in 18-layer transformers. TN-gram has the lowest validation loss and higher CORE than Engram and the raw GPT baseline, with its strongest CORE gain at larger orders.

B. Downstream Tasks

  • The downstream suite spans commonsense and scientific reasoning, reading comprehension, language modeling, symbolic reasoning, and factual knowledge.The benchmarks are used to test whether validation-loss improvements translate into broader downstream generalization.Source: B. Downstream Tasks · p. 12
  • CORE averages task accuracies after centering and normalizing scores so random guessing maps to 0 and perfect accuracy to 1.This makes results more comparable across tasks with different chance levels.Source: B. Downstream Tasks · p. 12

C. Scaling n-gram Order in a 19-Layer Transformer

D. Context-Aware Gating

The context-aware gate compares the current hidden state with the memory key, gates memory values according to alignment, then convolves them before adding the memory output to the residual stream.

D. Context-Aware Gating

  • The gate compares the current hidden state with the memory key and suppresses memory values that are poorly aligned with the hidden state.The hidden state and memory key are the inputs to the alignment comparison.Source: D. Context-Aware Gating · p. 13
  • A normalized alignment score is transformed into a scalar gate, and a signed square-root transformation reduces extreme dot-product magnitudes while preserving their sign.The gated memory value is the scalar gate multiplied by the memory value.Source: D. Context-Aware Gating · p. 13
  • A short depthwise causal convolution with kernel size 3, dilation N, and SiLU activation processes gated values before the memory output is added to the residual stream.The convolution is applied to the sequence of gated values over time.Source: D. Context-Aware Gating · p. 13

E. Hyperparameters

The models share a 9-layer Transformer configuration and are trained for 6000 steps, with optimizer choices and learning-rate schedules specified for different parameter groups.

E. Hyperparameters

  • Models train for 6000 steps on length-1024 sequences with a global batch of 524,288 tokens, using a 9-layer Transformer with hidden size 512.Engram and TN-gram memory modules are inserted in layers 1 and 7; training uses either 4 × NVIDIA A100 (80GB) GPUs or one L40S GPU (48GB), depending on model size.Source: E. Hyperparameters · p. 13, Table 4 · p. 14
  • Muon optimizes matrix-shaped Transformer parameters, while Adam optimizes embeddings, scalar/control parameters, and the language-model head.Muon uses learning rate 0.04; Adam uses β1 = 0.9, β2 = 0.95, and ϵ = 10^-8.Source: E. Hyperparameters · pp. 13–14, Table 5 · p. 14
  • Learning rates remain constant until the final 1200 steps, then linearly decay to zero.The tied embedding table uses learning rate 0.05, while scalar/control parameters use learning rate 0.04.Source: E. Hyperparameters · p. 14

Related papers

Loading 2606.08347v1…