An Empirical Study on Normalization in Mamba

Thank you for your detailed review and valuable feedback. We have carefully considered your comments and made several revisions to the paper. Below are our detailed responses to the key points you raised:

1. On Task Selection and Dataset Expansion

• Our Response:
  • We understand your concern regarding the generalizability of our findings. In our study, we divided tasks into two major categories: sequence tasks and vision tasks. We chose one representative task from each category for the experiments. To better assess the applicability of the normalization techniques, we expanded our experiments in the revised version to include ImageNet-1k and the Breakfast dataset, covering a broader range of tasks and datasets.
  • Due to computational cost constraints, we were unable to cover all domains. However, we have provided fully reproducible experimental settings, and we encourage other researchers to apply our proposed methods to different datasets and models. We also welcome readers to validate our normalization strategies on additional datasets and models, particularly for sequence and vision tasks.

2. On Contradictory Results between Sequence and Vision Tasks

• Our Response:
  • The contradictory results between sequence and vision tasks are indeed a significant finding in our work. These results suggest that there is no universal "best" normalization method, and different tasks require different normalization strategies based on their characteristics (e.g., sequence data vs. image data).
  • However, we have identified a normalization combination that, while not optimal, performs well across tasks and can serve as a generally applicable solution in many cases. We believe this is an important contribution, as it offers a practical starting point for Mamba networks across various tasks.
  • In the revised version, we provide a more detailed analysis explaining why normalization after the S6 Module benefits sequence modeling but not image classification. The differences arise due to the inherent properties of sequence-based data and the specific requirements of sequence modeling.

3. On Experimental Settings and Reproducibility

• Our Response:
  • We fully agree with your concern about reproducibility and have updated the paper with more detailed information about our experimental setup:
    • We have clarified the hyperparameters used, including learning rate, batch size, and training steps.
    • We added details on the accuracy metrics used in our results (such as those in Tables 2 and 3), ensuring replicability.
    • We also included a section on data preprocessing steps and evaluation metrics to ensure transparency and reproducibility.

4. On Inconsistent Visualization

• Our Response:
  • We have addressed this issue by improving the consistency of the visualization in Figure 3:
    • The x-axes have been standardized across all subfigures and labeled more clearly to make comparisons of weight norms across different configurations easier.
    • We have ensured that the plotting styles are consistent across all subfigures, improving the clarity and interpretability of the results.

5. On the Weight Domination Hypothesis

• Our Response:
  • We agree with your observation and have provided additional empirical evidence to support the Weight Domination hypothesis in the revised version. Specifically:
    • We conducted a more detailed analysis of the BN→None configuration and discussed the experimental conditions that may have led to contradictory results. By considering task-specific factors, we have clarified the discrepancies and provided a more nuanced interpretation of the results.
    • We believe that the additional empirical evidence strengthens the Weight Domination hypothesis and explains the variations observed under different experimental conditions.

6. On Limited Generalizability and Task-Specific Conclusions

• Our Response:
  • We acknowledge that our work primarily focuses on task-specific insights, but we believe these insights are valuable for understanding how normalization strategies affect different types of tasks.
  • In the revised version, we further discuss how our findings contribute to a broader framework for normalization across tasks. While normalization strategies should be tailored to specific tasks, the combination we propose offers a generalizable starting point for both sequence and vision tasks.
  • We understand the limitations of relying on task-specific conclusions, and we plan to explore more generalizable insights in future work. However, the methods we provide for sequence and vision tasks offer valuable guidance, and we hope they will assist future research in this area.

For the listed issues, we have addressed them in the revised version of the paper. Please refer to the revised version for the answers.

Dear reviewer tLtQ,

As the discussion phase is nearing its conclusion, we would like to know if our responses have addressed your concerns. Looking forward to further discussion with you. We are happy to address any questions if such occur. Wish you have a nice day!

Best regards,

Authors

Thank you for your valuable feedback on our paper. We appreciate the time and effort you have taken to review our work. Based on your suggestions, we have made several revisions to address the concerns you raised. Please refer to our revised paper. Below are our detailed responses to the key points:

1. On the Contradictory Results between Vision and Sequence Tasks

