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

We thank the reviewer for their feedback, useful comments, and address their concerns below.

We appreciate the opportunity to provide a clarification regarding the reviewer’s summary. Our work does not aim to provide “a contrasting view” of recent MoE research, nor are we trying to suggest that a single expert suffices. Rather, our analyses reveal that multiple experts are not being fully utilized in the specific context of online deep RL settings studied in [1]. This finding has an important implication on encouraging further research to increase their usage for further performance improvement as discussed in Section 1 and Section 8 and also acknowledged by Reviewers gzLL and 3aV6. With an awareness of the potential benefits of multiple experts, our work has already started to explore techniques from existing literature to enhance network utilization, as presented in Section 6, though this is not the paper's primary focus.

W1: 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.

A: Given that our work stems from the ideas explored in [1], we used the same scaled baselines used there. In contrast to [1], we studied the effect of the upscaling of each expert and found that agents do not experience the performance drop observed in the scaled baseline. Further details can be found in Section 4.2.

W2: Could the authors address this discrepancy and provide additional analysis or experiments to clarify the effectiveness of single-expert SoftMoE across different algorithms?

A: Indeed, as discussed in section 5.1 (and in [1]), DQN appears to benefit less from the use of SoftMoEs in general, which may explain why SoftMoE-1 yields little gain. We hypothesize this may be due DQN’s use of regression versus Rainbow’s classification (C51) loss; we are currently running experiments with SoftMoE-1 (x4) with the C51 loss and will report back here once they have made more progress.

W2: a more reasonable interpretation would be that both tokenization and multiple experts contribute to SoftMoE’s performance

A: Certainly, multiple experts do contribute to SoftMoE’s performance, as demonstrated by the higher performance of multiple experts compared to a single scaled expert in our figures. However, the relatively small performance difference between these two, compared to the larger difference between the scaled baseline and the single scaled expert, suggests that multiple experts are not the primary driver of the observed performance improvement; more importantly, it suggests that we are under-utilizing multiple experts, as we discussed in Section 6.

W3: In Section 6, the authors aim to develop new techniques to improve expert utilization. However, only two tricks are listed, and the number of baselines compared is too limited. If the goal is to provide which trick may improve certain algorithm, a broader set of techniques should be explored.

A: As we clarified, the paper's main focus is to understand the reasons behind the observed efficiency of SoftMoEs in deep RL. In Section 6, we take the first step towards a promising future direction for improving expert utilization by studying existing techniques.

Thank the authors for answering each of my questions in detail. However, my main concerns remain unresolved. Specifically:

