Don't flatten, tokenize! Unlocking the key to SoftMoE's efficacy in deep RL

TL;DR

This paper reveals that in deep reinforcement learning, tokenizing encoder outputs, rather than using multiple experts, is the main factor behind SoftMoE's performance improvements, even with a single expert.

Contribution

The study uncovers that tokenization, not the mixture of experts, drives SoftMoE's success in deep RL, challenging previous assumptions.

Findings

Tokenizing encoder outputs is key to SoftMoE's effectiveness.

Single expert models with tokenization can match SoftMoE performance.

Performance gains are largely due to tokenization rather than multiple experts.

Abstract

The use of deep neural networks in reinforcement learning (RL) often suffers from performance degradation as model size increases. While soft mixtures of experts (SoftMoEs) have recently shown promise in mitigating this issue for online RL, the reasons behind their effectiveness remain largely unknown. In this work we provide an in-depth analysis identifying the key factors driving this performance gain. We discover the surprising result that tokenizing the encoder output, rather than the use of multiple experts, is what is behind the efficacy of SoftMoEs. Indeed, we demonstrate that even with an appropriately scaled single expert, we are able to maintain the performance gains, largely thanks to tokenization.