🤖 machine learning

Structured Memory for Edge Language Models: Persistent Context and Corpus Retrieval via O(1) SSM State Injection

本論文は、ドキュメントコーパスを固定サイズの隠れ状態へと事前エンコードすることで $O(1)$ のコンテキスト注入を実現し、従来のTransformerベースのRAGに対して大幅な高速化を可能にする、状態空間モデル（State-Space Models）のための検索拡張生成システムであるPRECOGおよびSMCを紹介するものである。

原著者： Anusha Madan Gopal, Aras Pirbadian, Kristofor D. Carlson, M Anthony Lewis, Jonathan Tapson

公開日 2026-08-04

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Anusha Madan Gopal, Aras Pirbadian, Kristofor D. Carlson, M Anthony Lewis, Jonathan Tapson

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

あなたは、膨大な蔵書をロボットに読み込ませることで、質問に答えさせる方法を教えようとしている超スマートなロボットを想像してみてください。人工知能の世界では、最も一般的な方法は、ロボットに紙の束を渡し、「私の質問に答える前に、これらの紙のすべての単語を読みなさい」と言うようなものです。これは機能しますが、時間がかかります。もしその束が巨大であれば、ロボットは「こんにちは」と言う前に、コンテキスト（文脈）を読むだけで長い時間を費やすことになります。さらに、ロボットは今読んだすべての単語を覚えておく必要があり、それは百科事典一冊を頭の中に保持しながら会話をするようなもので、膨大なメモリを消費します。

科学者たちは、特にメモリが限られ、スピードがすべてであるスマートフォンやスマート家電のような小さなデバイスに住むロボットのために、このプロセスをより速く、より軽量にする方法を探してきました。彼らは、「状態空間モデル（SSM）」と呼ばれる異なる種類の脳のアーキテクチャを探索しています。SSMを、本をページごとに読み、すべての単語を暗記するロボットとしてではなく、本を読み、その物語全体を瞬時に一つの小さく魔法のような「要約カード」へと圧縮するロボットとして考えてみてください。このカードには、読まれた内容の「意味」は保持されていますが、単語そのものは保持されていません。また、物語が本のどこで起きたかは重要ではありません。研究者たちが問いかけている大きな疑問は、「この『要約カード』のトリックを使って、あの遅い読書プロセスを完全にスキップできるか？」ということです。もしこれが実現できれば、クラウド上の巨大なコンピュータを必要とせず、小さなデバイス上で即座に質問に答えるAIを作ることができるかもしれません。

この論文は、PRECOG（Pre-Computed Context Injection：事前計算されたコンテキスト注入）と呼ばれる巧妙なトリックを紹介しており、「イエス、可能です」と答えています。著者たちは、TENNs-LLMと呼ばれる特定のタイプのAIモデルを用いて、テキストの塊を取り込み、一度AIに実行させてその小さな「要約カード（隠れ状態）」を作成し、保存する方法を編み出しました。その後、ユーザーが質問をすると、システムはテキストを再読させる代わりに、保存された要約カードを取り出し、それを直接AIの脳内に投入します。これは、テストの読書部分をスキップして、事前に勉強しておいたノートを学生に手渡すようなものです。

結果は驚異的なスピードです。チームがテストした特定のエッジハードウェアにおいて、従来の「読む」方法では、最初の回答を得るまでに約27秒かかりました。しかし、PRECOGを使用すると、その時間は6ミリ秒未満にまで短縮されました。これは約4,500倍の高速化であり、使い物にならなかったプロセスを、即時的でインタラクティブなものへと変貌させました。この論文は、このショートカットがAIを「愚かに」することはないと数学的に証明しています。なぜなら、この特定のタイプのAIの仕組みにより、要約カードがテキストの完全な代替品となるため、回答の質はテキストを逐一読み上げた場合と同等なのです。

また、論文ではSMC（Structured Memory Consolidation：構造化メモリ統合）と呼ばれる第二の機能も紹介しています。PRECOGが「図書館の蔵書を覚えること」に関するものであるなら、SMCは「人の一生を覚えること」に関するものです。これは、AIの長期記憶（あなたの好みや過去の会話など）を、「感情」「事実」「場所」といった異なる「フォルダ」やクラスターに整理します。これにより、会話を始める際に、適切なフォルダを瞬時にAIの脳にロードすることができ、あなたが毎回説明しなくても、AIはあなたが誰で、何を好むかを記憶していられます。このシステム全体は、小型の低電力デバイスで動作するように設計されており、ポケットの中に、高度な記憶力を持つスーパースマートなAIアシスタントを持つことを可能にします。

