We appreciate your thorough review of our paper. Our key contribution addresses the challenge of learning long-tail domain data during the pretraining stage, substantially reducing the need for an expensive, labor-intensive second domain-specific pretraining stage in practical applications.

Q1 ..setup is not clear

Both BERT and GPT undergo both pretraining and domain-specific pretraining, while our method only pretraining. All models then serve as embedding generators with all params frozen for downstream tasks.

Q2 ..language modeling not covered

Our BERT-based MoE uses MLM for pretraining, while the GPT-based MoE employs causal language modeling. All downstream tasks are NLP tasks. We are unclear about the specific meaning of 'language modeling' you mentioned and appreciate further clarification.

Q3 .. MoE pre-trained with Wikipedia.

We compared MoE with 1) only pretraining 2) additional domain-specific pretraining and 3) our method. Results show that while domain-specific pretraining boosts baseline performance, our single-stage pretraining method surpasses both, efficiently learning long-tail domain.

Q4 .. 1-gram frequency too simple.

Global rebuttal Figure3 displays sentence frequency distribution using 2/3-gram averages, alongside gradient consistency. Results align with 1-gram analysis, confirming its effectiveness despite simplicity.

Q5 Compared with Figure2a, consistency in Figure3a..not..better..

Both figures were obtained on a baseline model. We'll clarify this in the main text.

Q6 cluster distance not meaningful..norm..larger.

In global rebuttal Figure2, we introduce a new metric—the ratio of cluster distances to cluster radii—to measure cluster structure clarity independent of norm values. This aligns with our observation that clusters become more distinct in deeper layers.

Q7 line 152:"..high domain coherence within the smaller clusters.." not specified.

In global rebuttal Figure4, we detail the content of sentences within each small cluster. For example, 'Java Network Debug' cluster frequently includes HTML, Javascript, and <link>, demonstrating high domain coherence.

Q8 line 157:"..frequency scores below 0.2 shows random distribution.."..random distribution unclear.

The term "random distribution" describes how extremely low-frequency sentences are irregularly scattered across clusters without clear semantic correlation. In Global rebuttal Figure4, we detailed the distribution and contents of these sentences, which often contains misprinted formulas or non-English.

Q9 not that surprising that samples in the same cluster have similar gradients.

Our main contribution is that we are the first to apply gradient similarity to identify and tackle the challenges of learning long-tail data during pretraining, significantly reducing the need for subsequent domain-specific pretraining.

Q10 Section 2.2 analyzes their CSE module .. formally introduced in Section 3..hard to follow.

Our goal was to ensure readers first understand the fundamental challenges of long-tail data learning, which sets the stage for appreciating the effectiveness of our method. We'll include a more detailed illustration of CSE in analysis section of revised version.

Q11 ..line 1 to line 12 in Figure3c?

We'll correct "line" to "layer".

Q12 The font size..in the figures is too small.

We'll increase the font size in figures.

Q13 ..discuss why..choose the testing datasets..domain of those datasets in the main text.

Our research identified long-tail distribution as a prevalent phenomenon in many common knowledge datasets, posing a challenge for pretraining. We chose Wikipedia for pretraining because it mirrors this distribution pattern. We also chose specific domains within the long-tail distribution with publicly available training corpora and downstream tasks for further testing. A detailed explanation of dataset choices will be included in the main text in revised version.

Q14 ..line 128:"optimization requirements for long-tail..overlooked .." not substantiated.

Existing researches [A, B] has shown that gradient conflicts lead to suboptimal learning outcomes for the affected data. In standard pretraining, since the gradient descent direction is dominated by common data, the differences in gradient direction between long-tail and common data result in the optimization of long-tail data being overlooked.

Q15 ..line 146: model's inability..long-tail..begins early..

Gradient consistency indicates how much a model focuses on learning specific data. Figure3a in the paper reveals that the GC of long-tail data declines sharply at the beginning of the pretraining stage and remains low, suggesting that the model's inability to effectively learn from long-tail data starts early in the training process.

Q16 ..line 148:"CSE layer..at..GC value peaks"

