PaperlayerView exact revision on arXiv

Source-linked AI summary

Lngram v2: Latent N-Gram Memory with Interpretable Discrete Representations

Yunao Zheng, Bin Wen, Xiaojie Wang

arXiv:2609.03426v1cs.CL

TL;DR

Transformers lack native lookup for recurring local patterns, while Lngram v1 couples memory capacity to backbone width. Lngram v2 decouples these dimensions with context-aware readout and hard discrete routing, improving scaled vision–language models while also making discrete IDs a semantic interface.

  • Problem

    Transformers lack a native lookup mechanism for local static patterns, and Lngram v1 couples route count with backbone width while leaving the semantic content of its IDs unclear.

  • Method

    Lngram v2 decouples route count, memory dimension, and backbone width, using grouped-query attention, a zero-value Sink, and counterfactual surrogate gradients with hard discrete addressing.

  • Results

    Lngram v2 consistently improves vision–language models across scales, reaches the 30B-parameter regime, reduces total and activated memory parameters by 82.6% and 95.2% versus Lngram v1, and preserves semantic information in discrete IDs.

  • Takeaways & Limitations

    Lngram v2 provides scalable latent conditional memory and a structured, queryable interface for analyzing internal model representations.

  • Takeaways & Limitations

    Lngram v1’s scalability is limited because its route count is coupled with backbone width, increasing memory size and readout cost with model dimension.

Abstract

from arXiv · show

Transformers lack a native lookup mechanism, requiring repeated dense computation to recognize and reuse local static patterns. Lngram v1 introduces tokenizer-independent conditional memory through discrete latent n-gram addressing, but its memory capacity is coupled with the backbone width, limiting scalability due to high parameter and activation costs. We propose Lngram v2, which decouples the number of routes, memory dimension, and backbone width, and introduces a context-aware grouped-query attention readout to scale memory capacity independently. A zero-value Sink and counterfactual surrogate gradients further improve readout selectivity and routing trainability while preserving hard discrete addressing. Experiments across vision--language models (VLMs) of different scales show consistent improvements, including successful scaling to a 30B-parameter model. Compared with Lngram v1, Lngram v2 substantially reduces both total and activated memory parameters while maintaining or improving language modeling performance. Further analysis shows that its discrete IDs preserve substantial semantic structure of continuous hidden states, enabling semantic recovery from IDs alone and stable ID--semantic associations across datasets. These results establish Lngram v2 as an efficient and scalable latent conditional memory mechanism whose discrete addresses also provide a structured interface for analyzing internal model representations.

1 INTRODUCTION

Lngram v2 addresses Transformers’ lack of native lookup by providing scalable latent conditional memory and a discrete interface for representation analysis. It decouples memory capacity from backbone width, improves vision–language models across scales, and preserves semantic information in discrete IDs.

  • Transformers lack a native lookup mechanism for recurring local patterns, so these patterns are handled through dense computation alongside context-dependent reasoning.
  • Lngram v1 uses learnable latent symbols for exact n-gram addressing, but its route count remains coupled to backbone width and its IDs’ semantic preservation is unclear.
  • Lngram v2 decouples route count, memory dimension, and backbone width, using grouped-query attention for context-dependent readout and counterfactual surrogate gradients for hard discrete routing.
  • 65.77%–84.27% of the excess semantic readout capability of continuous hidden states is retained by an ID-only semantic reader.
  • Across vision–language models, Lngram v2 consistently improves performance, scales to 30B parameters, and reduces Lngram v1’s total and activated memory parameters by 82.6% and 95.2%.

2 ARCHITECTURE

Lngram v2 maps hidden states to multiple hard discrete symbols, forms exact latent n-gram addresses, and retrieves route-wise memory tokens. Context-aware grouped-query attention selects among those tokens, while counterfactual lookup gradients train the discrete routing without changing the hard forward path.

  • 2.1 OVERVIEW: Lngram v2 inserts a conditional memory branch into Transformer decoder layers to offload local pattern matching and storage from the backbone.
  • 2.2 MULTI-ROUTE DISCRETE ADDRESSING: Each hidden state is mapped to multiple discrete symbols, whose local n-grams form exact route-specific memory addresses applicable across modalities.
  • 2.3 EXACT N-GRAM RETRIEVAL: Different routes use disjoint address ranges, yielding unique route–n-gram entries while masking invalid prefixes and sequence-boundary-crossing n-grams.
  • 2.4 CONTEXT-AWARE READOUT: Retrieved entries are organized as route-wise memory tokens, and the current hidden state contextually selects among them through grouped-query attention.
  • 2.4 CONTEXT-AWARE READOUT: A fixed zero-value Sink prevents unreliable memories from receiving forced probability mass during readout.
  • 2.5 BACKPROPAGATION THROUGH DISCRETE ADDRESSING: Counterfactual surrogate gradients train the discretization projection while preserving hard discrete addressing and exact lookup in the forward pass.