• Our Response:
  • We fully agree with your observation that the results for vision and sequence tasks appear contradictory. This is indeed a key finding in our work, and we acknowledge that there is no one-size-fits-all normalization method. This discrepancy reflects the fact that different tasks may require different normalization strategies based on their characteristics (e.g., image data vs. sequence data).
  • However, we have also found a combination of normalization methods that, while not necessarily optimal, is generally applicable across various tasks. This method shows strong performance and stability, making it a practical approach for Mamba networks. We believe this is a valuable contribution to the field, as it provides a more widely applicable normalization strategy that could be used as a starting point for various tasks.
  • In the revised version of the paper, we expanded the experimental scope by including ImageNet-1k and Breakfast datasets, which helps broaden the task and dataset coverage. Although we could not cover all domains due to computational constraints, we have provided fully reproducible experimental settings, encouraging other researchers to apply our proposed methods to different datasets and models.
  • For specific tasks, we suggest that researchers refer to our task-specific normalization methods as priors based on whether the task is vision-based or sequence-based. If they wish to find the optimal normalization strategy for a given task, they can follow our experimental methodology to explore this further.

2. On the Focus on Long Sequence Modeling for Normalization Position Analysis

• Our Response:
  • In response to your concern, we have added more detailed analysis of the placement of normalization layers in the revised version of the paper. Specifically, we now provide a more in-depth explanation of why long sequence modeling is an important and meaningful context for investigating the impact of normalization position.
  • Our rationale is that Mamba networks, with their focus on processing long sequences, face unique challenges in terms of training stability and convergence. Therefore, analyzing the effect of normalization in this context is critical for optimizing the network's performance. We hope the expanded discussion in the revised paper addresses your concerns.

3. On the Weight Domination Hypothesis

• Our Response:
  • We acknowledge your concerns regarding the Weight Domination hypothesis and have taken steps to address them in the revised version. Specifically, we have provided more empirical evidence supporting the hypothesis, with a focused discussion on the BN→None configuration.
  • We conducted a more thorough analysis of the experimental conditions that may have led to the contradictory results observed in BN→None performance. We have clarified the specific factors, such as task-dependent variations and network architecture, that contribute to these observations. This additional analysis provides stronger support for the Weight Domination hypothesis and helps explain the apparent contradictions. We believe these revisions significantly strengthen the credibility of the hypothesis.

4. On Inconsistent Visualization

• Our Response:
  • In response to your feedback, we have made significant improvements to the visualization in Figure 3. Specifically:
    • We standardized the x-axes across all subfigures and added clearer labels to facilitate easier comparison of weight norms across different configurations.
    • We ensured consistent plotting styles across all subfigures, which improves the clarity and interpretability of the results. These changes should make the figure more intuitive and easier to understand.

5. On the Lack of General Insights

• Our Response:
  • We acknowledge that our paper primarily focuses on task-specific findings, but we also believe that the insights we provide are valuable in understanding how normalization strategies affect different tasks. To address your concern, we have discussed in the revised version how our findings could contribute to a broader framework for normalization across tasks.
  • While our study shows that normalization strategies should be tailored to specific tasks, the combination we propose offers a generalizable starting point for both sequence and vision tasks, offering practical insights for researchers working with Mamba networks.

class Norm_Select(nn.Module):  
    def __init__(self, norm_type="BN", dim=0):  
        super(Norm_Select, self).__init__()  
        # Initialize different normalization layers based on the selected type  
        if norm_type == "BN":  
            self.norm = nn.BatchNorm1d(num_features=dim)  
        elif norm_type == "GN":  
            self.norm = nn.GroupNorm(num_groups=4, num_channels=dim)  
        elif norm_type == "IN":  
            self.norm = nn.InstanceNorm1d(num_features=dim, affine=True)  
        elif norm_type == "LN":  
            self.norm = nn.LayerNorm(normalized_shape=dim, elementwise_affine=True)  
        elif norm_type == "RMSN":  
            self.norm = RMSNorm(hidden_size=dim)  
        elif norm_type == "None":  
            self.norm = None  
        else:  
            raise ValueError(f"Unsupported normalization type: {norm_type}")  

    def forward(self, x):  
        if self.norm is None:  
            return x  
        # Ensure the correct shape for normalization layers  
        if isinstance(self.norm, (nn.BatchNorm1d, nn.GroupNorm, nn.InstanceNorm1d)):  
            x = x.permute(0, 2, 1)  
            x = self.norm(x)  
            x = x.permute(0, 2, 1)  
        else:  
            x = self.norm(x)  
        return x  