著者たちは自分たちの数学に非常に自信を持っています。彼らは、この特定のタイプのAIにとって、「要約カード」による手法は近似や推測ではなく、テキストを読み込むことと数学的に同一であることを証明しました。しかし同時に、このトリックはあくまでこの特定のAIアーキテクチャに対してのみ有効であることも指摘しています。より一般的なAIモデル（「トランスフォーマー」と呼ばれます）は構造が異なり、単語の順序に深く関わっているため、単に要約カードと入れ替えることはできません。それらのモデルの場合、毎回テキストを読み直す必要があります。しかし、将来の小型で高速なエッジベースのAIにとって、この論文は、遅い読書という重い荷物を背負うことなく、強力な記憶力を手に入れるための道筋を提示しています。

技術要約：O(1) SSM状態注入によるエッジ言語モデルのための構造化メモリ

1. 問題提起

Retrieval-Augmented Generation (RAG) のエッジデバイスへの実用的な展開は、現在、支配的なTransformerアーキテクチャが持つ2つの構造的な制限によって制約を受けている：

線形なプリフィル・コスト： インコンテキストRAGでは、検索されたコンテキスト・トークンをプリフィルフェーズ中にトークンごとに再取り込みする必要がある。これは、コンテキスト長（ $O(L_{context})$ ）に比例するレイテンシとエネルギーコストを課しており、帯域幅が制限されたエッジハードウェアにおいては極めて困難である。
位置情報が絡み合ったKVキャッシュ： TransformerのKey-Value (KV) キャッシュは位置エンコーディング（回転式または絶対的）を組み込んでいるため、異なるコンテキストや位置間での再利用ができない。さらに、キャッシュサイズはコンテキスト長に対して線形に増大するため、コーパスの規模に応じてスケーリングが悪化するというストレージのボトルネックを生じさせる。

状態空間モデル (SSM) は、事前のコンテキストを固定サイズの、位置に依存しない再帰的隠れ状態へと圧縮することで、理論的な代替案を提供する。しかし、既存のRAG実装は依然としてトークンごとの取り込みに依存しており、コンテキストの取り込みコストを完全に排除できる可能性を秘めたSSM独自の代数的な特性を活用できていない。

2. 手法

本論文では、コンテキストの取り込みをバイパスするために、SSMの時間並進不変性を利用した2つのメカニズム、PRECOG (Pre-Computed Context Injection) と SMC (Structured Memory Consolidation) を導入する。

PRECOG: 事前計算されたコンテキスト注入

PRECOGは、検索されたドキュメントの隠れ状態をオフラインで事前計算し、クエリ時にそれを直接注入することで、RAGの概念を再定義する。

オフライン・インデックス作成： ドキュメント・コーパスをチャンクに分割する。SSM言語モデルが各チャンクを処理し、すべてのレイヤーにおける最終的な隠れ状態（ $\bar{h}^{(i)}$ ）を捕捉する。これらの固定サイズの状態（1.2Bパラメータモデルで192 KB）は、軽量な文章エンコーダー・キーと共にフラッシュストレージに保存される。
クエリ時の注入： クエリを受信すると、システムはキーを介して一致する上位 $k$ 個のチャンクを検索する。テキストを再取り込みする代わりに、システムは対応する事前計算済みの隠れ状態をロードし、モデルの初期再帰状態（ $\bar{h}_{init}$ ）として注入する。その後、モデルはクエリ・トークンのみを処理する。
理論的保証 (定理1)： 本論文では、任意の初期状態 $h_0$ 、コンテキスト $c$ 、およびクエリ $q$ に対して、状態軌跡 $S(h_0, c \oplus q)$ は $S(S(h_0, c), q)$ と代数的に同一であることを証明している。SSMの更新は現在の状態とトークンのみに依存し（絶対的な位置には依存しない）するため、事前計算された状態を注入することは、コンテキスト・トークンを逐次処理することと数学的に等価である。これにより、PRECOGが（量子化ノイズを除いて）インコンテキストRAGと全く同じ回答品質を提供することが保証される。

SMC: 構造化メモリ統合

SMCは、状態注入メカニズムを、ユーザーの嗜好、インタラクション履歴、およびログを蓄積するための永続的なオンデバイスメモリへと拡張するものである。

階層的クラスタリング： インタラクション・チャンクは、認知領域クラスター（例：感情、時間、空間、事実）へとルーティングされ、さらに想起パターン（エピソード的かセマンティックか）に基づいてサブクラスターへと分類される。
忠実度・ストレージ・ダイヤル： パラメータ $K$ $K$ が、チャンクあたりの状態保持を制御する：
- $K=N_c$ : ロスレスなエピソード記憶（すべての状態を保持）。
- $K=1$ : セマンティック記憶（最終状態のみを保持。PRECOGのチャンクと同一）。
O(1) セッション初期化： セッションの開始時に、システムは支配的なドメインを特定し、対応する統合されたセマンティック状態をSSMに直接注入する。これにより、モデルは一切のトークンを再取り込みすることなく、インタラクションの履歴を備えた状態で初期化される。これにより、蓄積された履歴の長さに依存しない $O(1)$ のレイテンシが維持される。

