SimBa: Simplicity Bias for Scaling Up Parameters in Deep Reinforcement Learning

Dear reviewer 8rZo,

Thank you for your constructive feedback. To address your concerns, we have provided a detailed response to each comment. In addition, we have provided a common response to address key feedback from all reviewers. Please let us know if you have any further comments or feedback.

Question 4.1 "The plasticity analysis in Appendix C. could be confusing to the average reader. It shows a comparison to a basically ‘collapsed’ baseline MLP, as the effective rank is negligible. I believe it would therefore be much more informative to the reader to add an additional 2 rows of figures to Fig. 15:

• Row 1: Plasticity analysis on DMC - medium
• Row 2: Plasticity analysis with all the ablations of Fig. 4, to show the unique effects of every module."

We appreciate your suggestion to enhance the plasticity analysis for clarity. In response, we have conducted the following additional experiments:

• Added a plasticity analysis on the DMC-Medium environments.
• Included a plasticity analysis with all ablations from Fig. 4 to demonstrate the unique effects of each module.

These additions have been incorporated into Fig. 17 and 18 in Appendix C.

For the analysis on the DMC-Easy&Medium environment, SAC with the MLP architecture still exhibit high dormant ratios and large feature norms across both DMC-Easy & Medium indicating greater loss of plasticity.

For the component ablation study, integrating each component of SimBa individually shows improved average returns. However, ablating each component does not lead to significant reductions in plasticity measures, with only Post Layer Normalization showing an improved feature norm. This suggests that while individual components may not significantly impact plasticity measures on their own, their integration within the SimBa architecture is crucial for effectively mitigating plasticity loss.

Question 4.2 As important as it is to see that SimBa is computationally cheap, it would be also interesting to add some figures showing longer training (e.g. 5 million timesteps) on DMC hard and comparing to TD-MPC2 and the baseline SAC.

Thank you for highlighting the importance of evaluating SimBa over extended training periods. We have expanded our experiments to include training for 5 million timesteps on DMC-Hard environments. The results are presented in Appendix C under the additional ablations section.

The learning curves indicate that SimBa + SAC not only learns faster but also maintains more stable and consistent performance improvements over time, demonstrating better sample efficiency and reduced variance compared to the baselines.

Thanks to the reviewer's feedback, our new experiment confirm SimBa’s robustness and effectiveness in long-term training scenarios.

Question 4.3 I believe a mention of the improvements that the Impala network gave to RL should be in the paper, and maybe a short sentence about the correlation (Residuals connections -> Simplicity Bias -> better Convergence).

We agree with your suggestion and have updated the Related Work section on page 4 to include a citation of the Impala network.

Dear reviewer 2Hpz,

Thank you for your insightful and constructive comment. Your thorough review of each section of the paper has significantly improved its overall quality. In addition, we have provided a common response to address key feedback from all reviewers. Below, we have provided responses to each of your questions. Please let us know if you have any other concerns.

Question 3.1 Statistical significance is not apparent from the presented results. The standard error overlaps in some plots, e.g., Figures 1b, BRO, and SimBa on Figure 5b, and Figure 8 has no rages, but I guess it would be more challenging to add them since the analysis is in two dimensions, and also from Figure 14. The tables in the Appendix do not have standard deviations. Moreover, authors could use some t-test or Mann–Whitney U test to demonstrate statistical significance.

We appreciate the reviewer’s feedback regarding statistical significance and its importance in validating our results. To address this, we performed a Mann–Whitney U Test to assess the significance of the observed differences between SimBa and BRO across all environments. The results are summarized in the table below:

From the results, the performance differences between SimBa and BRO are statistically significant for Myosuite and HumanoidBench ($p$ < 0.05), highlighting SimBa's effectiveness compared to BRO in these environments. In contrast, the differences in the DMC Easy & Medium and DMC Hard environments are not statistically significant, as indicated by the higher p-values.

