Mixture-of-Experts ontmoet In-Context Reinforcement Learning

Mixture-of-Experts Meets In-Context Reinforcement Learning

June 5, 2025

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

cs.AI

Samenvatting

In-context reinforcement learning (ICRL) is naar voren gekomen als een veelbelovend paradigma voor het aanpassen van RL-agents aan downstream taken door middel van promptconditionering. Er blijven echter twee opmerkelijke uitdagingen bestaan bij het volledig benutten van in-context learning binnen RL-domeinen: de intrinsieke multi-modaliteit van de staat-actie-beloning data en de diverse, heterogene aard van beslissingstaken. Om deze uitdagingen aan te pakken, stellen we T2MIR (Token- en Task-wise MoE voor In-context RL) voor, een innovatief framework dat architectonische vooruitgang van mixture-of-experts (MoE) introduceert in transformer-gebaseerde beslissingsmodellen. T2MIR vervangt de feedforward-laag door twee parallelle lagen: een token-wise MoE die verschillende semantiek van invoertokens over meerdere modaliteiten vastlegt, en een task-wise MoE die diverse taken doorverwijst naar gespecialiseerde experts voor het beheren van een brede taakverdeling met verminderde gradientconflicten. Om de task-wise routing te verbeteren, introduceren we een contrastieve leermethode die de wederzijdse informatie tussen de taak en zijn routerrepresentatie maximaliseert, waardoor een nauwkeurigere vastlegging van taakgerelateerde informatie mogelijk wordt. De uitvoer van de twee MoE-componenten wordt samengevoegd en doorgegeven aan de volgende laag. Uitgebreide experimenten tonen aan dat T2MIR de in-context leercapaciteit aanzienlijk vergemakkelijkt en verschillende soorten baselines overtreft. We brengen het potentieel en de belofte van MoE naar ICRL, en bieden een eenvoudige en schaalbare architectonische verbetering om ICRL een stap dichter bij de prestaties in de taal- en visiegemeenschappen te brengen. Onze code is beschikbaar op 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.