Research map / Reinforcement learning / Offline RL
Don't flatten, tokenize! Unlocking the key to SoftMoE's efficacy in deep RL
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Analyzes why SoftMoEs help value-based deep RL on Atari and finds that tokenizing the encoder output (not multiple experts) drives the gains, so a single scaled expert matches multi-expert SoftMoE.
Problem
Performance degradation of deep RL networks as model size grows; unclear why SoftMoEs improve scaling in online value-based RL, and why they fail to help actor-critic methods.
Approach
Ablates each SoftMoE component (combining tokens, expert specialization, width, depth, number of experts, tokenization) and shows that tokenizing the convolutional encoder output (PerConv, keeping spatial structure) instead of flattening is the key driver, so a single scaled expert with tokenization matches multi-expert SoftMoE.
Key results
- A single scaled expert SoftMoE-1 with tokenization achieves performance comparable to multi-expert SoftMoE-4 on Rainbow over 20 Atari games.
- Tokenized Rainbow-lite with sum/mean pooling yields statistically significant gains over the flattened baseline, showing tokenization drives scaling.
- Shuffled tokenization performs worst while PerConv performs best, indicating spatial structure preservation matters; experts show little specialization and pruning two of four experts does not hurt performance.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Tokenize the encoder output instead of flattening it (ablation) | Rainbow-lite baseline · ALE 20 games | IQM | Helped |
| Use a single scaled expert with tokenization instead of flattening the baseline representation (ablation) | Rainbow · ALE 20 games | IQM | Helped |
| Use combined tokenization rather than sparse expert-choice tokenization (ablation) | SoftMoE-1 with scaled expert dimensionality · ALE 20 games | IQM | Helped |
| Use PerConv tokenization rather than PerFeat tokenization (ablation) | SoftMoE-1 · ALE 20 games | IQM | Helped |
| Preserve spatial structure with PerConv tokenization rather than shuffle the encoder output (ablation) | SoftMoE-1 · ALE 20 games | IQM | Helped |
| Use PerPatch tokenization instead of flattening (ablation) | SoftMoE-1 · ALE 20 games | IQM | Helped |
| Increase the number of slots to the number of tokens (ablation) | SoftMoE and ExpertChoice · ALE 20 games | IQM | No clear effect |
| Add a matching layer after the penultimate layer (ablation) | Rainbow baseline · ALE 20 games | IQM | No clear effect |
| Scale up expert width (ablation) | SoftMoE-4 · ALE 20 games | IQM | No clear effect |
| Use a single scaled expert rather than multiple experts (ablation) | SoftMoE · ALE 20 games | IQM | No clear effect |
| Reduce SoftMoE-1 to 10% of the default number of slots (ablation) | SoftMoE-1 · ALE 20 games | IQM | No clear effect |
| Replace flattening with global average pooling (ablation) | Rainbow baseline · ALE 20 games | IQM | Helped |
Surprising findings
- Single-expert SoftMoE with scaled width matches multi-expert SoftMoE, so multiple experts are not the main source of gains.
- Tokenization gains do not carry over to DQN, where SoftMoE-1 gives little improvement over the baseline.
Limitations
- Results mainly on Rainbow and Atari; DQN shows little gain from tokenization.
- Actor-critic experiments (SAC) are only initial and on continuous control CALE.
- Resets and shrink-and-perturb helped DER but degraded Rainbow, so expert utilization remedies are not general.
Experimental setup
Datasets: Arcade Learning Environment (20 games) Arcade Learning Environment (60 games) Atari100k (26 games) Procgen CALE
Benchmarks: Arcade Learning Environment (20 games) Arcade Learning Environment (60 games) Procgen CALE
Metrics: IQM human-normalized score median mean optimality gap
Baselines: Rainbow SoftMoE-4 ExpertChoice DQN DER
Builds on: SoftMoE Rainbow IMPALA ResNet DQN Dopamine
Closest papers by meaning
- Mixtures of Experts Unlock Parameter Scaling for Deep RL (ICML 2024 · similarity 0.87)
- Reinforcement Learning Finetunes Small Subnetworks in Large Language Models (NeurIPS 2025 · similarity 0.86)
- Do Not Let Low-Probability Tokens Over-Dominate in RL for LLMs (ICLR 2026 · similarity 0.86)
- Stop Regressing: Training Value Functions via Classification for Scalable Deep RL (ICML 2024 · similarity 0.85)
- Compute-Optimal Scaling for Value-Based Deep RL (NeurIPS 2025 · similarity 0.85)
- Network Sparsity Unlocks the Scaling Potential of Deep Reinforcement Learning (ICML 2025 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.