The peak GC value of a baseline model signals when it struggles to learn long-tail data. We then introduce CSE layer at this point to improve the model’s ability to learn long-tail data in later phases.

Q17 title..over-claiming

We'll replace 'finetuning' with 'domain-specific pretraining' to narrow down the specific training stage we aim to reduce.

Q18 read/cite this paper ..

We find this paper relevant and will cite it.

Q19 Do you add a classification head and fine-tune the whole model?

For each classification task, we train a classification head with all other params frozen.

References

[A] T. Yu, et al. Gradient Surgery for Multi-Task Learning. NeurIPS 2020.

[B] G. Shi, et al. RECON: REDUCING CONFLICTING GRADIENTS FROM THE ROOT FOR MULTI-TASK LEARNING. ICLR 2023.

Thank you for your comments.

DQ1 One of the most common usages of pretrained LMs is to use them without any fine-tuning. People use them to solve tasks by prompting the LMs. Unfortunately, it is not discussed in this paper.

DA1 We share your view that eliminating fine-tuning is highly valuable. While prompting is a common approach, domain-specific fine-tuning is also widely adopted, particularly for downstream tasks that require specialized domain knowledge[J, K, L, M, N, O]. In these cases, models often struggle to achieve the desired performance without specific expertise in the relevant domain. Fine-tuning is therefore essential to unlock the full potential of pre-trained language models for various downstream tasks.

Our work aims to minimize the need for cumbersome domain-specific fine-tuning, especially when models have already encountered relevant long-tail domain data during pretraining but have not effectively learned from it. In such cases, prompting alone may fall short in recovering lost knowledge without fine-tuning. Our method addresses this challenge by enhancing the model's ability to learn and retain long-tail domain knowledge during pretraining, thereby reducing the necessity for extensive domain-specific fine-tuning. Consequently, our approach provides a more robust solution, preserving and enriching critical domain knowledge, and offering a distinct advantage over purely prompting-based methods.

DQ2 For Figure 3(c), the authors should also specify what criteria/threshold they used for choosing the "long-tailed" data.

DA2 Thank you for your question. We applied the Elbow method to determine a threshold of 9.37% on the curve of domain proportions, classifying any domain with a data proportion at or below this level as a long-tail domain. By plotting the curve of Wikipedia domains and their corresponding proportions, we identified the point (9.37%) where the slope changes significantly, marking the transition from the head of the distribution to the tail. Following this, we randomly selected sentences from these long-tail domains for further analysis.

DQ3 "For each classification task, we train a classification head with all other params frozen." This is a uncommon setup now.

DA3 We share your perspective that there are other prevalent paradigms, but our approach is also widely adopted. The original BERT paper [I] mentions the feature-based fine-tuning approach, where "fixed features are extracted from the pre-trained model" and a "task-specific model architecture" is added and fine-tuned for specific tasks, such as a classification head for classification tasks. Similarly, the standard paradigm for training domain-specific models also follows this fine-tuning approach after a domain-specific pretraining stage. For instance, Section 4.1 in [A], Section 3 in [C], Section 3.3 in [G], Section 2.1 in [F], and other references [B, D, E, H] all explicitly state that they adhere to this paradigm.

DQ4 I am not sure if the next revision can incorporate all the clarifications and the new results/metric well. In my opinion, the current writing quality is not ready for publication for a top conference such as NeuRIPs.

DA4 Thanks for your question. We've summarized review comments received below and have already completed them in the revised version. We will incorporate further refinements as new comments come in.

Continued on the next post

Q1 Typo in line 189: Should be "Cluster" and not "Clsuter"; Typo in Table 1, 4th row: Should be "legal" and not "lgeal".

A1 Thank you for your comments. In our revised version, the error on line 189 will be corrected to Cluster, and the typo in Table 1, fourth row, will be updated to legal.

Q2 Its better if the limitations are in the main paper and not in the appendix.

A2 Thanks for your suggestion. We agree that including the limitations in the main paper rather than the appendix will enhance clarity for readers. We will make this adjustment in our revised document.