3 ANALYSIS

This analysis tests whether Lngram’s discrete IDs preserve and expose the semantics of continuous hidden states without using those states directly. The results show substantial semantic recovery, reproducible route-level associations, and stronger compositional information in complete multi-route signatures.

  • Research questions: Lngram v2 maps continuous hidden states to discrete IDs and tests whether their semantics remain recoverable from IDs alone.The analysis also examines whether semantic information is localized to individual routes or distributed across multiple routes.
  • Method: The study uses nearest-neighbor retrieval over 128-route signatures, with continuous hidden-state retrieval as the reference and nuisance-only predictions as the baseline.For each held-out query, K = 10 nearest distinct reference documents are selected by route-Hamming distance.
  • Semantic readout: 82.21%–84.27% of continuous-state excess AP is retained at Layer 1, while 65.77%–73.84% is retained at Layer 13.Across all six settings, discrete IDs recover a substantial fraction of the readable semantics in continuous hidden states, although discretization loses some information.
  • Route-level associations: 9,220 of 14,002 evaluable candidate associations satisfy held-out consistency criteria, yielding a 65.85% overall replication rate.Candidate associations and effect directions are determined on the reference set and independently evaluated on held-out data.
  • Route-level associations: Individual route codes show reproducible concept associations, including giraffe and person, under the held-out validation protocol.The giraffe association rises from 2.37% to 21.70%, while the person association rises from 53.60% to 63.25%.
  • Compositional structure: 31.79%–54.82% of continuous-state excess AP is retained by accumulated marginal route effects, versus 65.77%–84.27% for complete-signature ID-KNN.Using all 128 routes also outperforms the single best route, with AP ratios ranging from 2.02 to 12.32, indicating distributed joint encoding.

4 EXPERIMENTS

Experiments show that Lngram v2 improves vision–language performance across model scales while reducing conditional-memory costs and enabling flexible capacity–deployment trade-offs.

  • Main results: Lngram v2 improves Keye2B’s average score from 47.25 ± 0.15 to 47.93 ± 0.18, a mean gain of 0.68 points across three seeds.It exceeds PKM by 0.69 points and outperforms a parameter-matched Sparse FFN, whose average score is 45.86.
  • Main results: Lngram v2 improves Keye30B’s ten-benchmark average from 78.23 to 79.71, a gain of 1.48 points, demonstrating scaling to the 30B regime.The 95% confidence interval for the average improvement is [0.86, 2.11].
  • Ablations: Increasing routes from r16 to r64 raises the average from 54.20 to 55.71 while activated parameters increase only from 18.121M to 18.932M.At similar total parameter counts, increasing discrete routes outperforms increasing memory dimension alone by 1.02 points with roughly half the activated parameters.
  • Ablations: Readout capacity matters: r121/m4/KV16 achieves the best average score of 56.16, whereas a configuration adding more than 1B parameters remains below r64/m4/KV8.The results indicate diminishing returns from simply enlarging memory capacity.
  • Ablations: The zero-value Sink raises the Keye30B ten-benchmark average from 79.39 to 79.71 and improves MMBench-EN from 85.91 to 86.92.The Sink is enabled by default.
  • Inference efficiency: For r64/m4/KV8, prefill latency at length 1024 rises 6.1% and decode latency rises 8.0%, while the three-benchmark average improves by 2.07 points.The lighter r16/m4/KV8 configuration incurs approximately 3.7% prefill overhead at length 1024.
  • Parameter efficiency: Relative to Lngram v1, Lngram v2 reduces total module parameters by 82.6% and activated parameters per token by 95.2% while achieving slightly lower validation loss.At R = 16, validation loss is 2.8584 versus 2.8609 for v1.

5 CONCLUSION

Lngram v2 combines scalable hard discrete memory with surrogate-gradient training and exposes a route-based interface for analyzing internal representations.

  • Conclusion: Lngram v2 decouples route count, memory dimension, and backbone width, then reads retrieved n-gram memories with context-aware grouped-query attention.A zero-value Sink suppresses unreliable retrievals, while route-wise streaming controls intermediate activations at large scale.
  • Surrogate gradients: The exact surrogate expectation changes one local n-gram symbol at a time while keeping the other n −1 symbols hard, reducing enumeration to K candidate symbols.With 16 candidates per local position, the computation remains practically manageable.
  • Surrogate gradients: The one-bit approximation reduces counterfactual candidates from 16 to 4 when M = 4 by flipping one bit at a time.This lowers additional lookup overhead during training.
  • Implementation: In streaming mode, route chunks are processed with online-softmax accumulation while using the same surrogate formula as full-lookup mode.The resulting gradients accumulate into the complete routing logits.
  • Training procedure: Counterfactual lookups replace only routing-logit gradients: the hard forward lookup and its selected memory-entry gradients remain unchanged.Counterfactual table entries are treated as constants during surrogate computation.

