Don't flatten, tokenize! Unlocking the key to SoftMoE's efficacy in deep RL
Authors: Ghada Sokar
Johan Samir Obando Ceron
Aaron Courville
Hugo Larochelle
Pablo Samuel Castro
Venue: ICLR 2025 Spotlight (2025)
Keywords: Reinforcement learning
Deep reinforcement learning
Mixture of experts
Primary area: reinforcement learning
TL;DR: We demonstrate that one of the main reasons for the efficacy when using mixtures of experts comes from the choice of tokenization, and we demonstrate that even with a single expert, tokenization is enough to yield the previously observed benefits.

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.

Published as a conference paper at ICLR 2025
DON'T FLATTEN, TOKENIZE! UNLOCKING THE KEY TO SOFTMOE'S S EFFICACY IN DEEP RL
Ghada Sokar
Johan Obando-Ceron
Aaron Courville
Google DeepMind
Mila, Université de Montréal jobando0733 gma 1 com
Mila, Université de Montréal courvi 1 a@mi 1 quebec
bookar@google.com
Hugo Larochelle
Pablo Samuel Castro
Google DeepMind
Google Deep.ind
hugglaarcchelleeegogiiee.ccm com
Mila, Université de Montréal psccgoogle, com
ABSTRACT
The use of deep neural networks in reinforcement learningg (RL) often suffers from performance degradation as model size increases. While soft mixtures experts (SoftMoEs) have recently shown promise in mitigating this issue for online RL, the reasons behind their effectiveness remain largely unknown. this work we provide in-depth analysis identifying the key factors driving this performance gain. We discover the surprising result that tokenizing the encoder output, rather than the of multiple experts, is what behind the efficacy of SoftMoEs. In- deed, we demonstrate that even with an appropriately scaled single expert, we are able maintain the performance gains, largely thanks to tokenization.
INTRODUCTION
Deep networks have been central to many of the successes of reinforcement learning (RL) since their initial use in DQN (Mnih et al., 2015a), which has served as the blueprint (in particular with regards to network architecture) for most of the followup works. However, recent works have high- lighted the need to better understand the network architectures used, and question whether alternate
Baseline
SoftMoE
SoftMoE-1
8.0
6.0
10M
4.0
2.0
2
4
8
Num experts / layer multiplier
Comparison of IQM (Agarwal et al., 2021) for Rainbow (Hessel et al., 2018), using the Impala architecture (Espeholt et al., 2018), with penultimate layer scaled, and with SoftMoE with varying numbers of experts, and single, scaled, expert. IQM scores computed over 200M envi- ronment steps across 20 games, with independent runs each, higher better. Error bars represent 95% stratified bootstrap confidence intervals. A single scaled expert matches the performance of multiple experts.
Figure 1:

---

Published as a conference paper at ICLR 2025
choices can be more effective (Graesser et al., 2022; Sokar et al., 2022; 2023; Obando Ceron al., 2024; Lee et al., 2025). Indced, Obando Ceron et al. (2024) and Willi* et al. (2024) demon- strated that the use of of (MoEs) (Shazeer et al., 2017), and in particular soft mixtures-of-experts (SoftMoFs) (Puigcerver et al., 2024), yields models that are more parameter efficient. Obando Ceron R et al. (2024) hypothesized that the underlying cause of the effectiveness of SoftMoEs was in the induced structured sparsity in the network architecture. While some of their analyses did seem to support this, the fact that performance gains were observed even with a single expert seems to run counter to this hypothesis.
In this work we seek to better understand the true underlying causes for the effectiveness of Soft- MoEs in RL settings. Specifically, we perform a careful and thorough investigation into the many parts that make up SoftMoEs aiming to isolate the impact of each. Our investigation led to the sur- prising finding that the core strength of SoftMoEs lies in the tokenization of the outputs of the convolutional encoder.
The main implications of this finding are two-fold. On the one hand, it suggests that we are still under-utilizing experts when incorporating MoEs in deep RL, given that we are able to obtain com- parable performance with a single expert (see Figure 1). On the other hand, this finding has broader implications for deep RL agents trained on pixel-basee based environments, suggesting that the common practice of flattening the multi-dimensional outputs of the convolutional encoders is ineffective.
Our main contributions can be summarized as follows:
We conduct series of amalyses to identify the key factors that allow soft Mixture of Experts (SoftMoEs) to effectively scale RL agents.
We demonstrate that tokenization plays a crucial role in the performance improvements.
We demonstrate that experts within the SoftMoFs are not specialized to handle specific subsets of tokens and exhibit redundancy. We explore the use of recent techniques for improving network utilization and plasticity to mitigate this redundancy.
The paper is organized as follows. We introduce the necessary background for RL in Section 2 and MoEs in Section 3; we present our experimental setup in Section 4. and present our findings from Section 4.2 to Section 6; finally, we discuss related works in Section 7 and provide a concluding discussion in Section 8.
2 ONLINE DEEP REINFORCEMENT LEARNING
In online reinforcement learning, an agent interacts with an environment and adjusts its behaviour based on the reinforcement (in the form of rewards or costs) it receives from the environment. At timestep I, the agent inhabits state Ii E X' (where X' is the set of all possible states) and selects an action dt E A (where A is the set of all possible actions) according to its be- haviour policy π X + B(A); the environment's dynamics yield new state I(+1 X and a numerical reward r; E R. The goal of an RL agent is to find a policy n² that max- imizes V"(x)
Ee= hr'r//0 x., a ~ πt (xt))], where 7 E [0. 1) is discount factor. In value-based RI, an agent maintains an estimate of Q" X' x A + R, defined as Q" (x, a) = Ee (7/re/20 xr, ao a, at T(æ:]]: from any state E X, policy и can then be improved via (T) = arg max.eA Q (I, a) (Sutton & Barto, 1998).
In its most basic form, value-based deep reinforcement learning (DRL) extends this approach by using neural networks, parameterized by 0, to estimate Q. While the network architectures used may vary across environments, the one originally introduced by Mnih et al. (2015a) for DQN has served as the basis for most. Given that its inputs consist of frames of pixels from Atari games (Bellemare et al., 2013). this network consisted of a series of convolutional layers followed by a series of fully-connected layers; for the remainder of this work we will refer to the convolutional layers as the encoder. The top panel of Figure 2 displays a simplified version of this network, where the encoder produces 3-dimensional tensor with dimensions (h, w, d); this tensor is flattened to vector of length * w d, which is fed through fully-connected layers to produce the estimates for Q. One notable variant to the original DQN encoder is Impala's 's use of a ResNet encoder (Espeholt et al., 2018), which has been shown to outperform the former. In this work, we mainly focus on the more recent ResNet encoder, but also include some results using the original DQN encoder.
2

---

Published as a conference paper at ICLR 2025
Flatten
Q-values
Encoder output
Weighted toliens
Takenize
Flatton
Q-values
Assign one per
Tokenize
Router
Expsor
B
Scale
Q-values
procosses slots
Figure 2: The various architectures considered in this work. Top: baseline architecture; Middle: SoftMoE architecture; Bottom: Top-k architecture.
3 MIXTURES OF EXPERTS
While there have been number of architectural variations proposed, in this work we focus on Mixtures of Experts used by Obando Ceron" et al. (2024). Mixtures of experts (MoEs) were originally proposed as a network architecture for supervised learning tasks using conditional com- putation, which enabled training very large networks (Shazeer al., 2017; Fedus al., 2022). This was achieved by including MoE modules which route a set of imput tokens to set of sub- nut-networks, or experss; the output of these experts are then combined and fed through the rest of the outer net- work. The routing/gatii we consider two methods for pairing tokens and experts. In token choice routing, each token gets an assignation to k experts, each of which has slots (the number of processed inputs by each expert); this may result in certain tokens being dropped entirely by the MoE module and, more critically, certain experts may receive tokens (Gale et al., 2023). To overcome having load-imbalanced experts, Zhou et al. (2022) proposed expert choice routing, where each expert chooses p tokens for its p slots, which results in better expert load-balancing balancina.. See the bottom row of Figure 2 for an illustration of this. Unless otherwise specified, use value of p equals to #tokens In most of #ezperts"
mechanism and the expert sub-networks are all trainable. In top-k routing,
our analyses, focus expert choice routing demonstrates higher performance than token choice routing (see Appendix B. for full empirical comparison between the two across varying number of experts).
Puigcerver et al. (2024) introduced SoftMoEs, which replace the hard token-expert pairing with "soft" assignation. Specifically, rather than having each expert slot assigned specific token, each slot is assigned a (learne)) weighted combination of all the tokens. Thus, each expert receives weighted averages of all the input tokens. See the middle row of Figure 2 for an illustration of this.
Obando Ceron et al. (2024) proposed replacing the penultimate dense layer of a standard RL net- work with an MoE module; the authors demonstrated that, when using SoftMoEs, value-based on- line RI. agents exhibit signs of "scaling law", whereby adding more parameters (in the form of more experts) monotonically improves performance. This is in stark contrast with what is observed with the original architectures, where extra parameters result in performance degradation (see Fig- ure I). The authors argued that the structured sparsity induced by the MoE architecture what was primarily responsible for their strong performance.
In order to maintain the shape of the input tokens (as originally done with MoE architectures), Obando Ceron et al. (2024) add projection layer after the "combine" step (&c in Figure 2). Another important detail is the proposed tokenization scheme. MoEs are most commonly used in NLP transformer architectures where tokens are clearly defined (typically encodings of words),
3

---

