Beyond N-gram: Data-Aware X-GRAM Extraction for Efficient Embedding Parameter Scaling

1. Problem Statement

Large Language Models (LLMs) typically scale capacity by increasing model size, which linearly increases inference-time compute (FLOPs). An alternative approach is token-indexed memory, where models retrieve features from a large lookup table based on token IDs. However, prior approaches (e.g., MoLE, Retoken, Engram) face three critical limitations:

1. Zipfian Under-training: Token frequencies follow a long-tail (Zipfian) distribution. In naive lookup tables, high-frequency tokens dominate training updates, while the vast majority of "long-tail" rows remain under-trained or near initialization, leading to poor parameter efficiency.
2. Slot Collapse & Redundancy: Simply adding more parallel lookup slots (fixed $n$-gram tables) often results in "slot collapse," where different slots learn highly correlated, redundant representations because they are trained on overlapping token occurrences.
3. Memory vs. Compute Mismatch: Naive scaling requires massive tables (often terabytes) that do not fit in High-Bandwidth Memory (HBM), forcing reliance on host memory or SSDs. This creates a memory traffic bottleneck that cannot be hidden behind computation, especially during decoding.

The core question addressed is: How can we design a lookup-based scaling axis that is trainable under Zipfian statistics, diverse in representation, and efficient in memory usage?

2. Methodology: X-GRAM

The authors propose X-GRAM, a data-aware framework that converts static token-indexed retrieval into dynamic, diverse local $x$-gram features. The framework consists of three core stages:

A. Frequency-Aware Hybrid Hash Mapping (Stage 1)

To address the long-tail under-training problem, X-GRAM replaces uniform table allocation with a probability-balanced routing mechanism:

• VIP Reservation: The most frequent tokens (the "head") are assigned dedicated physical rows in a reserved VIP region to ensure they receive sufficient updates.
• Balanced Bucketing: The remaining long-tail tokens are partitioned into logical buckets based on smoothed frequency mass ($p(\omega)^\alpha$), ensuring each bucket has roughly equal total probability mass.
• Hybrid Hashing: Within dense buckets, tokens are mapped to physical rows using multi-path local hashing and alias mixing. This compresses the tail (e.g., 50% compression ratio) while distributing collision and update pressure more evenly, preventing "hotspot" domination.

B. Gated ShortConv Extraction (Stage 2)

To prevent slot collapse and extract diverse features, X-GRAM introduces a lightweight refinement module:

• Multi-Scale ShortConv: Instead of using raw retrieved vectors, the model applies depthwise convolutions with varying kernel sizes ($k$) to the retrieved sequence. This transforms static 1-gram embeddings into dynamic local $x$-gram features (e.g., capturing "not good" vs. "good").
• Gated Mechanism: A multiplicative gate (SiLU) modulates the convolution output. This acts as a context-selective filter, suppressing noise from hash collisions and amplifying relevant local patterns.
• Diversity: Different views use different kernel sizes, breaking symmetry across parallel memory slots and encouraging diverse subspaces.

C. Depth-Aware Contextual Injection (Stage 3)

To align static memory with dynamic context efficiently:

• Injection Sites: The refined signals are injected primarily into the Attention Value stream (where dimensionality $d_{kv}$ is smaller, reducing activation/compute costs) and optionally into Inter-layer Residuals (as a stable identity anchor).
• Depth-Aware Gating: The injection strength is modulated by a learnable gate that scales with layer depth ($\sqrt{\ell+1}$) and a warmup schedule. This ensures the injected signal remains stable and proportional as the network deepens, preventing early training instability.

3. Key Contributions

1. Diagnosis of Scaling Limits: The paper identifies and quantifies that naive lookup scaling fails due to Zipfian under-training and representational redundancy (slot collapse), rather than a lack of parameters.
2. X-GRAM Architecture: A novel framework combining frequency-aware hashing (for trainability), gated ShortConv (for diversity and $x$-gram extraction), and depth-aware injection (for stability).
3. Memory-Centric Scaling: X-GRAM decouples model capacity from FLOPs. It achieves significant performance gains with substantially smaller lookup tables (e.g., 50% of the size of baselines) by optimizing parameter efficiency rather than just raw size.
4. System Efficiency: By injecting into the value stream and using compression, X-GRAM reduces memory traffic and activation footprints, making it feasible to offload tables to host memory without severe performance penalties.

4. Experimental Results

The authors evaluated X-GRAM on 0.73B and 1.15B parameter models, pre-trained on 41B tokens.

• Performance Gains:
  • At the 0.73B scale, X-GRAM-50% (using 50% of the table size of baselines) achieved an average downstream accuracy of 48.5, outperforming strong baselines like MoRT (45.8) and Engram (47.2) by 2.7–3.2 points.
  • At the 1.15B scale, X-GRAM-50% reached 48.2 (1× config) and 50.8 (4× config), surpassing baselines by 2.0–2.7 points.
• Parameter Efficiency: X-GRAM-50% achieved higher accuracy than baselines using half the lookup table parameters. For example, at 0.73B, X-GRAM-50% (1.17B extra params) beat Engram (2.33B extra params).
• Scaling Behavior: Unlike baselines which often plateau as table size increases, X-GRAM continues to improve with larger capacity, indicating that additional memory is converted into useful, distinct features rather than redundant slots.
• Training Stability: The model showed stable convergence with a persistent loss advantage over baselines and reached baseline-quality performance using only 57.25% of the training data.

5. Significance

X-GRAM offers a practical and scalable paradigm for memory-augmented architectures. By shifting the focus from "more parameters" to "better parameter utilization," it solves the fundamental bottlenecks of token-indexed scaling:

• It makes large lookup tables trainable under realistic data budgets by addressing the long-tail problem.
• It ensures diversity in learned representations, preventing the diminishing returns associated with fixed-slot scaling.
• It provides a compute-decoupled scaling path that is memory-efficient, potentially enabling future LLMs to scale capacity without the prohibitive cost of increased FLOPs or massive on-device memory requirements.

In summary, X-GRAM demonstrates that with the right data-aware design (frequency-aware routing + local feature extraction), token-indexed memory can be a highly effective and efficient scaling axis for next-generation LLMs.