[論文レビュー] Conditional Memory via Scalable Lookup: A New Axis of Sparsity for Large Language Models

Xin Cheng, Shengding Hu+12 more|arXiv (Cornell University)|Jan 12, 2026

Topic Modeling被引用数 0

ひとこと要約

要旨: 本論文は Engram を導入する。Engram は MoE を補完する条件付きメモリモジュールで、拡張可能な N-gram メモリルックアップを提供し、メモリと計算の間に U 字型のスパース性割り当てを見出し、推論、コード/数学、長文コンテキストタスクで強力な利得を達成する。

ABSTRACT

While Mixture-of-Experts (MoE) scales capacity via conditional computation, Transformers lack a native primitive for knowledge lookup, forcing them to inefficiently simulate retrieval through computation. To address this, we introduce conditional memory as a complementary sparsity axis, instantiated via Engram, a module that modernizes classic $N$-gram embedding for O(1) lookup. By formulating the Sparsity Allocation problem, we uncover a U-shaped scaling law that optimizes the trade-off between neural computation (MoE) and static memory (Engram). Guided by this law, we scale Engram to 27B parameters, achieving superior performance over a strictly iso-parameter and iso-FLOPs MoE baseline. Most notably, while the memory module is expected to aid knowledge retrieval (e.g., MMLU +3.4; CMMLU +4.0), we observe even larger gains in general reasoning (e.g., BBH +5.0; ARC-Challenge +3.7) and code/math domains~(HumanEval +3.0; MATH +2.4). Mechanistic analyses reveal that Engram relieves the backbone's early layers from static reconstruction, effectively deepening the network for complex reasoning. Furthermore, by delegating local dependencies to lookups, it frees up attention capacity for global context, substantially boosting long-context retrieval (e.g., Multi-Query NIAH: 84.2 to 97.0). Finally, Engram establishes infrastructure-aware efficiency: its deterministic addressing enables runtime prefetching from host memory, incurring negligible overhead. We envision conditional memory as an indispensable modeling primitive for next-generation sparse models.

研究の動機と目的

• MoE（条件付き計算）を超える補完的スパース軸を、静的知識ルックアップを活用してLLMに導入する。
• N-gram 埋め込みを、Transformer に統合されたスケーラブルで微分可能なメモリ機構として再検討する。
• 固定予算下でメモリと計算のバランスを取る Sparsity Allocation 問題を定式化する。
• Engram は数十億パラメータ規模へ拡張可能で、推論、知識、コード/数学タスクで利得を生み出すことを示す。）

提案手法

• Engram を、静的 N-gram 埋め込みをマルチヘッドハッシュと決定論的アドレッシングで取得する条件付きメモリモジュールとして提示する。
• 語彙サイズを削減し堅牢な N-gram サフィックスを実現するトークナイザ圧縮を実装する。
• 現在の隠れ状態と取得メモリを調整する文脈認識ゲーティング機構を用いる。
• 取得メモリを動的バックボーンと多分岐対応統合および残差接続で融合する。
• メモリと計算を分離してホストメモリオフロードと決定論的プリフェッチを可能にし、効率化を図る。
• 固定計算量の下で MoE 専門家と Engram メモリの最適分割を見つける Sparsity Allocation フレームワークを用いて割り当てを分析し、U字型のスケーリング法則を明らかにする。）

Figure 1 : The Engram Architecture. The module augments the backbone by retrieving static $N$ -gram memory and fusing it with dynamic hidden states via context-aware gating. This module is applied only to specific layers to decouple memory from compute, leaving the standard input embedding and un-em

実験結果

リサーチクエスチョン

• RQ1 固定されたパラメータおよび計算予算の下で、条件付き計算（MoE）と条件付きメモリ（Engram）間でモデル容量をどのように割り当てるべきか？
• RQ2 拡張可能なメモリプリミティブの導入は、知識・推論・長文タスクのパフォーマンス改善において MoE を補完するか？
• RQ3 無限メモリまたは大規模メモリ予算下で Engram のスケーリング挙動はどうなるか？
• RQ4 ストレージと計算を分離しホストメモリオフロードを可能にすることで Engram は効率を維持・向上できるか？
• RQ5 Engram は実際の内部表現と長文リトリーブにどのような影響を及ぼすか？

主な発見

• Engram は総パラメータ 27B の場合、等パラメータ・同等 FLOPs の MoE ベースラインを多様なタスクで上回る。
• 固定予算下では最適割り当ては U 字型のカーブになり、スパース容量の一部を Engram に割り当てたとき（スパース予算の約 20–25%）に最良の性能を示す。
• Engram は無限メモリ領域でべき乗律の改善をたどり、追加の計算なしで顕著な利得を可能にする。
• Engram は長文コンテキストベンチマーク（LongPPL と RULER）で顕著な利得を達成し、リトリーブ重視およびコード/数学タスクを改善する。
• 機械的分析により Engram は初期層の静的再構成を減少させ、推論の深さを実質的に深め、注意をグローバルコンテキストへ解放する。

Figure 2 : System implementation of Engram. (a) Training Phase: The massive embedding tables are sharded across available GPUs. An All-to-All communication primitive is employed to retrieve active embedding rows across devices. (b) Inference Phase: Engram tables are offloaded to host memory. By expl