However, we would like to highlight that SimBa's strengths are not limited to achieving high performance for diverse tasks. Its design emphasizes computational efficiency and simplicity, setting it apart from BRO. Unlike BRO, which incorporates RL-specific techniques such as periodic resets, quantile regression losses, and complex explortaion techniques, SimBa achieves competitive results solely through architectural enhancements. This approach avoids any modification to the learning objective or loss function of the underlying RL algorithm, ensuring that SimBa remains both straightforward to implement and broadly applicable across diverse scenarios.

Question 3.2 It would be valuable to see reversed to Fig. 4 ablation, i.e., simplicity and Return when you turn off every single component. I would suggest visualizing this similarly to Figure 4.

We appreciate your observation regarding the importance of providing experimental results for the removal of each component in SimBa. To address this, we conducted additional ablation studies and included the following analyses:

• Evaluated the impact of removing individual components from SimBa on mean performance for DMC Hard.
• Assessed the simplicity measure of the architecture after the removal of each individual component from SimBa.

We have included these results in our revised version of the manuscript in Appendix C.1 and are illustrated in Figure 15.

As shown in Figure 15, the removal of any single component results in a degradation of overall performance, underscoring the necessity of each component for SimBa’s effectiveness. Additionally, we observe that, except for RSNorm, the removal of residual connections, pre-layer normalizations, or post-layer normalizations leads to a reduction in the simplicity score. This further emphasizes the critical role of each individual component in SimBa’s design.

Regarding RSNorm, although its removal significantly impacts the agent's performance, the overall simplicity score of the architecture increases rather than decrease. We believe this discrepancy arises because our Fourier Transform-based complexity measure heavily relies on evaluating the network's outputs over a uniformly distributed input space. RSNorm is specifically designed to normalize activations based on the statistics of non-uniform or dynamically changing input distributions, excelling with inputs that exhibit variance and complexity not captured by a uniform distribution. Therefore, using a uniform input distribution in our simplicity analysis does not fully showcase RSNorm's advantages or accurately reflect its impact on the simplicity score.

Regarding the answer to Question 3.1 Thank you for the detailed explanation and for conducting the Mann–Whitney U test on SimBa and BRO. I also appreciate the authors' effort to highlight that SimBa’s strengths are not solely limited to achieving high performance but also include computational efficiency and simplicity, as well as its independence from RL-specific techniques such as periodic resets and quantile regression losses.

Therefore, regarding:

SimBa demonstrates superior scaling performance in terms of average return (...)

I would like to suggest extending the statistical analysis to the architectural comparisons presented in Section 5.2. Specifically, in Figure 5, there is an overlap of standard errors for larger architectures, which raises questions about the statistical significance of SimBa’s reported scaling superiority. Since all architectures are tested in the same algorithmic setup (i.e., without BRO’s RL-specific techniques), it would be valuable to assess the statistical significance of the above claim.

Regarding the answer to Question 3.2

As shown in Figure 15, the removal of any single component results in a degradation of overall performance, underscoring the necessity of each component for SimBa’s effectiveness. Additionally, we observe that, except for RSNorm, the removal of residual connections, pre-layer normalizations, or post-layer normalizations leads to a reduction in the simplicity score. This further emphasizes the critical role of each individual component in SimBa’s design.

and also in the manuscript one can find:

Stronger simplicity bias correlates with higher returns for overparameterized networks.

I appreciate the inclusion of Figure 15 and the authors’ effort to quantify the impact of removing individual components. However, I would like to understand better the claim that “stronger simplicity bias correlates with higher returns,” as stated in the manuscript. While this correlation seems evident in Figure 4, it is less apparent in Figure 15. SimBa without RSNorm is one case that questions this correlation. The second is the SimBa without PostLN also has a higher simplicity score than configurations without residual or preLN components but still worse returns.