Published as a conference paper at ICLR 2025
Shuffle
hyp wilp tokens
Io tokens
sokens
Flatten
PerConv
PerFeat
PerPatch
Shuffled
Figure 3: Tokenization schemes considered in this work, specified in the bottom row.
but this tokenization is not something readily available for online RL agents. Thus, Obando Ceron al. (2024) proposed a number of mechanisms for tokenizing the output of the encoder.
Tokenization As mentioned in section 2, the output of standard RL encoder is 3-dimensional tensor with dimensions (h, w, d); Obando Ceron et al. (2024) proposed a few possible ways of creating tokens from this output. PerConv tokenization creates how tokens of dimension d; PerFeat tokenization creates tokens of dimension h*w; see Figure for an illustration. Although necessary for compatibility with transformer MoEs, this tokenization is departure from the standard practice of flattening the encoder output, as discussed in Section 2.
4 UNDERSTANDING THE IMPACT OF THE SOFTMOE COMPONENTS
While Obando Ceron et al. (2024) demonstrated the appeal of SoftMoEs for online RL agents, was not clear whether all of the components of SoftMoEs were important for the performance gains, or whether some were more critical. Further, it was observed that SoftMoEs failed to provide similar gains in actor-critic algorithms such as PPO (Schulman et al., 2017) and SAC (Haarnoja et al., 20.8). This section aims to isolate the impact of each of these components.
4.1 EXPERIMENTAL SETUP
The genesis of our work to provide deeper understanding of SoftMoES in online RL. For this reason, we will mostly focus on SoftMoEs applied to the Rainbow agent (Hessel et al., 2018) with 4 experts and the Impala achitecture (Espeholt et al., 2018), as this was the most performant setting observed by Obando Ceron et al. (2024). We explore variations of this setting where useful in clarifying the point being argued, We evaluate performance on the same games of the Arcade Learning Environment (ALE) (Bellemare et al., 2013) benchmark used by Obando Ceron* et al. (2024) for direct comparison, with 5 independent seeds for each experiment. However, we include the results of our main findings the full suite of 60 games. We will use the notation "SoftMoE-n" to denote a SoftMoF. architecture with n experts. All experiments were run on Tesla P100 GPUs using the same Dopamine library (Castro al., 2018) used by Obando Ceron et al. (2024); each run with 200 million environment steps took between six and eight days. Following the guidelines suggested by Agarwal al. (2021), we report human-normalized aggregated interquartile mean (IQM), regular mean, median, and optimality gap, with error bars indicating 95% stratified bootstrap confidence intervals. Performance gains are indicated by increased IQM, median, and mean scores, and reduced optimality gaps. See Appendix A for further details.
4.2 ANALYSIS OF SOFTMOE COMPONENTS
At a high level, SoftMoE modules for online RL agents consist (i) the use learnable tensor & used to obtain the dispatch (&p) and combine ($c) weights; (ii) processing p input slots per expert; (iii) the architectural dimensions (network depth and width); (iv) the use of n experts; and (v) the tokenization of the encoder output. The middle row of Figure provides illustration of these. In
4

---