1. Baseline unscaled consistently outperforms Baseline scaled ×4 in multiple metrics (as shown in Fig. 5: CNN and Impala's median, IQM, and mean). Therefore, I believe it is reasonable and necessary to compare against Baseline unscaled.

2. If we compare Baseline unscaled with tokenized sum/avg unscaled, the conclusions of the paper would not hold. The results suggest that tokenization is not the primary driver of performance improvement. (Refer to Fig 5: CNN and Impala's median, IQM, Mean; Appendix B.6.1 Assault, BeamRider, Boxing, CrazyClimber, Frostbite,Gopher,Hero...) In 25 environment of total 40 environment, Baseline unscaled performs better then tockenized sum unscaled and tockenized avg unscaled in the same time, which means Tokenization is not the primary driver of the performance improvement.

Thank you for your reply!

As previously mentioned, [1] proposes SoftMoEs that enables scaling RL networks with performance increases. Our primary objective in this paper was to investigate why SoftMoE is so effective at enabling this scaling, which is why most of our experiments have focused on the scaled baseline.

However, we do agree that it is useful to evaluate the impact of tokenization on unscaled models. We have added Figure 16 in the appendix which only compares the (unscaled) baseline against the tokenized counterparts. Even in this case, we do see gains with tokenization, especially when using the CNN architecture. With the Impala architecture we see strong gains with tokenization as measured by the Mean and comparable performance with IQM. We report these 4 metrics as they provide a clearer picture of the difference between the methods (see [2] for an in-depth discussion of their differences). Furthermore, as we observe in Figures 5 and 17 with 4x scaling, PerFeat seems to be a stronger tokenization approach than PerPixel when using without SoftMoE; this suggests that the gains we observe with the unscaled baseline can be made larger with PerFeat tokenization. Finally, in Figure 15 we explored replacing flattening with global average pooling, as suggested by reviewer gzLL, which shows significant gains on both the unscaled and scaled baseline.

We hope this addresses your remaining concern?

[1] Ceron, Johan Samir Obando, et al. "Mixtures of Experts Unlock Parameter Scaling for Deep RL." Forty-first International Conference on Machine Learning.

[2] Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C Courville, and Marc 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 InformationProcessing Systems, volume 34, pp. 29304–29320. Curran Associates, Inc., 2021.

We thank the reviewer for their feedback! We are glad that the reviewer found the “analysis in-depth” and the “paper is well-written”.

W1, W2: primarily focuses on why SoftMoEs work well when scaled up…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. … 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.

A: As the reviewer mentioned, we focus in this work on understanding the significant performance improvement observed by prior work [1] on the studied benchmarks. Although we are running on single-task settings, there is a high-degree of non-stationarity due to the evolving policy. Our findings have implications on future research mainly by 1. revisiting the architectural choices and replacing the flattening operation, and 2. improving expert utilization in MoEs for further performance improvements. We think that studying the effectiveness of SoftMoEs in OOD generalization is an orthogonal, but valuable future line of research. We are currently setting up infrastructure to evaluate PPO on ProcGen, as suggested by the reviewer, and will report back here when we have results.

W3: 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.

A: Indeed, as discussed in section 5.1 (and in [1]), DQN appears to benefit less from the use of SoftMoEs in general, which may explain why SoftMoE-1 yields little gain. We hypothesize this may be due DQN’s use of regression versus Rainbow’s classification (C51) loss; we are currently running experiments with SoftMoE-1 (x4) with the C51 loss and will report back here once they have made more progress.

Q1: 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?”

A: We thank the reviewer for their nice suggestion! We have run the full analyses in Section 4.2 on DER and added the results in Appendix B.3 and pointed to it in Section 4.2. We find that all the observations are consistent with our previous results as shown in Figure 16.

Q2: 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 [hw, d], which I believe corresponds to PerConv tokenization. Would similar trends hold if [d, hw] (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.

A: We thank the reviewer for this suggestion. We are currently running experiments with tokenization and 2x scaling, to verify whether we observe a scaling trend. With regards to tokenization, we selected PerConv tokenization for Figure 5 given that we found it to be the most performant (see Figure 7). However, we are currently running experiments with PerFeat tokenization to see if the same trend holds, as suggested by the reviewer. We will report back once the experiments are completed.

Q3: 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?

A: We are not selecting 10% of the slots. The number of slots p is a predefined constant number which determines the capacity of each expert and it is typically fixed throughout training. In this experiment, we just reduced the default value of this number by 90%. We refer the reviewer to the added description to Section 3 regarding the choice of $p$ (number of slots), and have added more details to section 5.1 to avoid confusion.

Q4: Why was the analysis in Section 6 (Figure 11) not conducted for DQN?

A: Similar to the rest of the paper, we focused on the cases where SoftMoE was observed to achieve significant performance gains over the baseline. Following your suggestion, we are running some analysis on DQN and will include the results in Appendix B.5.

[1] Ceron, Johan Samir Obando, et al. "Mixtures of Experts Unlock Parameter Scaling for Deep RL." Forty-first International Conference on Machine Learning.

We thank the reviewer for their feedback! We are pleased to hear that they find the “results significant” and that our finding on expert underutilization has “an important implication, encouraging more research in this area”.

W1: 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.

A: Thank you for your suggestion, we have added more clarification about the slot in the revised version, in addition to what is presented at the end of section 3 (and in Figure 3).

W2: DQN … further investigation

A: Indeed, as discussed in section 5.1 (and in [1]), DQN appears to benefit less from the use of SoftMoEs in general, which may explain why SoftMoE-1 yields little gain. We hypothesize this may be due DQN’s use of regression versus Rainbow’s classification (C51) loss; we are currently running experiments with SoftMoE-1 (x4) with the C51 loss and will report back here once they have made more progress.

Q1: 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?

A: Your interpretation is correct! Another benefit of combined tokens is reducing the computational time without performance drop. We can achieve almost the same results using only 10% of expert slots if slots contain combined tokens unlike the case of sparse tokens. Using 10% of the slots means that the number of processed inputs of each expert is reduced by 90%. We have compared the wall-time in the two cases and observed time saving, but it is relatively marginal given the size of networks we are using.

Q2: can the authors run a baseline where the cnn encoder output is actually a vector?

A: Thank you for this insightful suggestion. We have conducted the requested experiments using global average pooling in the default and scaled baselines and included the results in Appendix B.2. Interestingly, consistent with our findings, replacing the flattening operation with global average pooling leads to performance gains in all cases. We appreciate the reviewer’s question which further strengthens the suggestion for revisiting the common practice of flattening the outputs of the convolutional encoders.

[1] Ceron, Johan Samir Obando, et al. "Mixtures of Experts Unlock Parameter Scaling for Deep RL." Forty-first International Conference on Machine Learning.

We thank the reviewer for their feedback! We are glad that the reviewer found the “analysis solid” and the paper is “well-presented”.

W1: unclear whether the same observation is universal applicable to other simulation platforms

A: Interesting question! In this work, our primary focus is to understand the underlying reasons for the significant performance improvements observed by [1] on discrete tasks. Our findings turn out to help explain why MoE does not yield performance gains in Mujoco environments, as discussed in Section 8. We agree with the reviewer that extending this work by investigating other domains would be a valuable future direction. We have included this in the discussion section in the revised version.

W2: Would the conclusions be applicable to different numbers of experts?

A: We did provide the same analyses on 8 experts in the appendix. As presented in Appendix B.2, the findings and conclusions are consistent with the case of 4 experts. Following your suggestion, we have also included the results for 2 experts and observed similar trends. We added a reference to this appendix in Section 5 of the revised version.

W3: In Figure 6, it will be beneficial to visualize the performance with one unscaled expert to understand the results better.

A: We have added this to Figure 6, thank you for the suggestion!

[1] Ceron, Johan Samir Obando, et al. "Mixtures of Experts Unlock Parameter Scaling for Deep RL." Forty-first International Conference on Machine Learning.

We thank the reviewer for their feedback! We are happy that the reviewer found the paper has a “potentially significant impact”. It is great to hear that it could be “a common ground for future methodologies”.

W1: The scaled SoftMoE-1 resembles a similar performance/optimality gap from that of the scaled Baseline on DQN

A: Indeed, as discussed in section 5.1 (and in [1]), DQN appears to benefit less from the use of SoftMoEs in general, which may explain why SoftMoE-1 yields little gain. We hypothesize this may be due DQN’s use of regression versus Rainbow’s classification (C51) loss; we are currently running experiments with SoftMoE-1 (x4) with the C51 loss and will report back here once they have made more progress.

W3, Q1: Tokenization baseline: …How does the result depicted in Figure 5 reduce to the claim: “providing strong evidence..”

A: The only architectural change in experiment in Figure 5 was replacing the flattening operation with tokenization, allowing us to directly assess the impact of tokenization on enhancing the performance of scaled networks. Thanks for noting this, we clarified this sentence in the revised version by rephrasing it to “… plays a major role in the successful scaling of DRL networks”.

W2, Q2: Expert specialization (line 231): Unknown p value ...if the default value of p is close to the number of tokens...Would the authors be able to list the precise configurations?

A: Following the common practice in the MoE literature [2,3], we set $p$ to be the total number of tokens divided by the number of experts (unless otherwise specified). This results in a $p$ value that is $\frac{1}{numexperts}$ times smaller than having $p$ equal to the number of tokens, as discussed in the expert specialization section. Following your suggestion, we have clarified this in Section 3 in the revised version.

W4,Q3: Is a hyperparameter search for random resets and S&P conducted in section 6?

A: Yes, we searched for the reset period and the S&P values. We have included these details in the revised version in Appendix B.4 and added a link to it in Section 6.

W5: expert utilization (section 6) marks an interesting future direction….the paper would be improved by moving section 6 to the appendix

A: We thank the reviewer for their suggestion. We find it interesting to add a discussion on future work in the main paper, with positive empirical evidence. However, we will keep this suggestion in mind for the camera-ready version, if space becomes an issue.

W6-8: We thank the reviewer for providing formatting suggestions. We incorporated your comments in the revised version.

[1] Ceron, Johan Samir Obando, et al. "Mixtures of Experts Unlock Parameter Scaling for Deep RL." Forty-first International Conference on Machine Learning.

[2] Riquelme, Carlos, et al. "Scaling vision with sparse mixture of experts." Advances in Neural Information Processing Systems 34 (2021): 8583-8595.

[3] Gale, Trevor, et al. "Megablocks: Efficient sparse training with mixture-of-experts." Proceedings of Machine Learning and Systems 5 (2023): 288-304.