階層的メモリを用いた事前学習：ロングテール知識と共通知識の分離

Pretraining with hierarchical memories: separating long-tail and common knowledge

September 29, 2025

著者: Hadi Pouransari, David Grangier, C Thomas, Michael Kirchhof, Oncel Tuzel

cs.AI

要旨

現代の言語モデルの驚異的な性能向上は、現在のところパラメータのスケーリングに依存している：より大きなモデルはより多くの世界知識を保持し、より優れた推論能力を発揮する。しかし、すべての世界知識をパラメータに圧縮することは不必要であり、各プロンプトで使用されるのはその一部に過ぎないため、推論時のメモリと計算リソースが限られたエッジデバイスにとっては非現実的である。我々はこの欠点を、メモリ拡張アーキテクチャと既存のハードウェアパラダイムに沿った事前学習戦略によって解決する。我々は、世界知識をエンコードした大規模な階層的パラメトリックメモリバンクにアクセスする小型言語モデルを提案する。事前学習および推論中に、コンテキストに依存した小さなメモリブロックをフェッチし、モデルに追加する。我々の事前学習は、長尾の世界知識をメモリパラメータに保存することを学習し、一方で小型言語モデルは共通知識と一般的な推論能力を捉えるアンカーとして機能する。兆トンスケールの実験を通じて、我々は顕著な性能向上を示す：4.6Bのメモリバンクからフェッチされた18Mパラメータのメモリを拡張した160Mパラメータのモデルは、パラメータ数が2倍以上の通常のモデルと同等の性能を達成する。広範な実験を通じて、我々はトランスフォーマーにおける最適なパラメトリックメモリのタイプとサイズを研究し、それらを21B以上のパラメータにスケーリングする。我々は、提案する階層的フィードフォワードメモリが、事前学習中または事後的に追加されるかどうかにかかわらず、トランスフォーマーアーキテクチャ全体で堅牢に機能することを発見した。

English

The impressive performance gains of modern language models currently rely on scaling parameters: larger models store more world knowledge and reason better. Yet compressing all world knowledge into parameters is unnecessary, as only a fraction is used per prompt, and impractical for edge devices with limited inference-time memory and compute. We address this shortcoming by a memory-augmented architecture and a pretraining strategy aligned with existing hardware paradigms. We introduce small language models that access large hierarchical parametric memory banks encoding world knowledge. During pretraining and inference, we fetch a small, context-dependent memory block and add it to the model. Our pretraining learns to store long-tail world knowledge in the memory parameters, while the small language model acts as an anchor capturing common knowledge and general reasoning abilities. Through trillion-token-scale experiments, we show significant gains: a 160M-parameters model augmented with an 18M-parameters memory fetched from a 4.6B memory bank obtains comparable performance to a regular model with more than 2x the parameters. Through extensive experiments, we study the optimal type and size of parametric memories in transformers, scaling them to over 21B parameters. We find that our proposed hierarchical feed-forward memories work robustly across transformer architectures, whether added during pretraining or post-hoc.