[論文レビュー] Predictive auxiliary objectives in deep RL mimic learning in the brain

Ching Fang, Kimberly Stachenfeld|arXiv (Cornell University)|Oct 9, 2023

Neural dynamics and brain function被引用数 5

ひとこと要約

この論文は、深層強化学習（RL）における予測補助目的が、脳で観察されるニューロンの表現ダイナミクスを模倣することを示している。予測ヘッドは海馬の機能に類似しており、エンコーダーは視覚皮質に類似している。このような目的は、リソース制限のあるネットワークにおける学習を安定化させ、より長い時間窓を用いた転移学習を改善し、生物学的データと整合する表現変化を引き起こす。

ABSTRACT

The ability to predict upcoming events has been hypothesized to comprise a key aspect of natural and machine cognition. This is supported by trends in deep reinforcement learning (RL), where self-supervised auxiliary objectives such as prediction are widely used to support representation learning and improve task performance. Here, we study the effects predictive auxiliary objectives have on representation learning across different modules of an RL system and how these mimic representational changes observed in the brain. We find that predictive objectives improve and stabilize learning particularly in resource-limited architectures, and we identify settings where longer predictive horizons better support representational transfer. Furthermore, we find that representational changes in this RL system bear a striking resemblance to changes in neural activity observed in the brain across various experiments. Specifically, we draw a connection between the auxiliary predictive model of the RL system and hippocampus, an area thought to learn a predictive model to support memory-guided behavior. We also connect the encoder network and the value learning network of the RL system to visual cortex and striatum in the brain, respectively. This work demonstrates how representation learning in deep RL systems can provide an interpretable framework for modeling multi-region interactions in the brain. The deep RL perspective taken here also suggests an additional role of the hippocampus in the brain -- that of an auxiliary learning system that benefits representation learning in other regions.

研究の動機と目的

• 予測補助目的が深層RLにおける表現をどのように形作るか、およびその変化が脳の神経ダイナミクスにどれほど類似しているかを調査すること。
• リソース制限のある深層RLアーキテクチャにおける予測目的の機能的利点を評価すること。
• 予測時間窓の長さが、タスク間での表現転移を可能にするメカニズムを探索すること。
• 表現変化に基づいて、深層RLモジュールと脳領域（例：海馬、視覚皮質、striatum）との類似性を考察すること。
• 海馬が計画とは独立して、他の脳領域における表現学習を強化する補助的学習システムとして機能する可能性を提唱すること。

提案手法

• 主なRL目的（Q学習）と、コントラスト予測を用いた予測補助目的を併用したマルチ目的設定で深層RLエージェントを訓練した。
• 現在の観測から将来の状態を予測する予測モデルを訓練するため、正例と負例を用いたコントラスト損失を用いた。
• 状態表現を抽出するエンコーダー・ネットワークと、Q値を推定する価値ネットワークを、両方とも補助目的に影響を受けるように訓練した。
• 予測と価値学習の表現に与える影響を分離するために、グリッドワールド採掘環境で制御された実験を実施した。
• 既知の実験的パラダイム（例：刺激の入れ替え、報酬調節によるチューニング）をシミュレートすることで、RLシステム内の表現変化を神経データと照合した。
• ニューラル活動パターンを用いてエンコーダーと予測ヘッドの表現シフトを定量的に評価し、視覚皮質および海馬からのin vivo記録と比較した。

Figure 1: A deep RL framework to model multi-region computation. A. In the deep RL model we use, reward is provided as a scalar input $r$ . Observations $o$ are 2D visual inputs fed into an encoder (green) that learns low-dimensional state space representations $z$ . The encoder is a convolutional n

実験結果

リサーチクエスチョン

• RQ1 予測補助目的は、特にリソース制限のあるモデルにおいて、深層RLにおける表現学習にどのように影響を与えるか？
• RQ2 予測時間窓の長さが、表現転移および学習安定性に与える影響は何か？
• RQ3 深層RLシステムにおける表現変化は、視覚皮質および海馬で観察されるものとどの程度類似しているか？
• RQ4 RLエージェントの予測ヘッドは、脳における海馬の機能的アナログと解釈できるか？
• RQ5 あるモジュールで予測表現を学習することで、エンコーダーや価値ネットワークのような他の相互接続されたモジュールに有益な表現変化が誘発されるか？

主な発見

• 予測補助目的は、特にリソース制限のあるアーキテクチャにおいて、表現の崩壊を顕著に低減する。
• 長い予測時間窓は、タスク間での表現転移をより良くし、短時間窓の予測を上回る性能を示す。
• スティムラスの順序を入れ替えた刺激系列にさらされた後、エンコーダー内のニューロンがチューニングを変化させ、マカクのIT皮質での結果（Li & DiCarlo, 2008）と一致する。
• スイッチ位置に近づくほど、エンコーダーのユニットの発火率が非好ましい刺激にシフトし、その効果が強まるが、予測目的が存在しないとその効果は消える。
• 価値学習はエンコーダー表現に影響を与える：初期のエンコーダー層のユニットは、報酬が与えられた刺激に対して選択的になる傾向が強く、マウスのV1での報告（Poort et al., 2015）と整合する。
• 予測ヘッドの活動は、特に予測的および系列符号化特性において海馬のダイナミクスに類似しており、海馬が表現形成のための補助的学習システムとして機能するとする仮説を支持する。

Figure 2: Gridworld performance with predictive auxiliary tasks. A. The model is tested on gridworld task in a 8x8 arena. The agent must navigate to a hidden reward given random initial starting locations. B. Average episode score across training steps for models without auxiliary losses (blue), wit