ロカス: あなたのモデルは局所的にサポートされたパラメトリックメモリの原理に基づく初期化装置である

要旨

本論文では、テストタイムトレーニングと、モデルパラメータから柔軟にオフロードまたはマージ可能な新たな種類のパラメトリックメモリを統合することを目指す。我々はLocasを提案する。これは、現代のトランスフォーマーにおけるFFNブロックの設計を共有する局所サポート型パラメトリックメモリであり、効率的な継続学習をサポートしつつ、モデルパラメータへ柔軟に恒久化（permanentized）できる。Locasの2つの主要なバリアントについて論じる：一つは、より明確な理論的保証を持つ従来の2層MLP設計のもの。もう一つは、SOTA大規模言語モデル（LLM）と同じGLU-FFN構造を共有し、既存モデルに容易に付加でき、パラメータ効率と計算効率の両方に優れた継続学習を実現する。決定的に重要なのは、モデルパラメータ、活性化、および/または勾配を再利用するという原理に基づいた方法で行う、このような低ランクの横方向FFN型メモリの適切な初期化が、高速な収束、改善された汎化、および破滅的忘却の防止に不可欠であることを示すことである。提案するメモリ機構を、PG-19全書籍言語モデリングタスクおよびLoCoMo長文対話質問応答タスクで検証する。最小の場合でわずか0.02%の追加パラメータのみで、Locas-GLUははるかに小さなコンテキストウィンドウを維持しつつ、過去のコンテキストからの情報を格納可能である。さらに、比較的MMLU評価を通じて、Locasを用いて書籍全体を記憶した後のモデルの一般的能力の損失も検証する。結果は、Locasがモデルの既存の内部知識の破滅的忘却を最小化しつつ、過去のコンテキストをパラメトリック知識として恒久化する有望な能力を示している。

English

In this paper, we aim to bridge test-time-training with a new type of parametric memory that can be flexibly offloaded from or merged into model parameters. We present Locas, a Locally-Supported parametric memory that shares the design of FFN blocks in modern transformers, allowing it to be flexibly permanentized into the model parameters while supporting efficient continual learning. We discuss two major variants of Locas: one with a conventional two-layer MLP design that has a clearer theoretical guarantee; the other one shares the same GLU-FFN structure with SOTA LLMs, and can be easily attached to existing models for both parameter-efficient and computation-efficient continual learning. Crucially, we show that proper initialization of such low-rank sideway-FFN-style memories -- performed in a principled way by reusing model parameters, activations and/or gradients -- is essential for fast convergence, improved generalization, and catastrophic forgetting prevention. We validate the proposed memory mechanism on the PG-19 whole-book language modeling and LoCoMo long-context dialogue question answering tasks. With only 0.02\% additional parameters in the lowest case, Locas-GLU is capable of storing the information from past context while maintaining a much smaller context window. In addition, we also test the model's general capability loss after memorizing the whole book with Locas, through comparative MMLU evaluation. Results show the promising ability of Locas to permanentize past context into parametric knowledge with minimized catastrophic forgetting of the model's existing internal knowledge.