Have the authors explicitly measured the correlation between simplicity bias and returns across different configurations? This is particularly important since this relationship forms a central theme of the paper. But, it is also interesting to know if it should be a universal rule of thumb for deep reinforcement learning practitioners. Clarifying this would significantly strengthen the manuscript’s conclusions and provide helpful guidance for the RL community.

Again, thank you for your constructive review of our manuscript. We appreciate your feedback and the opportunity to address your concerns. Below, we respond to your specific points:

3.1. Statistical Significance of Scaling Performance

We appreciate your suggestion to extend the statistical analysis to the architectural comparisons in Section 5.2. In response, we conducted additional Mann–Whitney U tests to evaluate the performance differences between SimBa and BRO across varying critic sizes. The results are summarized below:

The analysis reveals that SimBa significantly outperforms BRO for critic sizes ranging from 0.07M to 4.5M parameters (p < 0.05). However, for the largest architecture at 17.8M parameters, the performance difference is not statistically significant (p = 0.310). We speculate that at larger scales, both methods approach similar performance ceilings on the DMC-Hard tasks, thereby reducing the observable differences.

Based on these findings, we acknowledge that the claim of "superior scaling performance" for SimBa requires further nuance. Specifically, while SimBa demonstrates statistically significant scaling advantages for small to medium-sized architectures, its superiority becomes less pronounced for larger architectures. This refined interpretation offers a more balanced view of SimBa’s scaling performance.

Additionally, it is important to note that RSNorm, one of our key contributions, was applied to all experiments to ensure a fair comparison between architectures. Notably, the original BRO architecture without RSNorm exhibits significantly worse performance.

3.2. Correlation between Simplicity Bias Score and Performance

We recognize the importance of thoroughly investigating the relationship between simplicity bias and performance to substantiate our claims. To address your concerns, we conducted an explicit correlation analysis, detailed in Appendix D.

In our study, we computed the correlation between simplicity bias scores and the average returns of 12 distinct architectures, including four primary types: MLP, SimBa, BroNet, and SpectralNet. Additionally, we examined four SimBa variants with component reductions and four SimBa variants with component additions transitioning from MLP to SimBa. As discussed in Appendix C, the simplicity bias measure may not accurately capture the influence of RSNorm. Therefore, we separately evaluated the correlation for architectures with and without RSNorm. All twelve architectures have nearly identical parameter counts, approximately 4.5 million parameters each, with a maximum variation of 1%. Performance was assessed on the DMC-Hard benchmark, trained for one million steps.

The results indicate that architectures employing RSNorm exhibit a Pearson correlation coefficient of 0.79, demonstrating a strong positive correlation between simplicity bias score and average return in scaled architectures. In contrast, across all architectures, the Pearson correlation coefficient is 0.54, representing a moderate correlation. We attribute the lower correlation in the broader set to the current limitations of the simplicity bias measure.

We believe these findings support our hypothesis that higher simplicity bias scores correlate with improved performance in deep reinforcement learning with scaled-up architectures. We hope this work provides guidance to the RL community in designing effective new architectures.

Dear reviewer F5aT,

Thank you for your constructive feedback. To address your concerns, we have provided a detailed response to each comment. In addition, we have provided a common response to address key feedback from all reviewers. Please let us know if you have any further comments or feedback.

Question 2.1 It might be beneficial to see SimBa's performance on a wider range of RL domains, particularly with images.

Thank you for the suggestion to evaluate SimBa on image-based environments. We agree that this is an important direction and have included it in our future work section.

As an preliminary insight, we anticipate that applying SimBa's simplicity bias principles to vision-based control architectures, such as, on top of BBF [1], could provide performance enhancements. Specifically:

• Convolutional Encoder:
  • Normalization Layers: Incorporating normalization layers (e.g., layer-normalization, batch-renormalization) can stabilize training.
  • Reduced Non-Linear Activations: Minimizing non-linear activations in the non-linear residual branch aligns with practices from advanced vision architectures like ConvNeXtV2 [2].
