Source-linked AI summary
Mixtures of Experts Unlock Parameter Scaling for Deep RL
Johan Obando-Ceron, Ghada Sokar, Timon Willi, Clare Lyle, Jesse Farebrother, Jakob Foerster, Gintare Karolina Dziugaite, Doina Precup, Pablo Samuel Castro
TL;DR
Parameter scaling remains challenging in deep RL, where simply enlarging networks does not reliably improve performance. This paper tests Mixture-of-Experts modules, especially Soft MoE, in value-based networks and reports benefits across diverse training regimes, with Soft MoE gains increasing with expert count.
Problem
Deep RL network scaling is challenging, and parameters may be underused, limiting evidence that increasing parameter count improves performance.
Source: §1 · p. 1Method
The paper evaluates Mixture-of-Experts modules, particularly Soft MoE, in value-based deep RL networks as parameter counts increase.
Source: §1 · p. 2Results
Takeaways & Limitations
The findings support the importance of architectural design choices for the performance of RL agents.
Source: §7 · p. 9The study focuses on penultimate-layer MoEs in off-policy, single-task value-based agents and leaves many design choices and results beyond ALE open for further study.
Source: §7 · pp. 8–9
Abstract
from arXiv · showhide
The recent rapid progress in (self) supervised learning models is in large part predicted by empirical scaling laws: a model's performance scales proportionally to its size. Analogous scaling laws remain elusive for reinforcement learning domains, however, where increasing the parameter count of a model often hurts its final performance. In this paper, we demonstrate that incorporating Mixture-of-Expert (MoE) modules, and in particular Soft MoEs (Puigcerver et al., 2023), into value-based networks results in more parameter-scalable models, evidenced by substantial performance increases across a variety of training regimes and model sizes. This work thus provides strong empirical evidence towards developing scaling laws for reinforcement learning.
1. Introduction
The introduction frames parameter scaling in deep RL as difficult, with network parameters often underused, and asks whether Mixture-of-Experts modules can improve value-based agents as they grow. It presents Soft MoE as a promising approach and previews analyses and evaluations across additional training regimes.
1. Introduction
- Soft MoE improves DQN and Rainbow performance with scale, whereas Top1-MoE helps only Rainbow and does not improve with scale.Source: Figure 1 · p. 1
- Deep RL scaling is challenging, and underused parameters may help explain why simply increasing model size does not reliably improve performance.Source: §1 · p. 1
- MoEs have supported scaling in supervised learning; their modularity suits distributed computation, and structured sparsity has been associated with improved performance (Fedus et al., 2022; Evci et al., 2020; Gale et al., 2019).Source: §1 · pp. 1–2
- The paper investigates whether MoEs make value-based deep RL networks more parameter-scalable, reporting that Soft MoE improves performance and gains increase with expert count.It also examines gating, input tokenization, and expert representations to investigate the observed gains.Source: §1 · p. 2
- Additional evaluations cover Offline RL, low-data training, and alternative architectural designs as directions for further study.Source: §1 · p. 2
2. Preliminaries
This section reviews value-based reinforcement learning and Mixture-of-Experts methods. It introduces Soft MoE as a differentiable alternative to hard token routing.
2.1. Reinforcement Learning
- Deep RL uses neural networks to approximate action-value functions, with replay buffers supplying transitions for gradient updates.The replay ratio is defined as gradient updates relative to environment interactions.Source: §2.1 · pp. 2–3
- Rainbow extends DQN with components intended to improve learning stability, sample efficiency, and overall performance.The passage identifies Rainbow as an important deep-RL baseline.Source: §2.1 · p. 3
2.2. Mixtures of Experts
- MoEs comprise expert subnetworks activated by a gating network that routes incoming tokens to experts.Soft MoE's reported accuracy–computational-cost trade-off is favorable relative to other MoE methods.Source: §2.2 · p. 3
- Soft MoE replaces hard token-to-expert assignments with fully differentiable soft assignments, addressing challenges including training instability and dropped tokens.Its input tokens are mixed for experts, whose outputs are then combined using learned weights.Source: §2.2 · p. 3
3. Mixture of Experts for Deep RL
The section describes how MoE modules are incorporated into DQN-based architectures, including choices about placement, tokenization, and gating. It compares Soft MoE and Top1-MoE and reports tokenization-dependent results on Rainbow.
3. Mixture of Experts for Deep RL
- The networks replace the penultimate layer with an MoE module whose experts match the original dense layer's dimensionality, widening that layer by the expert count.The paper also discusses alternative placements in Section 5.3.Source: §3 · p. 3
- On Rainbow with 8 experts, Soft MoE performs best with PerConv, Top1-MoE with PerFeat, and PerSamp is worst for both.Figure 6 compares the tokenization methods across the two MoE variants.Source: Figure 6 · p. 4
- PerConv tokenization splits the convolutional encoder output into h × w tokens of dimensionality d, with a trainable projection preserving token size after each expert.The paper describes PerConv as the best-performing and most intuitive approach among those explored.Source: §3 · p. 4
- The paper compares top-k gating with k = 1, called Top1-MoE, against Soft MoE, which can activate multiple experts per token.The authors focus on these variants because Top1-MoE is prevalent while Soft MoE is simpler and shows evidence of improving performance.Source: §3 · p. 4
4. Empirical evaluation
Across DQN and Rainbow experiments, Soft MoE scales better as experts increase, while widening baseline networks can hurt; design analyses test alternative choices and training dynamics.
4. Empirical evaluation
- The evaluation focuses on DQN and Rainbow, mostly using the Impala encoder on 20 Atari games, and compares MoEs with widening the penultimate layer.The widening control increases that layer by the same factor as the number of experts.Source: §4 · p. 4, §4.1 · p. 4
4.1. Soft MoE Helps Parameter Scalability
- Rainbow’s 1-to-8 expert increase yields a 20% performance improvement, while widening its base layer 1-to-8 lowers performance around 40%; Soft MoE’s advantage persists at high replay ratios.The experiments varied expert count and compared MoEs with widening the base penultimate layer.Source: §4.1 · pp. 4–5
4.2. Impact of Design Choices
- Soft MoE benefits from adding experts, whereas Top1-MoE does not show the same expert-count scaling.The paper varies expert count to assess parameter scalability.Source: §4.2 · p. 5
- Reducing each Soft MoE expert’s dimensionality preserves performance, while single-expert gains and random-Φ results support a role for learned gating and combining.The authors suggest structured sparsity, rather than expert size alone, accounts for much of the benefit.Source: §4.2 · p. 5
- PerConv works best for Soft MoE while PerFeat benefits Top1-MoE; Soft MoE also helps with the standard CNN encoder and across all 60 Atari games.The 60-game study uses Rainbow with the Impala architecture and eight experts.Source: §4.2 · p. 5, Figure 8 · p. 5
- Adding load-balancing losses does not rescue Top1-MoE, suggesting RL agents benefit from weighted token combinations rather than hard routing.The losses were tested on Top1-MoE with eight experts using DQN and Rainbow.Source: Figure 9 · p. 5, §4.2 · p. 6
4.3. Additional Analysis
- MoE networks show higher empirical NTK rank and negligible dormant neurons and feature norms than the baseline, but these metrics do not establish a direct performance cause.The analysis examines Rainbow with the Impala architecture and eight experts; Figure 10 also reports greater robustness to dormant neurons and higher effective rank of features and gradients, with less feature-norm growth.Source: §4.3 · p. 6, Figure 10 · p. 6
5. Future directions
The section tests Soft MoE and other expert designs across offline RL, low-data training, and architectural variants, reporting both gains and boundaries on when they appear.
5. Future directions
- The section presents cross-regime empirical results as indicators for promising future research on MoEs in deep RL.The investigation focuses on DQN and Rainbow with Soft MoE and Top1-MoE on the standard ALE benchmark.Source: §5 · pp. 6–7
5.1. Offline RL
- Soft MoE achieved the best aggregate final performance with CQL and CQL+C51 on offline datasets; Top1-MoE improved CQL+C51 but not CQL.The evaluation used 17 games, five seeds, and datasets containing 5%, 10%, or 50% of DQN-collected samples.Source: §5.1 · p. 7
5.2. Agent Variants For Low-Data Regimes
- At 100k interactions, agent variants showed no real difference, while 50M-step training produced gains for DrQ(ϵ) and DER, especially DER.The observed gains were consistent with the results for their base agents, DQN and Rainbow.Source: §5.2 · p. 7
5.3. Expert Variants
- Among expert variants, regular Soft MoE performed best overall; normalization seemed to help Big experts, while normalized All variants were unable to learn.Normalization particularly improved Big experts with DQN, surpassing the Regular setup; the authors suggest shared versus non-shared final layers may explain the difference.Source: §5.3 · p. 7
6. Related Work
Related work describes MoEs as a tool for scaling supervised models, while noting that RL parameter use and MoE applications in RL remain open areas of investigation.
MoEs in supervised learning and RL
- MoEs have supported scaling in language and vision models, while prior RL work has explored MoEs in single- and multi-task settings using definitions and uses that differ from this paper’s.Related work also notes that MoEs have been used in vision for continual learning and multi-task problems.Source: §6 · pp. 7–8
Reported agent comparisons
- Soft MoE generally remained stable with more training and attained higher final performance in the reported CQL, CQL+C51, DrQ(ϵ), and DER comparisons.The figures report normalized performance across 17 games for the offline agents and 26 games for DrQ(ϵ) and DER.Source: Figure 11 · p. 8, Figure 12 · p. 8
Parameter Scalability and Efficiency in Deep RL
- Prior studies report under-used RL parameters, including inactive neurons, substantial pruning without performance loss, and high sparsity without performance loss.The cited work also reports performance improvements from periodic weight resetting.Source: §6 · p. 8
7. Discussion and Conclusion
The discussion reports benefits of MoEs for value-based agents across training regimes and examines how sparsity and additional parameters relate to those gains. It also identifies boundaries and open questions for extending the evidence.
Discussion and Conclusion
- MoEs had a beneficial effect on value-based agents across diverse training regimes, while naive increases in network parameters did not improve performance.The paper frames alternative network topologies as a route toward robustness to scaling.Source: §7 · p. 8
- The results suggest MoE gains reflect both structured sparsity and MoE modules: a single-expert Rainbow MoE improved performance, and Soft MoE’s extra parameters were necessary.The authors cite reduced expert dimensionality without performance loss as additional evidence in this analysis.Source: §7 · p. 8
- The evidence is limited to penultimate-layer MoEs for off-policy, single-task learning, leaving choices such as Top1-MoE k, tokenization, and router optimization open.The authors also note that evaluation beyond ALE could provide more comprehensive results and insights, potentially at lower computational expense.Source: §7 · pp. 8–9
- The results support further investigation of MoEs in deep RL and underscore the impact of architectural choices on agent performance.The discussion also points to sparsity’s role in parameter scalability as an area for further understanding.Source: §7 · pp. 8–9
A. Experimental details
The appendix specifies the evaluation reporting convention and game-set coverage, then points to default hyperparameter tables for the agent and network families.
A. Experimental details
- Experiments generally report interquantile mean after 40 million environment steps with 95% stratified bootstrap confidence intervals, using 20 ALE games, 26 Atari 100k games, or sometimes all 60 games.The appendix identifies the standard 26-game set for Atari 100k agents and notes that specific game lists follow.Source: A. Experimental details · p. 14
- The appendix lists the games in the 20-game, 26-game, and 60-game evaluation sets.The 60-game set consists of the 26-game set plus 34 additional games.Source: A. Experimental details · p. 14
B. Hyper-parameters list
- Default hyperparameter settings are tabulated for DER and DrQ(ϵ), DQN and Rainbow, CQL and CQL+C51, and the CNN and Impala-based ResNet architectures.The tables are identified as Tables 1–4, respectively.Source: B. Hyper-parameters list · pp. 14–16
C. Extra results
Extra Atari experiments report stable training and higher final performance for Soft MoE on DrQ(ϵ) and DER, while filter-size experiments find SoftMoE outperforms the baseline with larger Impala filters.
C. Extra results
- The extra-results figures cover DQN architectural ablations and Rainbow architectural exploration, including investigation of the original Soft MoE normalization.The captions identify the experiments but do not state their outcomes.Source: Figure 14 · p. 17, C. Extra results · p. 17, Figure 15 · p. 18
- Across 26 Atari games, Soft MoE remains generally stable with more training and attains higher final performance for DrQ(ϵ) and DER using ResNet with 4 experts.The figure reports interquartile mean performance with 95% confidence intervals.Source: Figure 16 · p. 18
- For CQL+C51, normalized performance is plotted across 17 Atari games over gradient steps, with 10% and 50% uniform replay and no new data collected.The figure reports IQM with 95% stratified bootstrap confidence intervals.Source: Figure 17 · p. 19
E. Measuring runtime
Runtime comparisons plot IQM of human-normalized Rainbow scores against wall time. SoftMoE and the baseline show no noticeable running-time difference, while Top1-MoE is slightly faster.
E. Measuring runtime
- SoftMoE and the baseline have no noticeable running-time difference, whereas Top1-MoE is slightly faster than both.The comparison plots IQM of human-normalized Rainbow scores against wall time rather than environment frames.Source: E. Measuring runtime · p. 20, Figure 19 · p. 20
- Runtime is compared against IQM of human-normalized Rainbow scores over 20 games using ImpalaCNN and CNN networks.Source: Figure 19 · p. 20
- Each runtime experiment uses three independent runs, with 95% confidence intervals.Source: Figure 19 · p. 20
F. Experiments with PPO
Initial PPO and SAC experiments on MuJoCo show no significant performance gains or degradation with SoftMoE, while Top1-MoE degrades performance. The authors suggest network type and environment complexity as possible reasons for the lack of SoftMoE improvement.
F. Experiments with PPO
- In initial MuJoCo experiments, SoftMoE shows no significant performance gains or degradation with PPO and SAC, while Top1-MoE degrades performance.The figure covers SAC on 28 MuJoCo environments and PPO on 9 MuJoCo-Brax environments.Source: F. Experiments with PPO · p. 21, Figure 20 · p. 21
- The authors suggest MoE-induced sparsity may be most effective with convolutional layers, whereas their MuJoCo PPO and SAC networks use dense layers.Source: F. Experiments with PPO · p. 21
- The authors also suggest MuJoCo environments may be less complex than ALE environments, leaving SAC and PPO performance somewhat saturated.Source: F. Experiments with PPO · p. 21