Q3 I understand that doing this experiment for bigger models on academic resources is a tough ask. But I still ask if the authors have any intuition on what might happen if this method is applied to larger models? My second question is in a larger model, do the authors think the GC for long-tailed knowledge will be even lower?

A3 Thank you for your comments. Our latest experiments, which included 330M GPT-style models and tasks from more domains as detailed in the global rebuttal, demonstrate that our method continues to outperform baseline methods at a larger scale. Your question about the gradient consistency (GC) of long-tailed knowledge in larger models is crucial, as it directly relates to how our method would adapt to larger scales. We've identified three key factors influencing this:

• Model Size Larger models have greater learning capacities, which can enhance the GC for long-tail knowledge given a fixed training data size.

• Window Size A larger window size allows the model to capture longer-range relationships within each training sequence, potentially highlighting more subtle distinctions and leading to a more long-tail distribution.

• Training Data Size Larger data sizes and larger window sizes might result in a higher proportion of common data relative to long-tail clusters. The increased presence of common data and the potential subdivision of long-tail clusters into smaller groups could, paradoxically, reduce the GC for long-tail knowledge.

Q4 Can this method be applied to both FFN and Attention modules? I am curious since in some cases MoA is better and in some cases, MoF is better, and was wondering whether using both would lead to the best results.

A4 Thank you for your suggestion. We tried this architecture and tested it in the same small-scale setting reported in the paper, and found that MoA+MoF yields improvements in the overall performance compared to MoA/MoF only.

We appreciate the time and expertise you have invested in reviewing our submission. Below, we outline our responses to the specific points raised in your reviews.

Q1 Typo: line 189/300 should be Cluster rather than Clsuter.

A1 Thank you for your comments. It has been corrected to be clusters in our revised document.

We appreciate the time and expertise you have invested in reviewing our submission. Below, we outline our responses to the specific points raised in your reviews.

Q1 The improvement in the GPT style model is subtle compared to the baseline MoE method, which may undermine the advantage of the proposed methods.

A1 Thank you for your comments. The experiments presented in the paper are conducted with 130M models, and we've conducted larger-scale experiments with 330M GPT-style models and detailed the results in the global rebuttal. Results indicate that our approaches match and outperform the baseline MoE model, achieving an average accuracy improvement of 3.26%.

Q2 The experiments are not sufficient, maybe more tasks should be tested. Is it possible to conduct the experiments on a pre-trained model, which would make the contribution stronger?

A2 Thank you for your comments. To further substantiate the efficacy of our method, we've conducted experiments on a wider array of tasks drawn from diverse domains using the 330M GPT-style model. As detailed in the global rebuttal, we added domain-specific tasks from the academic, environmental, and financial sectors to demonstrate the generalization capability of our method across multiple domains. Additionally, we have included results on general knowledge tasks to confirm that our dedicated long-tail learning method does not compromise general knowledge acquisition.

We also conducted experiments on a pre-trained 110M scale model as you suggested, wherein all methods continue training from a single pre-trained checkpoint. The outcomes are presented in the following table. Results show that directly applying our method to a pre-trained model still yields superior performance compared to the baseline and MoE models.

Q3 In Figure 3(c), the legend is lines 1-12. Maybe the author wants to show layers?

A3 Thank you for your comments. We'll correct it to "layers 1-12" instead of "lines" in our revised document.

Q4 Can you conduct computation time and memory analysis of the proposed methods as you claim that fine-tuning is time-costing?

A4 Our method introduces negligible overhead in time and memory.

• Memory The additional memory consumption incurred by our method compared with an MoE who has an equivalent number of experts is attributable to the storage requirements for $k$ distinct cluster centers, each endowed with a dimension of $d'$, which is $O(kd')$ in total. In most practical scenarios, the $k$ is lower than a hundred, and $d'$ tends to be around a hundred. Thus $O(kd')$ is negligible when compared with a neural network with often millions of parameters. Moreover, during a forward pass, our model mirrors both the Mixture of Experts and baseline architectures in the number of parameters activated, thus not introducing a memory burden.