3. 主な貢献

RAGのアルゴリズム的リファクタリング： 本論文は、事前計算されたSSM隠れ状態を注入することにより、コンテキスト取り込みコストを $O(L_{context})$ から $O(1)$ に削減できることを示す。これは、位置情報の絡み合いによりTransformerでは構造的に不可能なメカニズムである。
厳密な等価性の証明： 定理1は、PRECOGがインコンテキストRAGと同一の状態軌跡を生成することを形式的に保証し、高速化が回答品質の低下を招かないことを保証する。
統一されたメモリ基盤： 同一の状態注入プリミティブが、静的なコーパス検索（PRECOG）と動的な永続メモリ（SMC）を統一し、単一のクエリステップで検索された知識とユーザー履歴を融合させることを可能にする。
エッジ展開の実現可能性： TENNs-LLM（1.2Bパラメータのゲート付きSSMであり、総隠れ状態は192 KB）を用いてこれらの手法を実装することで、状態レベルの検索が、オンチップメモリの制約内に隠れ状態のフットプリントが収まるニューロモーフィック・エッジハードウェア上で実用的であることを示している。

4. 結果

システムは、1.2BパラメータのTENNs-LLMバックボーンを用いて、SQuAD v1.1デベロップメントセットで評価された。

品質： PRECOG（トップ1注入）は、Exact Match (EM) 58.0、F1 73.4を達成し、インコンテキストRAGのベースライン（EM 58.2、F1 73.6）とFP16量子化ノイズの範囲内で一致した。これは、理論的な等価性が経験的に検証されたことを意味する。
レイテンシ： エッジハードウェア（シミュレーションされたニューロモーフィック・プロセッサ）において、PRECOGは、512トークンのチャンクのインコンテキスト取り込みに必要な ~27秒 のプリフィル・レイテンシを、検索と注入のオーバーヘッド自体については <6 ms へと削減した。クエリ・トークンの取り込みを含む合計のTime to First Token (TTFT) は、約 585 ms である。論文で引用されている ~4,500倍の高速化 は、特にコンテキスト取り込みフェーズの約27秒の排除を指しており、全クエリ・レイテンシを指すものではない。
ストレージ効率： 512トークンのチャンクに対し、PRECOGの状態は 192 KB を必要とするが、同等のTransformer KVキャッシュは ~16 MB を必要とする（85倍の削減）。10,000チャンクのコーパスの場合、PRECOGは1.9 GBのストレージを必要とするのに対し、仮定上のKVキャッシュ・アプローチでは160 GBが必要となり、後者は消費者向けデバイスでは実現不可能である。
アブレーション研究：
- 注入深度： 下位12レイヤー（96 KB）のみを注入した場合、フルレイヤーのF1スコアの99%を維持した。これはさらなるストレージ最適化の可能性を示唆している。
- チャンク長： チャン_ク長が約600トークン（モデルの実効メモリホライゾン）までは、インコンテキストRAGの性能と一致する。これを超えると、SSMが自然に遠方のトークンを忘却する挙動に従い、PRECOGは緩やかに性能が低下する。これは、同じバックボーンを持つインコンテキストRAGの挙動を反映している。

5. 重要性と主張

本論文は、PRECOGとSMCが、エッジ言語モデルにおける検索とメモリの扱いの根本的な転換を表すと主張している。その重要性は以下の点にある：

アーキテクチャの独占性： このメカニズムは、SSMの時間並進不変性に依存しており、TransformerのKVキャッシュにとっては「アーキテクチャ的に不可能」なものである。
エッジでの生存能力： 本手法は、現在、帯域幅が制限されたデバイスでのインタラクティブなRAGを妨げている「プリフィル・ボトルネック」を解決し、クラウド依存の推論から、オンデバイスでの低レイテンシなインタラクションへの移行を可能にする。
設計原則： 著者らは、SSMがTransformerとの品質の差を埋めるにつれ、その構造的特性（固定サイズ、位置に依存しない状態）が、アテンション型アーキテクチャでは到達できない新しいクラスの検索、キャッシュ、およびメモリ拡張アルゴリズムを可能にすると断じている。

本論文は、限界についても謙虚であり、PRECOGがSSMの有限のメモリホライゾンを継承すること（バックボーンが忘却する情報を回復できないこと）、およびトップ $k$ の構成ヒューリスティックが、トップ1注入の厳密な保証を持たない近似であることを認めている。また、ニューロモーフィック展開の数値は、物理的なテープアウトによる検証を待たずに、FPGA測定と12 nm電力シミュレーションに基づいている。