A.5 DIRECT STE FAILS TO EFFECTIVELY OPTIMIZE DISCRETE ROUTING

Direct STE fails to optimize Lngram v2's discrete routing effectively because it ignores how changing addresses changes retrieved memory contents. The counterfactual surrogate gradient instead incorporates those lookup changes and yields stable training.

  • A.5 DIRECT STE FAILS TO EFFECTIVELY OPTIMIZE DISCRETE ROUTING: Direct STE ignores the memory contents retrieved after switching to a different discrete address, whereas the surrogate gradient uses candidate-address lookup results.
  • A.5 DIRECT STE FAILS TO EFFECTIVELY OPTIMIZE DISCRETE ROUTING: Direct STE quickly stalls at substantially higher training loss than Lngram v2's surrogate gradient.The controlled ablation keeps architecture, forward computation, data, and optimization configuration identical except for the gradient estimator.
  • A.5 DIRECT STE FAILS TO EFFECTIVELY OPTIMIZE DISCRETE ROUTING: Validation losses are 2.8585 with the surrogate gradient and 7.3349 with direct STE.
  • A.5 DIRECT STE FAILS TO EFFECTIVELY OPTIMIZE DISCRETE ROUTING: The surrogate-gradient advantage appears in both training and validation loss under the controlled comparison.

B.2 CROSS-CHECKPOINT STABILITY OF DISCRETE ADDRESSES

Discrete-address stability differs substantially across model families during training, especially for visual tokens at deeper layers. The newer VLM's lower stability reflects changing continuous representations and makes coherent memory formation more difficult.

  • B.2 CROSS-CHECKPOINT STABILITY OF DISCRETE ADDRESSES: At Layer 13, visual route-code agreement is 0.9246 for Keye2B versus 0.7702 for the new VLM.
  • B.2 CROSS-CHECKPOINT STABILITY OF DISCRETE ADDRESSES: At Layer 13, visual 3-gram survival is 0.7945 for Keye2B versus 0.4754 for the new VLM.
  • B.2 CROSS-CHECKPOINT STABILITY OF DISCRETE ADDRESSES: Relative to Keye2B, visual code agreement decreases by 0.1024 at Layer 1 and 0.1544 at Layer 13.The corresponding 95% confidence intervals are [−0.1035, −0.1013] and [−0.1568, −0.1519].
  • B.2 CROSS-CHECKPOINT STABILITY OF DISCRETE ADDRESSES: Visual 3-gram survival decreases by 0.2579 at Layer 1 and 0.3190 at Layer 13, while text differences are substantially smaller.
  • B.2 CROSS-CHECKPOINT STABILITY OF DISCRETE ADDRESSES: Using each checkpoint's own router reproduces the stability gap, indicating that changing continuous representations, not router parameters alone, primarily drives it.

B.3 LAYER-WISE PROPAGATION OF LNGRAM V2 RESIDUALS

Lngram v2's effect on final representations and predictions is much stronger in mature Keye2B than in the newly assembled VLM. The main difference is smaller initial injection, associated with less stable discrete addresses during cross-modal adaptation.

  • B.3 LAYER-WISE PROPAGATION OF LNGRAM V2 RESIDUALS: At Layer 13, the new model's immediate visual injection R is 5.99 × 10^-3, approximately 1/9.6 of the mature model's value.At text positions, it is 9.26 × 10^-5, approximately 1/68.7 of the mature-model value.
  • B.3 LAYER-WISE PROPAGATION OF LNGRAM V2 RESIDUALS: Disabling Lngram v2 changes normalized next-token logits by 0.1928 in Keye2B versus 0.0174 in the new model.
  • B.3 LAYER-WISE PROPAGATION OF LNGRAM V2 RESIDUALS: Keye2B has JS divergence 1.25 × 10^-2 after disabling Lngram v2, compared with 2.97 × 10^-4 for the new model.
  • B.3 LAYER-WISE PROPAGATION OF LNGRAM V2 RESIDUALS: Disabling Lngram v2 changes the top-1 prediction for 6/48 Keye2B documents but none of the 48 new-model documents.
  • B.3 LAYER-WISE PROPAGATION OF LNGRAM V2 RESIDUALS: The weaker new-model effect is not explained by rapid downstream attenuation, because its initial injections are already substantially smaller.
  • B.3 LAYER-WISE PROPAGATION OF LNGRAM V2 RESIDUALS: Pretrained vision and language components alone do not guarantee immediate Lngram v2 effectiveness while cross-modal representations remain unstable.