Thank you very much for your detailed review and valuable feedback. We appreciate your careful consideration of our work, and we have made several revisions to the paper based on your comments. Please refer to our revised paper. Below are our detailed responses to the key points you raised:

1. On Innovation and Novelty

• Our Response:
  • We understand your concern about the novelty of the work. We would like to clarify that, to our knowledge, no prior work has systematically investigated the use of normalization in the Mamba architecture. Our study is novel in that it systematically explores the impact of normalization type, position, and combinations on the performance and stability of Mamba models. We have found normalization methods that outperform the commonly used techniques in Mamba, and we provide practical guidelines for selecting optimal normalization strategies when designing new Mamba architectures.
  • Additionally, to address concerns about the novelty of the contributions, we have expanded the discussion of related works in the revised version, especially focusing on prior research on normalization methods (e.g., Xiong et al., 2020; Santurkar et al., 2018). We believe that these discussions highlight the unique contributions of our work to the Mamba architecture.

2. On Dataset Size and Experimental Validation

• Our Response:
  • Regarding the dataset size, we would like to clarify that the ImageNet-100 dataset, although smaller than the full ImageNet, contains 130,000 images and is still a substantial dataset for image classification tasks. Therefore, it should not be considered as a "small dataset."
  • In response to your comment, we have expanded our experiments in the revised version to include not only ImageNet-100 but also ImageNet-1k and the Breakfast dataset. This allows us to cover a broader range of tasks and datasets.
  • Regarding More Diverse Tasks: Due to computational cost limitations, we were unable to conduct experiments across all domains. However, we provide a fully reproducible setup, and we encourage researchers in other fields to apply our proposed methods. We also welcome readers to validate our normalization strategies on additional datasets and models, especially for tasks involving both vision and sequence data.

3. On the Lack of Discussion of Related Works

• Our Response:
  • Our primary focus in this paper is on the application of normalization techniques within the Mamba architecture. As a result, we did not extensively discuss related works on normalization in other architectures. However, we understand the importance of situating our work within the broader context, and we have added a discussion on relevant literature in the revised version.
  • Specifically, we have included references to important prior works on normalization (e.g., Xiong et al., 2020; Santurkar et al., 2018) and clarified how our work differentiates itself by addressing normalization strategies specific to Mamba.

4. On Experimental Settings and Reproducibility

• Our Response:
  • We take reproducibility seriously and have updated the revised version with more detailed information about our experimental setup. Specifically:
    • We have clarified the hyperparameters used in the experiments, including learning rate, batch size, and training steps.
    • We have added details about the accuracy metrics used in our results (such as those in Tables 2 and 3), ensuring that other researchers can replicate our experiments.
    • We have also included a section outlining the data preprocessing steps and evaluation metrics used in our experiments.

5. On Inconsistent Visualization

• Our Response:
  • We have addressed this issue by improving the consistency of the visualization in Figure 3. Specifically:
    • The x-axes have been standardized across all subfigures, with clearer labels to make the comparison of weight norms across different configurations easier.
    • We have also ensured that the plotting styles are consistent across all subfigures, improving the clarity and interpretability of the results.

6. On the Weight Domination Hypothesis

• Our Response:
  • We agree with your observation and have provided additional empirical evidence to support the Weight Domination hypothesis in the revised version. We have also discussed the apparent contradictions in the BN→None configuration results and provided a more detailed explanation of the experimental conditions that may have contributed to these observations.
  • We have now included more robust empirical support for the hypothesis and clarified how the results vary depending on the task and configuration.

Dear reviewer dm4d,

As the discussion phase is nearing its conclusion, we would like to know if our responses have addressed your concerns. Looking forward to further discussion with you. We are happy to address any questions if such occur. Wish you have a nice day!

Best regards,

Authors