Published as a conference paper at ICLR 2025
SoftMoE-4
Baseline (scaled x4)
[i] Combined tokenization
ExpertChoice
[ii] Expert specialization
SoftMoE-4 (Expert's #slots #tokens)
ExpertChoice-4 (Expert's #slots = #tokens)
[îii] Width
SoftMoE-4 (scaled x4)
[iil] Depth
Baseline (scaled x4) with extra layer
2.0 2.5
3.5
4.0
5.5 6.0 6.5 7.0
Interquartile Mean (IQM)
Figure 4: Understanding the impact of SoftMoE components. Using Rainbow with the Impala architecture as our base, we investigate key aspects of SoftMoE-4: (i) combining tokens, (ii) ex- pert specialization, and (ii) adjusting architectural dimensions. Reporting IQM (Agarwal et al., 2021), where higher is better. SoftMoE does not appear to suffer from the performance degradation observed in the baseline, even when increasing its width.
this section we discuss the first four with empirical evaluation demonstrated in Figures and 4, and devote separate section for tokenization, which turns be crucial for performance. The same analysis was performed on DER, yielding the same results, and can be found in Appendix B.3.
To assess the impact of SoftMoE's weighted combination of tokens via the learnable weights t we compare it with expert choice routing, where experts select sets of (non-combined) tokens. While combined tokens yield improved performance, expert choice rout- ing itself demonstrates substantial gains over the baseline and effectively scales the network (see group [i] in Figure 4). Thus, combined tokens alone do not explain SoftMoE's efficacy.
Processing combined tokens
Expert specialization Each expert in MoE architecture has limited set of slots: with to- ken/expert selection each slot corresponds to one token, while in SoftMoE each slot corresponds to particular weighted combination of all tokens. This constraint could encourage each expert to specialize in certain types of tokens in a way that is beneficial for performance and/or scaling. To investigate, increase the number of slots to be equal the number of tokens (effectively granting unlimited capacity). Group [ii] in Figure 4 demonstrates that for both SoftMoE and expert choice routing, the performance sees little change when increasing the number of slots, suggesting that expert specialization is not major factor in SoftMoE's 's performance.
Expert width Obando Ceron al. (2024) scaled up the penultimate layer of the baseline model with the intent of matching the parameters of the SoftMoF. networks, thereby demonstrating that naïvely scaling width degrades performance. While they demonstrated that performance drop not observed when scaling down the dimensionality of cach expert, they did not investigate what occurs when scaling up the experts dimensionality. This experiment makes sense if comparing the baseline's 's penultimate layer dimensionality with that of each expert, rather than with the combined experts dimensionality. Scaling up the width of experts in SoftMoE does not lead to the same performance degradation observed in the scaled baseline model shown in the first group [iii] of Figure
Network depth As mentioned Section SoftMoE models have an additional projection layer within each expert to maintain input token size, resulting in a different number of layers compared the baseline model. To determine whether this extra layer contributes SoftMoE's performance gains, we add layer of the same size to the baseline model after the penultimate layer. The second group [iii] Figure demonstrates that this added layer does not improve the baseline model's 's per- formance nor its ability to scale, suggesting that SoftMoE'ss performunce gains are not altributable to the presence of the extra projection layer.
Multiple experts Obando Ceron et al. (2024) demonstrated that multiple experts improves per- formance. However, the performance gains plateau beyond a certain point, with minimal improve- ment observed after utilizing four experts (see Figure 1). Consequently, the presence of multiple experts cannot fully account for the observed performance advantages over the baseline model, sug- gesting that other factors are also contributing the improved results.
5

---

Published as a conference paper at ICLR 2025
Median
I0M
Mean
Optimaiity Gap
Baseline
Tokenizedhep (PerCorv)
TakenizedSum (Perfeat
Takenzzdkug (ParFaaO
Bascline
TakencceSSum (PerCorv))
TakenizedArg (ParCort)
TolenizedSum (Perfeat)
Tekenizefogg (Perfeat)
Human Normllizdd Scorei
Impala
Median
IOM
Mean
Optimality Gap
Baseline
Take rizedSum (PerCorw)
Tocenizedffu (PerConv)
TntenizedSumm (PerTeaD)
beleiizeAeg (Perfeat))
CNN
Baseline
Take rizcoSum (PerConw)
Tokenizeddog (PerConv)
TekenizedSum (Perfeal)
bloeizzdAeg (PerFeal)
Human Normalzed Store
Figure 5: Evaluating impact of tokenizing Rainbow-lite with unscaled (top) and scaled (bottom) architectures, using sum/mean over convolutional features, and exploring both PerConv and PerFeat tokenization schemes. Reporting Median, IQM, Mean, and Optimaly Gap scores (Agarwal al., 2021). Higher is better for all except Optimality Gap. Tokenization with scaled representation can yield significant significant performance gains. Per-game results available Appendix C.I.
DON'T FLATTEN, TOKENIZE!
A erucial structural difference remains in how SoftMoE and the baseline model process the encoder's output: while the baseline model flattens the output of the final convolutional layer, SoftMoE divides it into distinct tokens (see Section 3). We conduct series of experiments to properly assess the impact of this choice, going from the standard baseline architecture to a single-expert SoftMoF.
Tokenized baseline Our first experiment simply adds tokenization to the baseline architecture. Specifically, given the 3-dimensional encoder output with shape [h, w, d]. replace the standard flattening (to a vector dimension h*w=d) with a reshape to matrix of shape [h=w, d] for PerConv tokenization [d, h* w| for PerFeat; then, prior to the final linear layer, either sum over the first dimension or take their average. Since this operation is not immediately applicable to Rainbow (due to Noisy networks), in Dopamine which only includes C51 (Bellemare et al., 2017), multi-step returns, and prioritized experience replay. Figure demonstrates that this simple change yield statistically significant improvements for Rainbow-lite, providing strong evidence that tokenization plays major role in the successful scaling of DRL networks.
explored this experiment with Rainbow-lite, version of Rainbow
Single expert processing all tokens sparsely To assess the influence of tokenization within the MoE framework, while ablating all other four components studied in Section 4.2, we explore using single expert to process all input tokens. Each token is sparsely processed in of the single expert slots. This corresponds to expert choice-1. In line with the "expert width" investigation from study expert choice-1 with the scaled expert dimensionality to match the scaled
the last section, baseline one. As Figure shows, even with single expert the performance is comparable to that of SoftMoE with multiple experts. This suggests that tokenization one of the major factors in the performance gains achieved by SoftMoE.
Combined tokenization a key driver We further investigate the impact of combined tokeniza- tion in the single expert setting. To this end, we study SoftMoE-1, with scaled expert dimensionality as before. As Figure shows, using combined tokens with a single expert brings the performance closer the multiple expert setting. Similar results were observed in SoftMoB. with 2 and

---

Published as a conference paper at ICLR 2025
experts (See Appendix B.2 for more details). This further confirms that (combined) tokenization is the key driver of SoftMoE's efficacy.
5.1 ADDITIONAL ANALYSES
The main takeaway from the previous section is that SoftMoE with single (scaled) expert as performant as the variant with multiple unscaled experts. To better understand this finding, this section we perform an extra set of experiments with this variant of SoftMoE.
Are all types of tokenization equally effective? We compare the different tokenization types pre- sented in section 3, addition two new types of tokenization, which serve as useful references: PerPatch tokenization splits the 3-dimensional [h, w, d] tensor into p patches, over which average pooling is performed resulting in new tensor dimensions [h/p: w/ p, d]: Shuffled is similar to the encoder output. See Figure 3 for an il-
PerConv, except it first applies fixed permutation lustration of these. The comparison the different tokenization types is summarized in Figure 7. Consistent with the findings of Obando Ceron et al. (2024), PerConv performs best, and in par- ticular outperforms PerFeat. PerConv, while Shuffled is the worst performing. The degradation from PerPatch is likely due to the average pooling and reduced dimensionality: the degradation from shuffling is clearly indicating that it is important maintain the spatial structure of the encoding output. Combining these results with our previous findings suggests that using tokenization maintains the spatial structure learned by the convolutional encoder, and is crucial to maximizing performance.
interesting to observe that PerPatch is close in performance to
Computational efficiency by combined tokens Puigcerver et al. (2024) highlighted that the time complexity of SoftMoE depends directly on the number of slots. As can be seen in Figure 8, using SoftMoE-1 with only 10% of the default number of slots (kept fixed throughout training), and even with a single slot, achieves comparable performance to using the regular number of slots. On the other hand, the reduced performance of ExperiChoice-1 further confirms the gains added by combining tokens in SoftMoE.
More games and different encoders Our results far were run the 20 games used by Obando Ceron et al. (2024) and with the Impala ResNet encoder (Espeholt et al., 2018). To en- our findings
not specific these choices, extend our evaluation all 60 games in the ALE suite, and with the standard CNN encoder of Mnih et al. (2015b). Our findings, as illustrated in Figure 9, demonstrate consistent results across all games and various encoder architectures. We further conducted extra experiments on Proegen (Cobbe et al., 2019), which are consistent with our findings (Figure 21 in Appendix C).
More agents We focused Rainbow for the majority of our investigation due the positive results demonstrated by Ohando Ceron et al. (2024), but it is important to also consider other agents. To this point, we evaluate the impact of single scaled expert on DQN (Mnih al., 2015b)
SoftMoE
4 experts
expert scaled
1 expert
ExpertChoice
1 expert scaled x4
Baseline
Penultimate layer scaled x4
3.0
3.5
4.0
5.0
5.5
6.0
7.0
7.5
Interquartile Mean (IQM)
Figure 6: Combined tokenization is a key driver of S Reporting IQM (Agarwal al., 2021), where higher better. SoftMoE with single scaled expert has comparable perfor- mance with using multiple experts. Even non-combined tokenization (e.g. ExpertChocce) with single expert achieves similar performance.

---

Published as a conference paper at ICLR 2025
Shuffled
Perfeat
PerPatch
PerConv
6.0
4.0
2.0
Num experts layer multiplier
Figure 7: Comparison of different tokenization schemes for SoftMoE-1 with the standard and scaled penultimate layer. Reporting IQM (Agarwal et al., 2021), where higher better. PerConv achieves the highest performance, suggesting that maintaining the spatial structure of the features plays a role on the observed benefits.
and DER (Van Hasselt et al., 2019). Note that we replace the original Huber loss used in DQN with mean-squared error loss, as it has been shown to be more performant (Ceron & Castro, 2021), and coincides with what was used by Obando Ceron al. (2024). Following the training regime used by Obando Ceron et al. (2024), we train DER for 50M environment frames. In the DQN setting, SoftMoE-1 does not appear offer significant improvements over the baseline, potentially due to the limited gains observed with SoftMoFF-4 (Figure 10, top). Further investigation is needed to understand how the inherent instability of DQN might limit the potential benefits of SoftMoE. On the other hand, consistent with finding in Rainbow, in DER, where SoftMoE-4 achieves superior performance over the baseline, SoftMoE-1 is able to offer the same gains (Figure 10, bottom). To further evaluate the generality of our claims, conducted initial experiments with SAC (Haarnoja et al., 2018) on the CALE (Farebrother Castro, 2024), which is a continuous action version of the original ALE. In Figure 22 in Appendix it can be seen that tokenization provides an improvement over flattening.
CALLING ALL EXPERTS!
Our findings in Sections 4 and 5 reveal that the main benefit of using SoftMoE architectures is in tokenization and their weighted combination, rather than in the use of multiple experts. The fact that we see evidence of expert specialization suggests that most experts are redundant (see Figure 4). Rather than viewing this as negative result, it's an opportunity to develop new techniques to improve expert utilization and maximize the of MoE architectures in RL.
As first step in this direction, we adopt existing techniques designed to improve network utilization: (1) periodic reset of parameters (Nikishin et al., 2022) and (2) Shrink and Perturb (S&P) (Ash & Adams, 2020), as used by D'Oro et al. (2022) and Schwarzer* et al. (2023). We perform these experiments Rainbow and DER with SoftMoE-4. We explore applying resets and S&P to either all experts or subset of experts, where the experts to reset are those that have the highest fraction of
SoftMoE
Using 10% of slots
SoftMoE
ExpertChoice
Single
CxpirtChoice
200
25
3.0
3.5
6.0 65 Interquartile Mean (IQM)
4.0
5.0 5.5
7.0
75
Figure 8: Evaluating the impact fewer slots with single scaled (by 4x) expert for both SoftMoE and ExpertChoice. Reporting IQM (Agarwal et al., 2021), where higher is better. Even with drastically fewer slots, SoftMoE-1 remains competitive.
8

---

Published as a conference paper at ICLR 2025
Saseliee oocaled wdi SottMoz
SoftMEE 1 localedd x4)
Endhatoo ond to gamss
Endiaton Chh MoL.coii
Figure 9: (Left) Evaluating the performance on all of the 60 games of Atari 2600. Reporting IQM (Agarwal al., 2021), where higher is better. Consistent with our results SoftMoE-1 matches the performance of multiple experts over the full suite. (Right) CNN Encoder. The performance gain of SoftMoF-1 present when training with the standard CNN encoder, revealing the importance of tokenization across architectures.
dormant neurons (Sokar et al., 2023). We reset every 20M and 200K environment steps for Rainbow and DER, respectively, More experimental details can be found in Appendix B.5.
Interestingly, the effectiveness of these techniques varied across different RL agents. While weight reset and S&P did improve expert utilization the Rainbow agent, they led enhanced perfor- mance in the sample-efficient DER agent (Figure 11). Within the DER agent, resetting a subset of experts proved more beneficial than resetting all, while applying S&P to all experts yielded most significant performance gains. These findings highlight the need for further research into methods that can optimize expert utilization and improve performance in RL with SoftMoE.
RELATED WORKS
Mixture of Experts Mixture of Experts (MoEs) have demonstrated remarkable success in scaling
language and computer vision models trillions of parameters (Lepikhin et al., 2020; Fedus al.,
2022: Wang et al., 2020; Yang et al., 2019; Abbas Andreopoulos, 2020; Pavlitskaya et al., 2020). Chung et al. (2024) examined the implicit biases arising from token combinations in SoftMoE and defined the notion of expert specialization in computer vision. Besides the success MoEs in single task training, similar gains have been observed in transfer and multitask learning (Puigcerver et al., 2020; Chen et al., 2023; Ye & Xu, 2023; Hazimeh et al., 2021). In addition to the works of Obando Ceron al. (2024) and Willi* al. (2024) already discussed, Farebrother al. (2024) proposed replacing regression loss with classification loss and showed its benefits on MoEs.
Network scalability in RL Scaling reinforcement learning (RL) networks presents significant challenges, largely due to issues with training instabilities (Hessel et al., 2018). To address these difficulties, several techniques have been proposed, with a focus on algorithmic improvements or hyperparameter optimization (Farebrother et al., 2022: 2024; Taiga et al., 2022; Schwarzer* al., 2023). Recently, a few studies have begun to focus on architectural innovations as a means of addressing these challenges. Kumar al. (2023a) showed that incorporating a leared spatial em- bedding with the output of a ResNet encoder, prior to flattening it, improves the scalability in offline RL. Obando Ceron al. (2024) demonstrated that gradual magnitude pruning enables the scaling of value-based agents, resulting in substantial performance gains.
Median
IQM
Mean
Optimality Gap
Baseline (scaled x4) SoftMoE-1 (scaled x4)
SoftMeE-1
2.8
3.0
3.4
3.00
3.50 Human Normalized
6.6
0.060 0.075 0.090
Median
IQM
Mean
Optimality Gap
Baseline (scalled
SoftMoE-1 (scaled x4)
S
1.2
1.6
2.0
1.2
1.6
2.0
1.2
1.6
2.0
0.20 0.24 0.28 0.32
Human Normal ized Score
Figure 10: Evaluating on DQN (top) and DER (bottom). Reporting Median, IQM, Mean, and Optimaly Gap scores (Agarwal et al., 2021). Higher is better for all except Optimality Gap. When combined with DQN, SoftMoE with a single scaled expert does not provide gains over the baseline. In contrast, when combined with DER, achieves performance comparable to 4

---

Published as a conference paper at ICLR 2025
Rainbow
Rezet 4 experts
Rexel esperts
No reets
S&P 4 experts
S&P esperts
No SGP
IQM Human Normalized Score
Figure 11: Evaluating the use of resets (top) and S&P (bottom) for improving expert utilization on Rainbow (left column) and DER (right column). Reporting IQM (Agarwal et al., 2021), where higher is better. All experiments run with SoftMoEE4. Resets and S&P seem to benefit DER, but degrade performance for Rainbow.
Network plasticity in RL It has been observed that networks trained on non-stationary distribu- tions tend to lose their expressivity over time, phenomenon referred to as the loss of plasticity (Kumar et al., ;021; Lyle et al., 2021). While several studies have attempted to explain this phe- nomenon (Lyle al., 2023; Ma et al., 2023; Lyle et al., 2024), the underlying causes remain un- certain. In a parallel line of research, numerous methods have been proposed to improve network trainability. These approaches involve using small batch size for gradient updates (Ceron et al., 2023), periodically resetting some of the network's 's parameters to random values (Sokar et al., 2023; Nikishin al., 2022; Igl al., 2020; Dohare et al., 2024), slightly modifying the existing weights by shrinking and perturbing them (D' Oro et al., 2022: Schwarzer* et al.. 2023), and restricting the large deviation from initial weights (Kumar al., 2023b; Lewandowski et al., 2024). Recently, Lee et al. (2024) have also demonstrated that combining reinitialization with teacher-student frame- work can enhance plasticity. Another technique involves preserving useful neurons by perturbing gradients to ensure smaller updates for important neurons (Elsayed & Mahmood, 2024). Liu et al. (2025) proposed dynamic network growth through neuron generation and dormant neuron pruning. In this work, explore some of these methods to improve plasticity and utilization of experts.
CONCLUSIONS
We began our investigation by evaluating the components that make up SoftMoE architecture, so as to better understand the strong performance gains reported by Obando Ceron et al. (2024) when applied to value-based deep RL agents. One of the principal, and rather striking, findings is that tokenizing the encoder outputs is one the primary drivers of the performance improvements. Since tokenization is mostly necessary for pixel-based environments (where encoders typically consist of convolutional layers), it can explain the absence of performance gains observed by Obando Ceron* al. (2024) and Willi* al. (2024) when using MoE architectures in actor-critic algorithms, since these were evaluated in non-pixel-based environments. Initial experiments illustrated in Figure 22 suggest that our findings carry over to actor-critic agents continuous control environments.
More generally, our findings suggests that the common practice of flattening the outputs of deep RL encoders is sub-optimal, it is likely not maintaining the spatial structure output by the con- volutional encoders. Indeed, the results with PerConv versus Shuffled tokenization in Figure demonstrate that maintaining this spatial structure is important for performance. This is further confirmed by our results in Figure 5, where we can observe that not flattening the convolutional outputs can lead to significant improvements. This idea merits further investigation, as the results vary depending on the architecture used and whether features are summed or averaged over.
somewhat surprising that our results do not seem to carry over to DQN. However, the fact that do see similar results with Rainbow-lite and DER (which is based on Rainbow) suggests that the use of categorical losses may have a complementary role to play. This is a point argued by Farebrother al. (2024) and merits further inquiry.
Although the use of cawxxtures-of-expeeee that most of the experts not being used to their full potential. The findings in Figure 11 present one viable mechanism for increasing their use, but further research is needed to continue increasing their efficacy for maximally performant RI. agents.
architectures for RL show great promise, our results suggest
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENTS
The authors would like to thank Gheorghe Comanici, Jesse Farebrother, Joan Puigcerver, Doina Pre- cup, and the rest of the Google DeepMind team for valuable feedback on this work. We would also like to thank the Python community (Van Rossum & Drake Jr, 1995; Oliphant, 2007) for developing tools that enabled this work, including NumPy (Harris et al., 2020), Matplotlib (Hunter, 2007) and JAX (Bradbury et al., 2018).
REFERENCES
Alhabib Abbas and Yiannis Andreopoulos. Biased mixtures of experts: Enabling computer vision inference under data transfer limitations. IEEE Transactions on Image Processing, 29:7656-7667, 2020.
Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C Courville, and Mare Bellemare. Deep reinforcement learning at the edge of the statistical precipice. In M. Ranzato, A. Beygelz- imer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 29304-29320. Curran Associates, Inc., 2021.
Jordan Ash and Ryan P Adams. On warm-starting neural network training. Advances in neural information processing systems, 33:3884-3894, 2020.
Marc G Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The arcade learning environ- ment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47: 253-279, 2013.
Mare G. Bellemare, Will Dabney, and Rémi Munos. A distributional perspective on reinforce- ment learning. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pp. 449-458, PMLR, 06-11 Aug 2017. URL https: http:///poccceeeeeeeeeeee www.mlrr.pp...... A.L..press/v00. bellemare117 7a. html.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Honderman-Milng et al. Jax: composable transformations of python+ numpy programs. 2018.
Pablo Samuel Castro, Subhodeep Moitra, Carles Gelada, Saurabh Kumar, and Mare G. Bellemare. Dopamine: A Research Framework for Deep Reinforcement Learning. 2018. URL http:
/ ///paxiv.rrgaaas/////////.o.oaaa /182. 06110.
Johan Samir Obando Ceron and Pablo Samuel Castro. Revisiting rainbow: Promoting more insight- ful and inclusive deep reinforcement learning research. In International Conference Machine Learning, pp. 3373-1383 PMLR, 2021.
Johan Samir Obando Ceron, Marc G Bellemare, and Pablo Samuel Castro. Small batch deep re- inforcement learning. seirty- Conference on Neural Information Processing Systems, 2023. URL https: openreview, net/ www.oooorooooroooorooorooo d Eh.
Zitian Chen, Yikang Shen, Mingyu Ding, Zhenfang Chen, Hengshuang Zhao, Erik G Learned- Miller, and Chuang Gan. Mod-squad: Designing mixtures of experts as modular multi-task learn- ers. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattrrn Recognition, 2023.
pp.
Youngseog Chung, Dhruv Malik, Jeff Schneider, Yuanzhi Li, and Aarti Singh. Beyond parameter count: Implicit bias in soft mixture of experts. arXiv preprint 0tt::24090000770 2024.
Karl Cobbe, Christopher Hesse, Jacob Hilton, and John Schulman. Leveraging procedural genera- tion to benchmark reinforcement learning. arXiv preprint arXiv: 19/2.0158. 2019.
Shibhansh Dohare, Fernando Garria, Qingfeng Lan, Parash Rahman, A Rupam Mah- mood, and Richard Sutton. Loss of plasticity in deep continual learning. Nature, 632(8026): 768-774, 2024.
11

---

Published as a conference paper at ICLR 2025
Pierluca D'Oro, Max Schwarzer, Evgenii Nikishin, Pierre-Luc Bacon, Mare G Bellemare, and Aaron Courville. Sample-efficient reinforcement learning by breaking the replay ratio barrier. In Deep Reinforcement Learning Workshop NeurlPS 2022, 2022. URL http:::: openreview. net/ foormm? AGBAAAGGGWWEEYE
Mohamed Elsayed and A Rupam Mahmood. Addressing loss of plasticity and catastrophic forget- ting in continual learning. In The Twelfth International Conference on Learning Representations, 2024.
Lasse Espeholt, Hubert Soyer, Remi Munos, Karen Simonyan, Volodymir Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, lain Dunning, et al. Impala: Scalable distributed deep-rl with importance weighted actor-learner architectures. In Proceedings of the International Conference on Machine Learning (ICML), 2018.
Jesse Farebrother and Pablo Samuel Castro. CALE: Continuous arcade learning environment. Ad- vances in Neural Information Processing Systems, 2024.
Jesse Farebrother, Joshua Greaves, Rishabh Agarwal, Charline Le Lan, Ross Goroshin, Pablo Samuel Castro, and Mare G Bellemare. Proto-value networks: Scaling representation learn- ing with auxiliary tasks. In The Eleventh International Conference on Learning Representations, 2022.
Jesse Farebrother, Jordi Orbay, Quan Vuong, Adrien Ali Taïga, Yevgen Chebotar, Ted Xiao, Alex Irpan, Sergey Levine, Pablo Samuel Castro, Aleksandra Faust, Aviral Kumar, and Rishabh Agar- wal. Stop regressing: Training value functions via classification for scalable deep rl. In Forty-first International Conference Machine Learning. PMLR, 2024.
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. The Journal of Machine Learning Research, 23(1): 5232-5270, 2022.
Trevor Gale, Deepak Narayanan, Cliff Young, and Matei Zaharia. MegaBlocks: Efficient Sparse Training with Mixture-of-Expert, Proceedings of Machine Learning and Systems, 5, 2023.
Laura Graesser, Utku Evci, Erich Elsen, and Pablo Samuel Castro. The state of sparse training in deep reinforcement learning. In International Conference on Machine Learning, pp. 7766-7792. PMLR, 2022.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In International confer- ence on machine learning. pp. 1861-1870. PMLR, 2018
Charles R Harris, K Jarrod Millman, Stéfan J Van Der Walt, Ralf Gommers, Pauli Virtanen, David Cournapeau, Eric Wieser, Julian Taylor. Sebastian Berg, Nathaniel J Smith, et al. Array program- ming with numpy. Nature, 585(7825):357-362, 2020.
Hussein Hazimeh, Zhe Zhao, Aakanksha Chowdhery, Maheswaran Sathiamoorthy, Yihua Chen, Rahul Mazumder, Lichan Hong, and Ed Chi. Dselect-k: Diffcrentiable selection in the mixture of experts with applications to multi-task learning, Advances in Neural Information Processing Systems, 34:29335-299333 2021.
Matteo Hessel, Joseph Modayil, H. V. Hasselt, T. Schaul, Georg Ostrovski. Will Dabney, Dan Hor- gan, Bilal Piot, M. G. Azar, and D. Silver. Rainbow: Combining improvements in deep reinforce- ment learning. In AAAL, 2018.
John D Hunter. Matplotlib: A 2d graphics environment. Computing in science & engineering, (03):90-95, 2007.
Maximilian Igl, Gregory Farquhar, Jelena Luketina, Wendelin Bochmer, and Shimon Whiteson.
Transient non-stations stationttonaatt and generalisation in deep reinforcement learning. In International Conference Learning Representations, 2020.
12

---

Published as a conference paper at ICLR 2025
Aviral Kumar, Rishabh Agarwal, Dibya Ghosh, and Sergey Levine. Implicit under- partterrrmmtriiitioonnn inhibits data-efficient deep reinforcement learning. In International Conference on Learning Rep- resentations, 2021.
Aviral Kumar, Rishabh Agarwal, Xinyang Geng, George Tucker, and Sergey Levine. Offline q-learning diverse multi-task data both scales and generalizes. In The Eleventh Interna- tional Conference on Learning Representations, 2023a. URL Btllnttss///peeeeeieee .eet/
Forum?id=4-k7kune
Saurabh Kumar, Henrik Marklund, and Benjamin Van Roy. Maintaining plasticity via regenerative regularization. arXiv preprint arXiv: 11958, 2023b.
Hojoon Lee, Hyeonseo Cho, Hyunseung Kim, Donghu Kim, Dugki Min, Jaegul Choo, and Clare Lyle. Slow and steady wins the race: Maintaining plasticity with hare and tortoise networks. In Forty-first International Conference on Machine Learning. 2024.
Hojoon Lee, Dongyoon Hwang. Donghu Kim, Hyunseung Kim, Jun Jet Tai, Kaushik Subrama- nian, Peter R. Wurman, Jaegul Choo, Peter Stone, and Takuma Seno. Simba: Simplicity bias for scaling up parameters in deep reinforcement learning. In The Thirteenth International Confer- ence on Learning Representations, 2025. URL https:/ httpss//ppeeeevew..oo www.nneee/eetteeeeeeeeeeeeeaeeee - jXLiDXsuDo.
Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. Gshard: Scaling giant models with conditional computation and automatic sharding. In International Conference on Learning Representations, 2020.
Alex Lewandowski, Haruto Tanaka, Dale Schuurmans, and Marlos C. Machado. Directions of curvature as an explanation for loss of plasticity, 2024. URL https:: httpaaaiiiviiiiaaaaaaaaa/aaeae/e 2312. 00246.
Jiashun Liu, Johan Samir Obando Ceron, Aaron Courville, and Ling Pan. Neuroplastic expansion in deep reinforcement learning. In The Thirteenth International Conference on Learning Repre- sentations, 2025. URL
Clare Lyle, Mark Rowland, and Will Dabney. Understanding and preventing capacity loss in rein- forcement learning. In International Conference on Learning Representations, 2021.
Clare Lyle, Zeyu Zheng, Evgenii Nikishin, Bernardo Avila Pires, Razvan Pascanu, and Will Dabney. Understanding plasticity in neural networks. In International Conference on Machine Learning, pp. 23120- 23211. PMLR, 2023.
Clare Lyle, Zeyu Zheng, Khimya Khetarpal, Hado van Hasselt, Razvan Pascanu, James Martens, and Will Dabney. Disentangling the causes of plasticity loss in neural networks. arXiv preprint arXiv:2402.18762. 2024.
G Ma, L,,, Zhang, Z Liu, Z Wang, Y Chen, L Shen, X Wang, and D Tao. Revisiting plasticity in visual reinforcement learning: Data. Modules and Training Stages, 2023.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A. Rusu, Joel Veness, Marc G. Belle- mare, Alex Graves, Martin Riedmiller, Andreas K. Fidjeland, Georg Ostrovski, Stig Petersen, Charles Beattie, Amir Sadik, Ioannis Antonoglou, Helen King, Dharshan Kumaran, Daan Wier- stra, Shane Legg, and Demis Hassabis. Human-level control through deep reinforcement learning. Nature, 518(7540):529-5333 February 2015a.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Belle- mare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 51 1(77540) 2015b.
Evgeni Nikishin, Max Schwarzer, Pierluca D' D'Oro, Pierre-Luc Bacon, and Aaron Courville. The primacy bias in deep reinforeement learning. In International Conference on Machine Learning, pp. 16828-16847. PMLRR, 2022.
13

---

Published as a conference paper at ICLR 2025
Johan Samir Obando Ceron, Aaron Courville, and Pablo Samuel Castro. In value-based deep rein- forcement learning, a pruned network is a good network. In Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 38495-38519, PMLR, 21-27 Jul 2024.
Johan Samir Obando Ceron*, Ghada Sokar*, Timon Willi*, Clare Lyle, Jesse Farebrother, Jakob Nicolaus Foerster, Gintare Karolina Dziugaite, Doina Precup, and Pablo Samuel Cas- tro. Mixtures of experts unlock parameter scaling for deep RL. Forty-first International Conference on Machine Learning, 2024. URL https: http///pppeeeeeeeeeeeeeeeeeeeeee net htt/fooummmee X9VMhfFxWI.
Travis E. Oliphant. Python for scientific computing. Computing in Science & Engineering, 9(3): 10-20, 2007. doi: 10.1109/MCSE2200..555
Svetlana Pavlitskaya, Christian Hubschneider, Michael Weber, Ruby Moritz, Fabian Huger, Peter Schlicht, and Marius Zollner. Using mixture of expert models to gain insights into semantic segmentation. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops, pp. 342-343, 2020.
Joan Puigcerver, Carlos Riquelme Ruiz, Basil Mustafa, Cedric Renggli, André Susano Pinto, Syl- vain Gelly, Daniel Keysers, and Neil Houlsby. Scalable transfer learning with expert models. In International Conference on Learning Representations, 2020.
Joan Puigcerver, Carlos Riquelme Ruiz, Basil Mustafa, and Neil Houlsby. From sparse to soft mixtures of experts. In The Twelfth International Conference on Learning Representations, 2024.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. CoRR, 0ts//77700634777 2017.
Max Schwarzer*, Johan Samir Obando Ceron*, Aaron Courville, Mare G Bellemare, Rishabh Agar- wal, and Pablo Samuel Castro. Bigger, better, faster: Human-level Atari with human-level effi- ciency. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference Machine Learn- ing, volume 202 of Proceedings of Machine Learning Research, pp. 30365-30380. PMLR, 23-29 Jul 2023.
Noam Shazeer, *Azalia Mirhoseini, *Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hin- ton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of- experts layer. In International Conference on Learning Representations, 2017. URL https: / //opereeview net wwtr:m/rmmmmmmmmmmm d-BldlBBBle Bioommmmi-
Ghada Sokar, Elena Mocanu, Decehal Constantin Mocanu, Mykola Pechenizkiy, and Peter Stone. Dynamic sparse training for deep reinforcement learning. In International Joint Conference on Artificial Intelligence, 2022.
Ghada Sokar, Rishabh Agarwal, Pablo Samuel Castro, and Utku Evei. The dormant neuron phe- nomenon deep reinforcement learning. International Conference on Machine Learning, pp- 32145-32168. PMLR, 2023.
Richard S. Sutton and Andrew G. Barto. Introduction to Reinforcement Learning. MIT Press, Cambridge, MA, USA, 1st edition, 1998. ISBN 0262193981.
Adrien Ali Taiga, Rishabh Agarwal, Jesse Farebrother, Aaron Courville, and Marc G Bellemare. Investigating multi-task pretraining and generalization in reinforcement learning. The Eleventh International Conference on Learning Representations, 2022.
Hado Van Hasselt, Matteo Hessel, and John Aslanides. When to use parametric models in rein-
forcement learning? Advances in Neural Information Processing Systems, 32, 2019.
Guido Van Rossum and Fred L Drake Jr. Python reference manual. Centrum voor Wiskunde en Informatica Amsterdam, 1995.
14

---

Published as a conference paper at ICLR 2025
Xin Wang, Fisher Yu, Lisa Dunlap. Yi-An Ma, Ruth Wang, Azalia Mirhoseini, Trevor Darrell, and Joseph E Gonzalez. Deep mixture of experts via shallow embedding. Uncertainty in artificial intelligence, pp. 552-562. PMLR, 2020.
Timon Willi*, Johan Samir Obando Ceron*, Jakob Nicolaus Foerster, Gintare Karolina Dziugaite, and Pablo Samuel Castro. Mixture of experts in mixture of RL settings. In Reinforcement Learn- ing Conference, 2024. URL https: /pwpreevee.eeeeeeeeeeeeee. wwt.:neeeteeeeeeeeeeee hett/oorum22eeSPPPRRL..
Brandon Yang, Gabriel Bender, Quoe V Le, and Jiquan Ngiam. Condconv: Conditionally parame- terized convolutions for efficient inference. Advances in neural information processing systems, 32, 2019.
Hanrong Ye and Dan Xu. Taskexpert: Dynamically assembling multi-task representations with memorial of In Proceedings of the IEEE/CVF International Conference on Com- puler Vision, pp. 21828-21837, 2023.
Yanqi Zhou, Tao Lei, Hanxiao Liu, Nan Du, Yanping Huang, Vincent Zhao, Andrew M Dai, Quoc V Le, James Laudon, et al. of with expert choice routing. Advances in Neural Information Processing Systems, 35:7103-71144 2022.
15

---

Published as a conference paper at ICLR 2025
A EXPERIMENTAL DETAILS
Tasks We evaluate Rainbow on 20 games from the Arcade Learning Environment (Bellemare al., 2013). We use the same set of games evaluated by Obando Ceron al. (2024) for di- rect comparison. This includes: Asterix, SpaceInvaders, Breakout, Pong, Qbert, DemonAttack, Seaquest, WizardOfWor, RoadRunner, BeamRider, Frostbite, CrazyClimber, Assault, Krull, Box- ing, Jamesbond, Kangaroo, UpNDown, Gopher, and Hero. In addition, we run some of our main results on all suite of 60 games.
On Atari 100k benchmark, we evaluate on the full 26 games. These are: Alien, Amidar, Assault, As- terix, BankHeist, BattleZone, Boxing, Breakout, ChopperCommand, CrazyClimber, DemonAttack, Freeway, Frostbite, Gopher, Hero, Jamesbond, Kangaroo, Krull, KungFuMaster, MsPacman, Pong, PrivateEye, Qbert, RoadRunner, Seaquest, UpNDown.
We the default hyper-parameter values for DQN (Mnih et al., 2015a), Rainbow (Hessel et al., 2018), and DER (Van Hasselt et al., 2019) agents. We share the details of these values in Table 1.
Hyper-parameters
Table 1: Default pareee-prrmeters setting for DQN, Rainbow and DER agents.
Atari
Hyper-parameter
DQN
Rainbow
DER
Adam's (e)
1.5e-4 6.25e-5
1.5e-4
0.00015
Adam's learning rate
6.25e-5
0.0001
Batch Size
32
32
Conv. Activation Function
ReLU
ReLU
ReLU
Convolutional Width
Dense Activation Function
ReLU
ReLU
ReLU
Dense Width
512
512
512
Normalization
None
None
None
Discount Factor
0.99
0.99
0.99
Exploration
0.01
0.01
0.01
Exploration decay
250000
250000
2000
Minimum Replay History
20000
20000
1600
Number of Atoms
0
51
51
Number of Convolutional Layers
3
Number of Dense Layers
2
2
2
Replay Capacity
1000000
1000000
1000000
Reward Clipping
True
True
True
Update Horizon
10
Update Period
4
1
Weight Decay
0
Sticky Actions
True
True
False
B ADDITIONAL EXPERIMENTS
B. COMPARISON BETWEEN TOKEN CHOICE ROUTING AND EXPERT CHOICE ROUTING
We focus our analysis on expert choice routing, where experts select tokens, as opposed to token choice routing, where tokens select experts. We compare these two routing strategies across different numbers of experts using Rainbow agent with a ResNet architecture.
Figure 12 demonstrates that expert choice routing consistently outperforms token choice routing across varying number of experts. Moreover, this performance gap widens as the number of experts increases, highlighting the advantages of expert choice routing addressing the limitations of token choice routing.
16

---

Published as a conference paper at ICLR 2025
Tiken Cholce-2 Export Choice 2
Tiken Cholce-4 Expert Choiie
T0M Ruman No ma
IOM Human No mr
Chaice-8 Experi Cholcc-8
Tilen Chaice-16 Expert Cholco
Normal
Number
Number Frames
Figure 12: A comparison between token choice routing and expert choice routing on Rainbow with ResNet encoder. Expert choice routing consistently outperforms token choice routing across different number of experts. We report interquantile mean performance with error bars indicating 95% confidence intervals, over 20 games with seeds each.
B.2 DON'T FLATTEN TOKENIZE
In this section, we provide further analysis on the impact tokenization and their weight com- bination on models with varying number of experts. To this end, investigate SoftMoE-2 and SoftMoE-8 and compare it against the scaled SoftMoB-1 and ExperiChoice-1 (non-combined to- kenization). Figure 13 shows that single scaled expert can reach performance close multiple experts, even in the case of large number of experts. In addition, consistent with our results, com- bined tokenization adds benefits over the sparse one. Figure shows a comparison between the learning curves of single scaled expert and multiple ones.
SoftMEE 2
SoftMoE-1 x2)
Expert Choice-1 (scaled x2)
Baseline (scaled x2)
3.0
3.5
5.5
7.0
7.5
Interquartile Mean (IQM)
SeftMoE-1 (scaled x8)
Expert Choice-1 (scaled
Baseline (scaled x8)
3.0
3.5
4.0
4.5
5.0
5.5
6.0
6.5
7.0
7.5
Interquartile Mean (IQM)
Figure 13:
Combined tokenization plays crucial role in performance. hidden layer achieves a comparable performance to the standard SoftMoE.
SoftMoE-1 with scaled
17

---

Published as a conference paper at ICLR 2025
SEtHOGL Ro Cseine Toaet
Sxpert Cheen Drvatd
Expelt thaced boaidd
25
55
125
156
206
23
50
73
100
225
25
53
160
125 150
315
Figure 14: Comparison between the learning curves of the scaled SoftMoE-1 and multiple experts. We report interquantile mean performance with error bars indicating 95% confidence intervals, over 20 games with 5 seeds each.
To further assess the effect of flattening in reducing performance of agents, replace the flattening operation the baseline by global average pooling. We performed these experiments on the baseline with the default dimensionality and the scaled one in Rainbow.
Median
IQM
Mean
Optimality Gap
Baseline (scaled x4) wr Avg Pooling
Baseline Iscaled x4] wi Fatttening
Baseline w Avg Popling-
Baseline Flattening
4.5
6.0
7.5
Tuumnni Normelized Scafe
6
10
0.030 0.045
0.060
Figure 15: Replacing the flattening operation by global average pooling improves the performance of the baseline and enables scaling.
B.3 ANALYSIS OF SOFTMOE COMPONENTS
In this section, present analysis of each of SoftMoF components the DER algorithm. As shown in Figure 16, we observe similar tends to our findings in Section 4.2 on Rainbow. Combined tokens offers an additional benefits over the sparse ones. Experts do seem not be specialized in subset of the tokens. Adjusting the architectural depth and width for direct comparison with the baseline doesn't explain the performance gains of SoftMoE.
M 4
Baseline (scaled x4)
Combined tokenization
ExpertChoice-4
Expert speciallzaatinn
SoftMoE-d (Expert's #slets #tokens)
ExpertChoicee (Expert's #slets #takens)
[iii] Width
SoftMoE-4 (scaled x4)
[iii] Depth
Baseline (scaled x4) with extra
0.0
1.0
1.5
2.0
Interquartile Mean (IQM)
Figure 16: Understanding the impact of SoftMoE components on DER.
Effeet of Expert width In section 5, we studied the effect of scaling up the size of the expert hidden layer of So Here, we present the results of scaling down the size of this layer by 4. As shown in Figure 17, SoftMoE-1 with scaled down layer still outperform the baseline despite that has less parameter count.
B.4 MORE AGENTS
Same as our finding for DQN using SoftMoE-4, we observe that single expert has closer perfor- mance to the scaled baseline than SoftMoE-8, as shown in Figure 18. Further investigation needed to fully understand the behavior of SoftMoE on DQN.
B.5 EXPERT UTILIZATION
In this section we provide more analysis to understand the utilization of experts in multi expert setting and the effect of existing techniques improving plasticity.
18

---

Published as a conference paper at ICLR 2025
JOH Human Norma tred Seore
Bassline
SoltMot-1 (xcaled sd]
SoltMob-2 (scated
25
50
75
100
125
150
175
200
Number of Frames iin milliosl
Figure 17: The effect of scaling down the size the expert hidden layer in 1-ftMoE-1.
Median
IQM
Mean
Optimality Gap
ExpertChoice-1 (scaled Baseline (scaled
SoftMoE-8
2.5
33
4,0
7.30.06
0.08
Human Normalized Score
Figure 18: Evaluating more agents. When combined with DQN, SoftMoE with single scaled expert does not provide large gains over the baseline.
B.5.1 SOME EXPERTS ARE REDUNDANT
In Section 4.2, we showed that experts are not specialized learning subset tokens. We expand this analysis by assessing the redundancy of experts. To this end, prune experts of SoftMoE-4 during training. We studied two pruning schemes: pruning the two experts once and gradually prune throughout training. As shown in Figure 19, pruning these experts does not change the performance of the agent.
6
JOM Human Nomalizd
4
2
No pruning
Prune one 25M, then S0M
75
Number of Frames (in millions)
Figure 19: Pruning subset of experts in SoftMoE-4 does not lead to performance drop, suggesting that some the experts are redundant. We report interquantile mean performance with error bars indicating 95% confidence intervals, over 20 games with 5 seeds each.
B.5.2 IMPROVING EXPERT UTILIZATION
Experimental details For hyper-parameter tuning, we used six games (Asterix, Demon Attack, Seaquest, Breakout, Beam Rider, Space Invaders). For the reset period, we searched for Rainbow over the grid [0.5 10a 10 10 50 x 10a, 100 x 10 125 x 10a, 500 10a, 1000 10*] gradient steps. The best found is 125 104, which corresponds to 20M environment steps. For S&P, we searched over the values [00.5, 0.5), (0.4, 0.1), (0.8, 0.2)] for shrink and perturb hyperparameters respectively, studied in (Schwarzer* et al., 2023). We used values of (0.5, 0.5).
Reset experts and router weights In investigating current techniques to improve expert utiliza- tion, we explore reset either experts weights only or both experts and router weights. As illustrated in Figure 20, both methods do not help in improving utilization in Rainbow and DQN, with resetting the router weights leads to worse performance.
19

---

Published as a conference paper at ICLR 2025
Optimality Gap
Reset-4 (experts router) Reset-2 (experts router)
Reset-4 (experts) Reset-2 (experts) SoftMoE 4
7
8
6.4
6.8
7.2
9.6
10.2
10.8
11.4
0.0210.0240 0.210.024.0077.0000
Human Normalized Score
Median
1QM
Mean
Optimality Gap
Reset-4 (experts router)
Reset 2 lexperts router)
Reset-4 (experts)
(experts)
SoftMoE-4
2.75
3.00
2.75
5.5
Human Normalized Score
Median
IQM
Mean
Optimallty Ity Gap
Reset-4 (experts router)
Reset-2 (experts router)
Reset-A (experts)
Reset-2 lexperts)
SoftMoE 4
2,00
2.25 2.50
2.75
2,00
2.25
2.50
2.75
2.4
2.8
3.2
0.12
0.15
0.18
0.21
Human Normalized Score
Figure 20: Comparison between resetting experts weights only versus experts and router weights in Rainbow (top), DQN (middle), and DER (bottom).
C EXTRA ENVIRONMENTS
To evaluate the generality of our claims, we conducted experiments with Rainbow on Procgen (Cobbe et al., 2019) (Figure 21) and with SAC (Haarnoja et al., 2018) on CALE (Farebrother & Castro, 2024).
Median
IQM
Mean
Optimality Gap
SoftMoE-1 (x4)
SoftMoE-4
Baseline (x4)
Baseline
0.16
0.24
0.32
0.40
0.20
0.25
0.30
0.35
0,28
0.32
0.36
0.64
0.68
0.72
Human Normalized Score
Figure 21: Evaluating Rainbow on the 16 easy Procgen (Cobbe et al., 2019) environments for SoftMoE-4 and SoftMoE-1 with the penultimate layer scaled by 4x. Normalization using the Rmin and Rman scores from Cobbe et al. (2019). Reporting Median, IQM, Mean, and Optimality Gap (Agarwal al., 2021), where higher is better for the first three. Consistent with the results in the paper, SoftMoe yields improvements, and the scaled SoftMoE-1 1 is comparable to SoftMoE.4.
done
Median
IQM
Mean
Optimality Gap
Tokenized (4x)- Tokenized Baseline (4x) Baseline-
0.30
0.35 0.40
0.45
0.42
0.48
0.54 0.60
1.04
1.12 1.20
0.510
0.525
0.540
Human Normalized Score
Figure 22: Evaluating tokenization on the SAC encoder (Haarnoja et al., 2018) on the CALE (Farebrother & Castro, 2024), done in Figure 5. In these experiments, summed the first dimension after tokenization. Reporting Median, IQM, Mean, and Optimality Gap (Agarwal et al., 2021), where higher is better for the first three. Consistent with the results in the paper, tokenization yields improvements over the baseline.
20

---

Published as a conference paper at ICLR 2025
C.I PER-GAME RESULTS
C.I.1 DON'T FLATTEN, TOKENIZE!
Assault
Asterix
BeamRider
400
3503
220000
Taannodiing Nec
voudomm
Raturn
3.0
233000
tha)
100c
Coperiiar
Guper
Gopetment
200
101
2006
Soettine
und daa (it)
-20
Beser dio
Toa vandiaaa H6C
100
359
200
389
150
20c
26
200
12c
Breakout
CrazyClimbeer
DemonAttack
Frostbite
180000
4CE05
Toonanding
p40000
Thernzeding
cotiiii the
22000
Tou oite dhin the
Retum
3sall lin Hild
Eanil 2c0
Joan maaltiaa tua
L13
4000
630
L13
40000
Maarbsckig Hed
15:0
1000
59-
5600
6:0
300
Ite
3ha
303
1:00
11/e
Gopher
Jamesbond
Kangaroo
1480
12303
50000
souettyy Ltxt
1230
13:03
Return
43000
Raccie
Exsii
Experiaat
D partmar
4.0
tho sumbeey der
Tree peeeenntien dic
10000
20C
Kanik
Base
3-
116
300
351
90
305
100
tie
Kruil
Pong
Qbert
RoadRunner
3.03
10.04
52.00
2.03
10-
43:0
4303
5.0
10:04
25:0
35303
Prperimed
-10
Skanbacsun
glig Hoi
3333
2.03
Revdiin 132
Buoline [1x)
Rese seclao
3.80
100 Seaquest
224
280
359
200
20c
100
i2c
200
Spacelrvaders
UpNDown
WizardorWor
thig ttal
sonding
Tabon liveltig (1x)
Tooessddove 120
Toannedioo 12d
100
Return
Toor thether 110
Base (ao
Stop
Stop
Stop
Stop
Figure 23: Per-game results for tokenized Rainbow-lite with CNN network (corresponding to Fig- ure 5).
21

---

Published as a conference paper at ICLR 2025
Assault
Astentx
BeamRider
Boxing
182
voading No
Retumn
2w3
Teomstnding 2rd Thomsnnderp Mc
Dase by deo
Stanting Hid
Edto (eo)
Breakout
CrazyClimbber
DemonAttack
Frostbite
150000
22303
12000
120000
Rotur
Troem toe shon Das
Taker aencian
socingg (hc
Rasti
Band
400-
Toansndiirg 2id
Tow soeakig HCC
X000
Toansoddiio H6d
2000
230
Gopher
Hero
Jamesbond
Kangaroo
Eqgatrees
2506
12:03
Taker aam
Relum
1X0
simm
1568
Axperieee coar
Kxperin and Tou semeeoddee Sec
wondeng See
Marbbreck ded
Bese d4a0
230
1790
Kruil
Pong
Obert
RoadRunner
1583
12.04
Saparree Maorbeckig Red
10004
13333
Deeeinedigg deg Hoo
Bate den life
sondhar
Suetine
100
202
230
200
150
26
100
Lic
200
Seaquest
Spacelrvaders
UpNDown
WizardOrWor
Eipattmen:
3403
23000
Tola voadiom 12-
Taker beddan Lao
3000
Retumn
11000
te Dhe
22303
Expermert
Exeerment
12000
1393
Sin sondday del
Toantodium 16d
5000
Toonsediia (aa)
Tow wondding InC
1nn Stop
1h6
380
174 Step
V0
105
190 Stop
Stop
Figure 24: Per-game results for tokenized Rainbow- with Impala network (corresponding to Figure 5).
22


Reviews:
[{"review_id": "IUtIhtXtvw", "reviewer": "Reviewer_3aV6", "summary": "The paper extensively studies the degree of contributions by each factor of SoftMoEs, a method mitigating the inverse proportionality between the performance and the architecture size of the online value-based deep reinforcement learning (DRL) methods. Despite the effectiveness of the approach, which component of SoftMoEs drives the improvement remains a mystery. The paper discovered that tokenization, the scheme of converting extracted features while maintaining its spatial structure, plays a significant role in SoftMoEs, bestowing the ability for the algorithm with a large network to perform well. From this insight, the paper further extends the argument to the redundancy of the experts, claiming that a single expert performs competitively against the multi-expert under certain conditions.", "questions": "- How does the result depicted in Figure 5 reduce to the claim: “providing strong evidence that tokenization plays a major role in the efficacy of SoftMoE” (line 295)? \n- As mentioned in the weaknesses section, some empirically supported claims are skeptical due to the lack of information about the default configurations of SoftMoEs. Would the authors be able to list the precise configurations? \n- Is a hyperparameter search for random resets and S&P conducted in section 6?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "The paper tackles the important problem: *which factors of SoftMoEs benefit the most in mitigating the performance degradation of DRL algorithms with the increase of network size?* The problem is significantly important as it falls into the category of questions that ask the fundamental mechanism of the algorithms. Questions asking the fundamental mechanism of “why” the existing approach is effective are often overlooked in the community but also broadly impact the community in multiple ways (provides the common ground for the future methodologies, encourages the community to rethink the current approaches, etc.). Successfully answering the fundamental question thus turns this paper into an invaluable work with a potentially significant impact on the community.\n\nThe main finding of the paper is that tokenization of the extracted features of observations majorly contributes to performance improvements. This argument has been drawn and supported by the series of experiments along with the rigorous empirical analysis (adoption of IQM, 95% confidence interval over the stratified bootstrap samples), which reinforces the legitimacy of the claim.", "weaknesses": "- The significance of tokenization is limited to particular settings and does not apply robustly over the value-based online DRL. In fact, Figure 10 depicts that the scaled SoftMoE-1 resembles a similar performance/optimality gap from that of the scaled Baseline on DQN, indicating the minor, near-zero effect on the performance improvement by the tokenization. \n- The paper lacks some crucial details on the experimental settings of SoftMoEs. One of the prominent ones is the number of slots $p$. Unknown $p$ value reduces the confidence in the conclusion drawn in paragraph **Expert specialization** (line 231). Here, the paper claims that the specialization of experts is not the primary factor contributing to the high performance by increasing the number of $p$ to the number of tokens. However, the analysis might not be valid if the default value of $p$ is close to the number of tokens, and there is no way the readers can notice this unless specified in the paper. The paper would benefit from clarifying the actual values of the default configurations of SoftMoEs. \n- Some ablation studies omit important explanations, making it hard to follow the arguments. For instance, in the paragraph **Tokenization baseline**, the paper replaces the feature flattening operation of the Baseline with a tokenize-and-aggregate (either by average or sum) operation. The results suggest that tokenization with scaled representation significantly improves the performance of the Baseline, Rainbow-lite architecture. However, it is hard to connect this finding to the claim that “tokenization plays a major role in the efficacy of SoftMoE” (line 295). An additional explanation bridging similar logical gaps in the ablation studies (especially from the results to the final statement) would gradually improve the clarity of the paper. \n- The effort towards the hyperparameter sweeps is not explicitly mentioned in the paper. In empirical studies, hyperparameter sweeps are necessary for fair comparison, especially when the purpose of comparison is to determine the effectiveness of approaches. Omitting this step weakens the arguments made in section 6, mitigating experts’ redundancy by parameter reset and S&P. \n- Although mentioning how to improve expert utilization (section 6) marks an interesting future direction, it also feels unnecessary. This is mainly due to the open-ended analysis and the fact that the argument is slightly out of the main focus of the paper. The consistency and logical flow of the paper would be improved by moving section 6 to the appendix. \n\nIn addition to these points, some minor formatting errors and ambiguous presentations caught my attention:\n\n- Some references miss the year of publication. For instance, line 469 contains two works cited without the publication years. \n- Inconsistent citation formats. Capitalization of titles, format of venues, etc. \n- Some figures lack explanations. Specifically, the target quantity is unclear in the bar plots: Figure 1, 7, and 9. While it is mentioned that a human-normalized score is the target of measurements in the main text (section 4.1), it would be convenient to indicate within the figure. Also, since other figures indicate human-normalized scores, clarifying the target quantity in all the figures would further improve consistency.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "dIhEvpbVUP", "reviewer": "Reviewer_9z4X", "summary": "This paper investigates the key factors that make SoFeMoE effective in visual reinforcement learning. Through comparisons, it demonstrates that tokenization and the combination of token weights, rather than simply using multiple experts, drive the performance improvements. The conclusions are based on experiments conducted in the Arcade Learning Environment, covering 60 games and using 5 seeds.", "questions": "Please refer to the weakness section.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is well-written and easy to read.\n - The figures effectively illustrate the different analyses and help understand the main results.\n - The analysis is solid and well presented.", "weaknesses": "The main concern is the generality of the claim.\n - The experiments are conducted in one simulation platform with all discrete actions. It is unclear whether the same observation is universal applicable to other simulation platforms, especially involving agents with continuous states and actions.\n - All the results are conducted with 4 experts. Would the conclusions be applicable to different numbers of experts?\n - In Figure 6, it will be beneficial to visualize the performance with one unscaled expert to understand the results better.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "7rO2rPjQMA", "reviewer": "Reviewer_gzLL", "summary": "This paper analyzes the effect of different components in the soft-of-mixture of experts (SoftMoEs) in online RL, the goal is to understand the key factors and design decisions that influence/derive the performance of (SoftMoEs), the analysis shows that tokenizing the output of the cnn encoder has the biggest effect on the performance, even when using a single expert or the baseline model, tokenizing the encoder output results.", "questions": "- In line 361, computational efficiency by combined tokens, I do not understand the point of the plot, is it to show that using fewer slots (which means better time complexity) still results in good performance? Can you add a plot that directly shows the relationship between the number of slots and time complexity?\n\n- The authors argue that tokenizing the encoder output preserves the spatial information unlike flattening, can the authors run a baseline where the cnn encoder output is actually a vector? This can be done by adding a global average pooling layer after the last conv layer, which will reduce the spatial dimension to one.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- Results are significant with small variance over the 60 Atari games. \n\n- The effect of tokenization seems to transfer between architectures, which might suggest that using tokenization would always be helpful, at least in the algorithms used in the paper (Rainbow and DER).\n\n- Showing that we are underutilizing the mixture of experts in (SoftMoEs) is an important implication of this paper, which encourages more research in this area.", "weaknesses": "- Authors should add a section that explains what a token is and what a slot it is, this is not defined in the paper and it will make the paper much more clearer. \n\n- The effect of tokenization in a single expert does not seem to transfer to DQN, which suggests there is something missing in the analysis, the authors suggested that it might be related to the categorical loss used in Rainbow and DER, but there is no further investigation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "d0ANu2sJ7O", "reviewer": "Reviewer_8sGc", "summary": "Obando-Ceron* et al. (2024) [1] demonstrated that SoftMoEs are effective architectures for scaling models in online RL. However, their paper did not explain the reasons behind this performance gain at scale. This submission analyzes factors that might contribute to the effectiveness of SoftMoEs at scale by ablating different components of the SoftMoE architecture within the Rainbow, DER, and DQN baselines on the Atari benchmark.\n\nOverall, the authors show that tokenizing the encoder output, rather than using multiple experts, is the primary factor driving SoftMoE’s effectiveness. They also demonstrate that a single scaled expert with tokenization can match the performance of multiple experts.\n\n[1]: Ceron, Johan Samir Obando, et al. \"Mixtures of Experts Unlock Parameter Scaling for Deep RL.\" Forty-first International Conference on Machine Learning.", "questions": "I have some questions based on the current state of submission:\n\n- For all trends observed in Section 4.2, do the authors anticipate that similar trends would hold for DER? Could the hypotheses be verified on an additional algorithm to confirm that the design choices assessed in Section 4.2 are generally applicable and not specific to one algorithm i.e. Rainbow?\n- In Figure 5, did the authors observe a clear scaling trend from 1x to 2x, 4x, and 8x using the tokenized_sum baseline with either CNN or IMPALA? Also, in Figure 5, the tokenized scheme was [h\\*w, d], which I believe corresponds to PerConv tokenization. Would similar trends hold if [d, h\\*w] (i.e., PerFeat tokenization, similar to Figure 7) were used instead? This would help confirm whether tokenization generally improves performance, even if PerConv outperforms PerFeat, as long as both do better than simple flattening on the baseline.\n- In Figure 8, when selecting 10% of the slots, are these slots chosen randomly, or is a specific heuristic used for pruning? Additionally, are these 10% slots fixed throughout training, or do they change randomly at each training iteration?\n- Why was the analysis in Section 6 (Figure 11) not conducted for DQN?\n\n\n**Typos:**\n\nPage 3 Line 134: propose —> proposed\n\nPage 8 Line 385: performance on all the 60 games —> performance on all **of** the 60 games", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The paper is well-written, easy to follow, and well-motivated.\n- The authors conducted numerous empirical ablations on 3 different algorithms—primarily on Rainbow, but also on DER and DQN—using the popular Atari benchmark.\n- I appreciate that the authors reported the IQM and Optimality Gap on 5 seeds for statistical significance.", "weaknesses": "While this paper provides an in-depth exploration of tokenization in SoftMoEs, there are areas where it could be improved to make it an accept.\n\n1. The paper’s scope feels somewhat narrow, as it primarily focuses on *why* SoftMoEs work well when scaled up on the Atari benchmark using the Rainbow, DER, or DQN algorithms alone. As a result, it remains unclear whether SoftMoEs would be effective in more challenging OOD generalization benchmarks, such as Procgen, which present greater difficulty than the IID singleton environments like Atari.\n2. Additionally, as shown in Figure 10, the SoftMoE-1 (scaled 4x) baseline performs significantly worse in the DQN setting. It would be beneficial for the authors to test their scaled-single-expert approach on another algorithm, potentially an actor-critic algorithm like PPO, in a pixel-based environment like Procgen.\n3. Given the popularity of DQN, further investigation into the differences in DQN’s behavior compared to the other two algorithms would also add more value to this submission.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "NleckLjVew", "reviewer": "Reviewer_XvpH", "summary": "MoE (Mixture of Experts) has demonstrated significant potential. The latest research on MoE, exemplified by softMoE, attributes the performance advantages of these algorithms to structural sparsity. However, the authors of this paper propose a contrasting view. They observe that even a single-expert version of softMoE can perform quite well. Additionally, they find no evidence of expert specialization occurring in softMoE.\n\nThrough ablation studies, the authors discover that \"tokenization\" of the output from the convolutional encoder plays a crucial role. As a result, they suggest that tokenization is a highly important operation.\n\nThe authors investigate the importance of different components of softMoE in Atari game experiments, including: (i) the use of a learnable tensor Φ for obtaining dispatch (ΦD) and combine (ΦC) weights; (ii) processing p input slots per expert; (iii) architectural dimensions (network depth and width); (iv) the use of n experts; and (v) the tokenization of the encoder output.\n\nA key finding is that tokenization significantly enhances performance. By applying only tokenization to the baseline, a notable performance improvement is observed, demonstrating that tokenization is a primary contributor to the effectiveness of softMoE. Furthermore, even a single-expert softMoE achieves relatively strong results.\n\nFinally, the authors explore methods to improve the performance of multi-expert MoE models. They attempt reset and S&P methods, testing them with Rainbow and DER algorithms, but do not find any approach that consistently enhances multi-expert MoE performance.", "questions": "1) The authors mentioned that \"they demonstrated that a performance drop is not observed when scaling down the dimensionality of each expert.\" Did you experiment with a setup where the dimensionality is reduced by 4 times in the SoftMoE model? Specifically, the experimental group would use SoftMoE with 4 experts, each scaled down by a factor of 4 to maintain the same overall model size as the baseline. The control group is baseline, which would not be scaled. This setup might highlight the importance of the number of experts in SoftMoE if the baseline performance remains unaffected.\n\n2) Could the authors provided training curves for each individual Atari game? Comparing the performance of the Unscaled baseline and Unscaled tokenize models on each task.\n\n3) Did the authors conduct experiments on environments like DeepMind Control (DMC) or Meta-World to thoroughly demonstrate the generalization capabilities of the tokenize?\n\n4) Why is \"Network Plasticity in RL\" discussed in the related work section? Which part of the paper is relevant to plasticity? Does it correspond to Section 6? If so, the experimental results presented in Section 6 do not seem to lead to any valuable conclusions.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 4, "contribution": 2, "strengths": "The authors clearly present three different classical methods for processing encoder features: Flatten (baseline), SoftMoE architecture, and the routing mechanism.", "weaknesses": "1) The paper has a significant issue: the baseline scaled by *4 used throughout the experiments appears to perform worse, potentially due to suboptimal parameter. The use of the *4 baseline is problematic, as it may unfairly weaken the baseline, thus exaggerating the benefits of tokenization. \n\nThe authors could run experiments comparing unscaled baselines to SoftMoE models with equivalent parameter counts. Specifically, It is crucial to compare the original-sized baseline with the SoftMoE where experts are scaled down by four times, and also compare it with a single expert SoftMoE scaled down by four times. \n\nWhen examining Figure 5, comparisons between unscaled results reveal that both the tokenized sum and tokenized average actually perform worse. This suggests that the effectiveness of tokenization depends on a balance between network structure, task complexity, and parameter settings, rather than indicating that tokenization is inherently beneficial. \n\n2) The paper also mentions that a single-expert SoftMoE performs well, which suggests that multiple experts may not be critical, and that tokenization is the key factor. However, Figure 10 indicates that combining a single expert with tokenization in SoftMoE does not improve performance in DQN. This contradict the claims that a single expert is effective. This indicates that the paper’s main assertions are not solid. Could the authors address this discrepancy and provide additional analysis or experiments to clarify the effectiveness of single-expert SoftMoE across different algorithms？\n\nThe truth might be that the authors' DQN setup is relatively underperforming, making the multiple-expert configuration beneficial to compensate the weakness of single expert. In contrast, stronger algorithms perform well enough with a single expert, making multiple experts redundant. Testing on more complex tasks might reveal a clearer advantage of using multiple experts over a single one. Therefore, a more reasonable interpretation would be that both tokenization and multiple experts contribute to SoftMoE’s performance, aligning with the fundamental concept of MoE rather than overemphasizing the importance of tokenization.\n\n3) In Section 6, the authors aim to develop new techniques to improve expert utilization and maximize the benefits of MoE architectures in RL. However, only two tricks are listed, and the number of baselines compared is too limited. The paper fails to present methods that consistently improve multi-expert MoE performance. If the goal is to provide which trick may improve certain algorithm, a broader set of techniques should be explored.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Spotlight)