エキスパートの混合と文脈内強化学習の融合

Mixture-of-Experts Meets In-Context Reinforcement Learning

June 5, 2025

著者: Wenhao Wu, Fuhong Liu, Haoru Li, Zican Hu, Daoyi Dong, Chunlin Chen, Zhi Wang

cs.AI

要旨

文脈内強化学習（In-context Reinforcement Learning, ICRL）は、プロンプト条件付けを通じてRLエージェントを下流タスクに適応させる有望なパラダイムとして登場しました。しかし、RL領域における文脈内学習を十分に活用する上で、2つの重要な課題が残されています：状態-行動-報酬データの本質的なマルチモーダル性と、意思決定タスクの多様で異質な性質です。これらの課題に取り組むため、我々はT2MIR（Token- and Task-wise MoE for In-context RL）を提案します。これは、トランスフォーマーベースの意思決定モデルに専門家混合（Mixture-of-Experts, MoE）のアーキテクチャ的進化を導入する革新的なフレームワークです。T2MIRはフィードフォワード層を2つの並列層に置き換えます：入力トークンの複数のモダリティにわたる異なるセマンティクスを捉えるトークンワイズMoEと、多様なタスクを専門家にルーティングして広範なタスク分布を管理し、勾配の衝突を緩和するタスクワイズMoEです。タスクワイズルーティングを強化するため、タスクとそのルータ表現間の相互情報量を最大化するコントラスティブラーニング手法を導入し、タスク関連情報をより正確に捉えることを可能にします。2つのMoEコンポーネントの出力は連結され、次の層に供給されます。包括的な実験により、T2MIRが文脈内学習能力を大幅に促進し、さまざまなタイプのベースラインを上回ることが示されました。我々はMoEの可能性と約束をICRLにもたらし、言語および視覚コミュニティでの成果に一歩近づけるためのシンプルでスケーラブルなアーキテクチャ的強化を提供します。コードはhttps://github.com/NJU-RL/T2MIRで公開されています。

English

In-context reinforcement learning (ICRL) has emerged as a promising paradigm for adapting RL agents to downstream tasks through prompt conditioning. However, two notable challenges remain in fully harnessing in-context learning within RL domains: the intrinsic multi-modality of the state-action-reward data and the diverse, heterogeneous nature of decision tasks. To tackle these challenges, we propose T2MIR (Token- and Task-wise MoE for In-context RL), an innovative framework that introduces architectural advances of mixture-of-experts (MoE) into transformer-based decision models. T2MIR substitutes the feedforward layer with two parallel layers: a token-wise MoE that captures distinct semantics of input tokens across multiple modalities, and a task-wise MoE that routes diverse tasks to specialized experts for managing a broad task distribution with alleviated gradient conflicts. To enhance task-wise routing, we introduce a contrastive learning method that maximizes the mutual information between the task and its router representation, enabling more precise capture of task-relevant information. The outputs of two MoE components are concatenated and fed into the next layer. Comprehensive experiments show that T2MIR significantly facilitates in-context learning capacity and outperforms various types of baselines. We bring the potential and promise of MoE to ICRL, offering a simple and scalable architectural enhancement to advance ICRL one step closer toward achievements in language and vision communities. Our code is available at https://github.com/NJU-RL/T2MIR.