C.2 ROUTE-CODE SEMANTIC ASSOCIATIONS AND REPRESENTATIVE EXAMPLES

Lngram route codes exhibit reproducible semantic associations, both across complete signatures and within individual routes. These associations are tested on held-out data and include interpretable concepts such as animals and people.

  • C.2 ROUTE-CODE SEMANTIC ASSOCIATIONS AND REPRESENTATIVE EXAMPLES: Lngram's discrete addresses are analyzed as native memory-generated codes rather than separately trained probes or dictionary units.
  • C.2 ROUTE-CODE SEMANTIC ASSOCIATIONS AND REPRESENTATIVE EXAMPLES: Candidate code–semantic relations are selected on reference data and independently tested on held-out data under a strict separation protocol.Of 16,304 reference candidates, 14,002 are evaluable on held-out data and 9,220 satisfy the reported statistical requirements.
  • C.2 ROUTE-CODE SEMANTIC ASSOCIATIONS AND REPRESENTATIVE EXAMPLES: A representative code raises held-out giraffe probability to 21.70% from a 2.37% nuisance-controlled baseline, a gain of 19.33 percentage points.
  • C.2 ROUTE-CODE SEMANTIC ASSOCIATIONS AND REPRESENTATIVE EXAMPLES: Another representative code yields a held-out person probability of 63.25%, 9.65 percentage points above baseline.
  • C.2 ROUTE-CODE SEMANTIC ASSOCIATIONS AND REPRESENTATIVE EXAMPLES: A route's different discrete values can distinguish related semantic states rather than representing one fixed concept.Codes 9, 10, and 11 on route 75 associate strongly with horse, giraffe, and zebra, respectively.

C.3 DETAILED ANALYSIS OF DISTRIBUTED SEMANTICS

Lngram’s individual route codes carry stable semantic signals, but complete multi-route signatures recover substantially richer semantics through complementary joint information. These discrete IDs can therefore support held-out semantic prediction and structured analysis of internal representations.

  • Individual route codes provide stable semantic tendencies, while complete IDs offer substantially stronger semantic readout.
  • Held-out semantic reading uses only discrete IDs and tests whether reference-set route associations compose into representations that generalize to unseen samples.
  • Summing independent single-code effects outperforms the nuisance baseline across all six settings and retains 31.79%–54.82% of Hmodel’s excess AP.
  • The complete-signature advantage reflects complementary information across multiple routes rather than one consistently dominant route.
  • Complete-ID KNN retains 65.77%–84.27% of the continuous representation’s excess AP, versus 31.79%–54.82% for independent route-code marginal effects.
  • Lngram IDs function both as exact memory addresses and as a statistically analyzable, queryable, and reproducible interface to internal representations.

D.2 COMPLETE INFERENCE-EFFICIENCY RESULTS

Lngram v2 adds modest online inference cost in standard configurations, while very large capacity settings impose substantially higher long-sequence overhead. Moderate configurations therefore offer the more favorable practical trade-off among performance, capacity, and computation.

  • At r64/m4/KV8 and sequence length 1024, prefill overhead is approximately 6.1% and decode overhead approximately 8.0%.Prefill latency rises from 56.99 ms to 60.49 ms, while decode latency rises from 48.26 ms/token to 52.14 ms/token.
  • The same r64/m4/KV8 configuration increases memory usage from 5.24 GiB to 5.46 GiB for prefill and from 4.89 GiB to 5.05 GiB for decode.
  • The extremely large r512/m4/KV8/mem256 configuration incurs substantially higher long-sequence memory usage and prefill overhead.
  • Moderate-scale configurations provide a more favorable trade-off among downstream performance, parameter capacity, and online computation cost.

E HYPERPARAMETER SETTINGS

The Keye2B and Keye30B experiments use shallow and intermediate Lngram insertions with 4-bit codes, 2/3-gram memory, and multi-head readout. The Keye30B setup uses distributed parallel training over 64 GPUs and a 32,768-token sequence length.

  • Lngram v2 uses 4-bit discrete codes, 2/3-gram memory, and multi-head readout in both Keye2B and Keye30B experiments.
  • In both model scales, Lngram v2 is inserted into a shallow and an intermediate backbone layer.
  • The Keye30B experiment uses 8 nodes and 64 GPUs with TP = 1, PP = 4, CP = 16, EP = 8, and ETP = 1.
  • Keye30B training uses sequence length 32,768, micro batch size 1, global batch size 32, and approximately 1.049M tokens per step.
  • The formal Keye30B comparison uses the step-5000 checkpoint after approximately 5.24B cumulative training tokens.
Loading 2609.03426v1…