• Time Throughout the training phase, the additional computation introduced by our method, relative to both the baseline and MoE, is the cluster initialization step. The time complexity of this operation will be $O(n\log n d')$ and is bounded by $O(n^2d')$ for the worst case, where $n$ represents the number of sampled points. Notably, during subsequent inference stages, the quantum of parameters activated per iteration by our algorithm aligns identically with that of both MoE and the baseline.

Since MoE and the baseline necessitate an additional fine-tuning phase, it becomes evident that the time cost of this phase far surpasses the $O(n^2d')$ complexity of cluster initialization. To provide more straightforward data, we list the wall-clock time of training & finetuning 330M GPT on 20B tokens with 2 NVIDIA GeForce RTX 3090 GPUs each in the table below.

We appreciate the time and expertise you have invested in reviewing our submission. Below, we outline our responses to the specific points raised in your reviews.

Q1 The baseline results in Tables 1 and 2 are surprising as, for example, BERT/med (further trained on medical) would be expected to outperform BERT/legal (further trained on legal) on medical datasets, but the two models are instead performing very similarly to each other, which suggests potential issues with baseline tuning.

A1 Thank you for your comment. Upon investigation, we discovered that the issue stemmed from catastrophic forgetting during the fine-tuning stage. Our training process strictly followed the basic framework of pretraining and domain-specific pretraining, and we didn't control the occurrence of forgetting, focusing more on the best performance. However, we monitored the forgetting phenomenon using the test loss on pretraining data. Specifically, BERT/med(further trained in medical), exhibited a more severe forgetting issue, with test perplexity on pretraining data increasing by 24.33, compared to a smaller increase of 8.41 for BERT/legal(further trained in legal). After correcting this imbalance by utilizing early-stop to select checkpoints where each model showed similar degrees of forgetting but not the best performance, the results came out that each domain-finetuned model outperformed on its respective domain tasks. This further underscores the challenges posed by the fine-tuning stage, affirming the value of our approach in not requiring domain-specific pretraining.

Q2 It is not clear if the proposed method retains its general knowledge as there are not sufficient benchmarks to judge general knowledge (e.g., WebQS, TriviaQA).

A2 Thank you for your comment. To address this, we've conducted experiments on general knowledge tasks, the results of which are detailed in the global rebuttal. These results demonstrate that our methods not only match but sometimes even surpass the performance of baseline models on a variety of general tasks. This is because our method enhances the model's capability to learn long-tail data by segregating the learning process between long-tail and non-long-tail data, allowing each category to be learned effectively without compromising one at the expense of improving the other.

Q3 Why is k-means clustering used for the analysis but DBSCAN is used for the main method? What is the value of K in k-means?

A3 Thank you for your comments. Initially, we used the K-means clustering algorithm with the elbow method to identify the optimal number of clusters (K). However, the elbow method's reliance on subjective judgment is a limitation. We then tested DBSCAN, which automatically determines the number of clusters based on density and is easier to set parameters compared to K-means. Consequently, we opted to employ DBSCAN in our solution design. In the global rebuttal Figure 1, we also presented cluster number and cluster distance calculated with DBSCAN, demonstrating that the choice of clustering algorithm does not impact the statements presented in Section 2.

Q4 Can the authors clarify the reasoning behind routing at sequence-level instead of token-level?

A4 Thank you for your comments. The definitions, analyses, and conclusions regarding long-tail data presented in Section 2 are all at the sequence level; consequently, routing is also performed at this same level. The rationale for defining and analyzing long-tail at the sequence level stems from the fact that when we discuss long-tails, we are essentially referring to semantics rather than individual words or tokens. Such semantics can only be observed within the context provided by sequences as a whole.

Q5 Are there any recommendations for adapting to more frequent (instead of long tail) domains?

A5 Thank you for your question. As discussed above, since data from frequent domains dominate the gradient direction of model optimization, our method separates the parameter of learning long-tail and frequent domains to improve the model's ability to capture long-tail domain knowledge. If we expect to improve the ability on frequent domains, we may need to design methods to analyze the gradient consistency within frequent clusters, and conduct for finer distinctions within the frequent domain.