Mixture-of-Experts Encontra Aprendizado por Reforço em Contexto

Resumo

O aprendizado por reforço em contexto (ICRL, na sigla em inglês) surgiu como um paradigma promissor para adaptar agentes de RL a tarefas subsequentes por meio de condicionamento por prompts. No entanto, dois desafios notáveis permanecem para aproveitar plenamente o aprendizado em contexto em domínios de RL: a multimodalidade intrínseca dos dados de estado-ação-recompensa e a natureza diversa e heterogênea das tarefas de decisão. Para enfrentar esses desafios, propomos o T2MIR (Token- and Task-wise MoE for In-context RL), uma estrutura inovadora que introduz avanços arquitetônicos de mistura de especialistas (MoE, na sigla em inglês) em modelos de decisão baseados em transformadores. O T2MIR substitui a camada feedforward por duas camadas paralelas: um MoE token-wise que captura semânticas distintas dos tokens de entrada em múltiplas modalidades, e um MoE task-wise que direciona tarefas diversas para especialistas específicos, gerenciando uma ampla distribuição de tarefas com conflitos de gradiente reduzidos. Para aprimorar o roteamento task-wise, introduzimos um método de aprendizado contrastivo que maximiza a informação mútua entre a tarefa e sua representação de roteamento, permitindo uma captura mais precisa de informações relevantes à tarefa. As saídas dos dois componentes MoE são concatenadas e alimentadas na próxima camada. Experimentos abrangentes mostram que o T2MIR facilita significativamente a capacidade de aprendizado em contexto e supera vários tipos de baselines. Trazemos o potencial e a promessa do MoE para o ICRL, oferecendo uma melhoria arquitetônica simples e escalável para avançar o ICRL um passo mais próximo das conquistas nas comunidades de linguagem e visão. Nosso código está disponível em 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.

Mixture-of-Experts Encontra Aprendizado por Reforço em Contexto

Mixture-of-Experts Meets In-Context Reinforcement Learning

Resumo

Support