• Predictor Network:
  • Integrating SimBa: Replacing the MLP with a SimBa network could help to alleviate overfitting present in vision-based training.

[1] Bigger, Better, Faster: Human-level Atari with human-level efficiency., Schwarzer et al.

[2] ConvNeXt V2: Co-designing and Scaling ConvNets with Masked Autoencoders., Woo et al.

Question 2.2 The paper doesn't discuss potential limitations or scenarios where SimBa might not be as effective

Are there any specific types of RL problems or environments where SimBa might not be as effective?

While SimBa consistently outperforms or matches naive MLP architectures in over-parameterized regimes across various experiments, its effectiveness may be limited in scenarios with constrained model capacity, such as under-parameterized or on-device settings. In these cases, SimBa's simplicity bias might restrict its ability to effectively fit the target function and could negatively impact parameter count and VRAM usage. Nonetheless, since the majority of the problems we address operate in over-parameterized regimes, we believe SimBa remains highly effective overall.

Question 2.3 Why are the standard errors in Figure 10 so high?

The high standard errors in Figure 10 are due to the binary nature of task outcomes in Craftax, where each trial results in either success (1) or failure (0).

Dear reviewer 9P9G,

We appreciate your insightful questions and positive support. We have provided a detailed response to each comment. In addition, we have provided a common response to address key feedback from all reviewers. Please let us know if you have any further comments or feedback.

Question 1.1 The proposed RSNorm is too similar to the RMS Norm, while the author has no discussion about the RMS Norm.

We apologize for causing confusion between Running Statistic Normalization (RSNorm) and Root Mean Square Normalization (RMSNorm). Despite the similarity in their names, RSNorm and RMSNorm differ fundamentally in two key aspects:

1. Normalization Axis:
   • RSNorm normalizes each feature dimension independently across the batch, similar to Batch Normalization. For example, if one feature is sampled from N(0,1) and another from N(10,1), RSNorm scales each to N(0,1), preventing any single feature from dominating due to significant differences in scale.
   • RMSNorm normalizes features within each sample, which may result in features with larger variances to disproportionately influence the model.
2. Statistics Used:
   • RSNorm utilizes cumulative running statistics for normalization. This approach is more stable and less noisy compared to using mini-batch statistics, which is particularly beneficial in reinforcement learning settings where batch estimates can be unreliable.
   • RMSNorm does not rely on any mini-batch based statistics, so it doesn't need to track running statistics.

To further address your concern, we have now included a discussion of RMSNorm in our manuscript and conducted comparative experiments on DMC-Hard averaged over 5 random seeds. The results are summarized below:

As shown, RSNorm significantly outperforms RMSNorm and other normalization methods, highlighting the effectiveness of RSNorm. We have also revised Figure 12 to include RMSNorm as a baseline for clearer comparison.

Question 1.2 Unfair comparison between SAC+SimBa vs others. Since SAC+SimBa introduces some other layers, it might include more parameters for the neural networks. The author can make some ablations with different sizes of SAC to exclude the capacity issue.

Missing details of the model size. I didn't find the detailed sizes of each model.

We appreciate your concern regarding potential capacity differences between SimBa and other architectures (MLP, BRONet, SpectralNet). To ensure a fair comparison, we have provided detailed parameter counts with a hidden-dimension $d$ and num-blocks $2$.

• MLP: $17d^2$.
• SpectralNet: $17d^2 + 8d$.
• BroNet: $17d^2 + 10d$.
• SimBa: $17d^2 + 8d$.

Additionally, we present the proportion of parameters contributed by normalization layers in SimBa:

As shown, normalization layers contribute less than 1% of the total parameters across all model sizes, making their impact on model capacity negligible. Therefore, the superior performance of SAC+SimBa stems from the effective integration of normalization layers rather than an increased model capacity. To improve clarity, we have updated Figure 5 to display the number of parameters on the x-axis instead of the hidden dimension, accurately reflecting our comparisons.