Meta社は最近、画期的な研究成果を発表しました。大型言語モデル（LLM）の事実の正確性を大幅に向上させ、パラメータ規模において前例のない拡張を実現する、新型のメモリ層技術を開発したのです。この技術は、従来のニューラルネットワークの拡張方法に挑戦するだけでなく、将来のAIアーキテクチャ設計に新たな方向性を提供します。
この研究の中心は、学習可能なキーバリュー検索メカニズムを利用して、計算量（FLOPs）を増やすことなくモデルに追加パラメータを与えることにあります。この手法の核心は、計算集約的なフィードフォワード層を補完する、スパース活性化のメモリ層によって、専門的な情報の保存と検索能力を提供することです。
従来の高密度ネットワークと比較して、メモリ層は情報の保存処理においてより効率的です。例えば、言語モデルは人名と誕生日、国の首都などの簡単な関連情報を学習する必要がありますが、メモリ層は単純なキーバリュー検索メカニズムで実現でき、フィードフォワードネットワークを使用するよりも効率的です。
この研究の主要な貢献は、メモリ層の規模を前例のない1280億パラメータに拡張したことでしょう。実験結果によると、改良型メモリ層を搭載した言語モデルは、下流タスクにおいて、計算量が2倍の稠密モデルを上回り、計算量とパラメータ数が一致した場合でも、混合専門家モデルを凌駕します。特に事実に関するタスクにおいて、性能向上が顕著です。
Metaの研究者たちは、Transformerネットワークの1つ以上のフィードフォワードネットワーク（FFN）をメモリ層に置き換えることで、この目標を達成しました。この置き換え方法は、様々な基礎モデルサイズ（1.34億〜80億パラメータ）とメモリ容量（最大1280億パラメータ）において一貫して優位性を示しました。実験結果によると、メモリ層は言語モデルの事実の正確性を100％以上向上させ、同時にコード記述と一般的な知識についても顕著な向上を示しました。多くの場合、メモリ層を搭載したモデルは、計算量が4倍必要な高密度モデルと同等の性能に達しました。
研究者たちは、大規模な応用における課題を克服するために、メモリ層にいくつかの改良を加えました。
積キー検索メカニズムの採用：大規模メモリ層におけるクエリキー検索のボトルネックを解決するために、学習可能な積量子化キーを採用し、各クエリキーペアの比較を回避しました。
メモリ層の並列化：マルチGPU環境でのメモリ層の並列化を実現するために、埋め込み検索と集約操作を複数のGPUに分散しました。
共有メモリメカニズム：パラメータの共有を最大限にするために、すべてのメモリ層間で共有メモリパラメータプールを使用しました。
性能と安定性の最適化：カスタムCUDAカーネルを使用してEmbeddingBag操作を最適化し、メモリ帯域幅の利用率を大幅に向上させました。さらに、silu非線形性を持つ入力関連ゲート制御メカニズムを導入し、トレーニングの性能と安定性を向上させました。
実験結果からは、以下の重要な発見が明らかになりました。
メモリ層のサイズは性能に著しい影響を与える：メモリ層のサイズが大きくなるにつれて、事実に関する質問への回答の性能が継続的に向上します。
複数のメモリ層は単一のメモリ層よりも優れている：共有パラメータを持つ複数のメモリ層を使用すると性能が向上しますが、メモリ層が多すぎると性能が低下します。最適なメモリ層の数は3つです。
メモリ層は事実をより速く学習する：トレーニング初期において、メモリ層を搭載したモデルの性能向上はより速く、メモリ層がモデルによる事実の学習を加速させることを示唆しています。
メモリ層と高密度層は補完関係にある：スパースなメモリ層と高密度なフィードフォワード層はどちらも不可欠であることが実験で示されました。
メモリ層技術の有効性を検証するために、研究者たちは、以下の複数のベンチマークテストで評価を行いました。
事実に関する質問への回答（NaturalQuestions、TriviaQA）
多段質問への回答（HotpotQA）
科学と常識的な知識（MMLU、HellaSwag、OBQA、PIQA）
コード記述（HumanEval、MBPP）
その結果、メモリ層を搭載したモデルはこれらのテストですべてのベースラインモデルを上回り、特に事実に関する質問への回答において性能向上が最も顕著でした。
Metaのこの研究は、AIモデルの拡張に新たなアイデアを提供するだけでなく、事実に関する問題を解決し、モデルの性能を向上させるための新たな道を切り開きました。研究者たちは、メモリ層技術は高い拡張性を持ち、将来様々なAIアプリケーションで広く利用される可能性があると考えています。また、メモリ層はハードウェアアクセラレーションにおいてまだ課題に直面していますが、継続的な研究と最適化を通じて、従来のフィードフォワードネットワークと同等、またはそれを超える性能を実現できると信じています。
さらに、Metaの研究チームは、新たな学習方法によってメモリ層の性能をさらに向上させ、モデルの忘却や幻覚を減らし、継続学習を実現することを目指しています。
この研究の発表は、間違いなくAI分野に新たな活力を吹き込み、将来のAIの発展に期待を抱かせます。
論文：https://arxiv.org/pdf/2412.09764
