[論文レビュー] Think Before You Act: Decision Transformers with Working Memory

Jikun Kang, Romain Laroche+4 more|arXiv (Cornell University)|May 24, 2023

Explainable Artificial Intelligence (XAI)被引用数 5

ひとこと要約

本稿では、コンテンツアドレス可能な作業メモリモジュールを用いて、タスク固有の経験を明示的に格納・取得するメモリ拡張型強化学習アーキテクチャであるDecision Transformers with Memory（DT-Mem）を提案する。低ランク微調整（LoRA）を統合することで、モデルパラメータの10%のみでAtariゲームおよびMeta-Worldタスクにおいて最先端の一般化性能を達成し、サンプル効率と適応性が著しく向上した。

ABSTRACT

Decision Transformer-based decision-making agents have shown the ability to generalize across multiple tasks. However, their performance relies on massive data and computation. We argue that this inefficiency stems from the forgetting phenomenon, in which a model memorizes its behaviors in parameters throughout training. As a result, training on a new task may deteriorate the model's performance on previous tasks. In contrast to LLMs' implicit memory mechanism, the human brain utilizes distributed memory storage, which helps manage and organize multiple skills efficiently, mitigating the forgetting phenomenon. Inspired by this, we propose a working memory module to store, blend, and retrieve information for different downstream tasks. Evaluation results show that the proposed method improves training efficiency and generalization in Atari games and Meta-World object manipulation tasks. Moreover, we demonstrate that memory fine-tuning further enhances the adaptability of the proposed architecture.

研究の動機と目的

• 大規模言語モデルベースのエージェントがパrameterベースの暗黙的記憶に依存し、深刻な忘却を起こすという非効率性を解消すること。
• 人間の認知プロセスを模倣して作業メモリを明示的にモデル化することで、タスク間での訓練効率と一般化性能を向上させること。
• メモリモジュールにおける低ランク微調整（LoRA）を用いて、最小限の微調整データで新しいタスクへの効率的適応を可能にすること。
• 従来の手法と比較して計算リソースとデータ要件を削減しながら、性能を維持または向上させること。
• 意思決定変換器におけるパrameterベース記憶の代替として、より解釈可能でモジュラーな手法を提供すること。

提案手法

• 内部作業メモリモジュールを、タスク固有の経験を格納するコンテンツアドレス可能な行列として実装する。
• メモリの更新は、入力シーケンスとメモリキー間の相関を計算するアテンションベースのメカニズムを用い、注意重みに基づいて値を更新する。
• メモリの取得にはコンテンツベースのアドレス指定を用い、クエリと格納済みキーとの類似度に基づいて関連情報を特定して取得する。
• メモリモジュールの微調整には低ランク微調整（LoRA）を用い、新しいタスク用にメモリ出力を調整する小さな学習可能なパラメータを適用する。
• 一般知識を保持する事前学習済みのDecision Transformerを活用し、微調整時にはメモリモジュールのみを適応させるため、パラメータ効率を維持する。
• 二重メカニズムを採用：新規経験の格納を目的としたメモリ更新と、過去のタスク関連情報に基づいた意思決定の条件付けを目的としたメモリ取得。

実験結果

リサーチクエスチョン

• RQ1 明示的な内部作業メモリは、マルチタスク強化学習における一般化性能とサンプル効率を向上させることができるか？
• RQ2 暗黙的パrameterベース記憶と比較して、コンテンツアドレス可能なメモリの取得は、深刻な忘却と性能の観点でどのように異なるか？
• RQ3 低ランク微調整（LoRA）は、最小限のデータでメモリモジュールの新しいタスクへの微調整をどの程度効率的に可能にするか？
• RQ4 提案されたメモリ拡張アーキテクチャは、AtariおよびMeta-Worldベンチマークにおいて、MDT、HDT、PDTといった既存手法を上回る性能を示すか？
• RQ5 メモリサイズと微調整データ量の変化が、モデルの性能と適応性に与える影響は何か？

主な発見

• DT-MemはMeta-World ML45ベンチマークで最先端の性能を達成し、微調整後テストスコア0.95±0.10を記録し、HDT、PDT、MDTを上回った。
• Meta-Worldにおいて、DT-Memは微調整なしでも8%、微調整後でも3%の性能向上をHDTに対して達成した。
• Atariゲームでは、モデルパラメータの10%のみでMDT-Top3と同等または優れた一般化性能を達成した。
• Pongデータセットの20%を200kステップ微調整した際、DT-Mem Top3はMDT-Top3よりも1.154%の性能向上を示し、効果的な適応を示した。
• 微調整データ量の増加に伴い、性能低下が10%データ時0.978%から20%データ時1.154%の向上に変化し、微調整データ量の増加が一般化性能の向上に寄与することが示された。
• LoRA適応を施したメモリモジュールは147Kパラメータのみを必要とし、HDTの147Kの適応パラメータおよび230万のハイパーネットパラメータと比較して、計算オーバーヘッドを著しく削減した。