The sparsely activated mixture of experts (MoE) model presents a promising alternative to traditional densely activated (dense) models, enhancing both quality and computational efficiency. However, training MoE models from scratch demands extensive data and computational resources. Moreover, public repositories like timm mainly provide pre-trained dense checkpoints, lacking similar resources for MoE models, hindering their adoption. To bridge this gap, we introduce MoE Jetpack, an effective method for fine-tuning dense checkpoints into MoE models. MoE Jetpack incorporates two key techniques: (1) checkpoint recycling, which repurposes dense checkpoints as initial weights for MoE models, thereby accelerating convergence, enhancing accuracy, and alleviating the computational burden of pre-training; (2) hyperspherical adaptive MoE (SpheroMoE) layer, which optimizes the MoE architecture for better integration of dense checkpoints, enhancing fine-tuning performance. Our experiments on vision tasks demonstrate that MoE Jetpack significantly improves convergence speed and accuracy when fine-tuning dense checkpoints into MoE models. Our code will be publicly available at https://github.com/Adlith/MoE-Jetpack.
Paper
Similar papers
Peer review
Summary
The authors propose MoE Jetpack, a method to transform dense checkpoints into a SoftMoE-like network. They introduce checkpoint recycling to utilize dense checkpoints as initial weights for MoE models. Additionally, the authors enhance the SoftMoE with the proposed Expert Regularization and adaptive Dual-path mechanisms. Extensive experiments and ablation studies are conducted to demonstrate the effectiveness of these proposed mechanisms.
Strengths
How to transform a dense checkpoint to an MoE is surely an important problem. The proposed checkpoint recycling method seems promising. Extensive experiments and ablations are given to prove the effectiveness of the proposed method.
Weaknesses
The comparison in the main results appears to be unfair. From my understanding, MoE Jetpack is initialized with a pre-trained model and then fine-tuned on the evaluation dataset. However, the baseline (SoftMoE) is trained from scratch solely on the evaluation dataset. For instance, in the ViT part (Table 1a), MoE Jetpack is initialized with a ViT pre-trained on ImageNet-21k and then fine-tuned on ImageNet-1k, whereas SoftMoE is trained from scratch on ImageNet-1k. This unfairness is even more pronounced on smaller datasets like STL-10 or Pets, as MoE Jetpack already benefits from prior knowledge gained from ImageNet-21k.
Questions
1. Could you compare MoE Jetpack with a Soft MoE pre-trained on ImageNet21k and fine-tuned on the evaluation dataset? 2. Could you compare MoE Jetpack with a model pre-trained on ImageNet21k, upcycling to an MoE, and fine-tuned (only router and experts' MLP layer) on the evaluation dataset?
Rating
5
Confidence
5
Soundness
3
Presentation
2
Contribution
3
Limitations
Limitation addressed in the paper
Response to Reviewer hQZu
Thank you for your feedback and for acknowledging the contributions of checkpoint recycling in our work. The statement “If a stronger dense checkpoint (e.g., pre-trained on CC12M) were used, the gap would be even larger” indeed confirms the effectiveness of our checkpoint recycling, indicating that a better pre-trained dense model can further enhance the performance of the MoE model. However, it appears that your concerns differ from those of Reviewer R7BZ9. Reviewer R7BZ9 suggested **integrating the comparisons in Table 2 into Table 1 to provide a more comprehensive comparison for the readers**. We find this suggestion reasonable and are considering it. We will also discuss further with Reviewer R7BZ9 on how to organize the tables in the paper. On the other hand, you seem to believe that the comparisons in our paper are unfair and that Table 1 is meaningless, as you stated, “such a comparison would not be meaningful.” This is fundamentally different from Reviewer R7BZ9's concern. **Table 2 in our paper provides a detailed ablation study comparing the results of training SoftMoE from scratch, SoftMoE with checkpoint recycling weights, and SephoreMoE with checkpoint recycling weights.** Table 1, on the other hand, showcases the overall improvements of the entire MoE Jetpack (checkpoint recycling + SephoreMoE). We hope that the ablation study in Table 2 addresses your concerns. It is worth noting, as mentioned in line 203 of our paper, “Baseline Implementation. We follow the implementation details outlined by Xu et al.” Our experiments were inspired by the work of Xu et al. [1], and Table 1 of our paper adopts a similar experimental setup as theirs. But if you feel that including the comparative results from Table 2 in Table 1 would be beneficial, we would be happy to consider your suggestions. **Regarding the comparison of convergence speeds in Figure 5, the main improvements are indeed attributed to checkpoint recycling.** Utilizing checkpoint recycling to **transform dense checkpoints into MoE initialization weights significantly enhances both convergence speed and performance, which is a key contribution of our work.** We will further emphasize that checkpoint recycling can significantly accelerate convergence and improve model performance and will also provide a comparison of the convergence speeds of SoftMoE with random initialization and SoftMoE with checkpoint recycling weights to make our paper more comprehensive. Thank you for your valuable suggestion. Additionally, **I have some questions regarding the reviewer's comments: "At the very least, the authors should include a baseline of a dense model pre-trained on ImageNet-21k and then fine-tuned on Pets."** Our Table 1, Row 8 already includes this result. Moreover, this result has no connection to MoE. I would appreciate it if the reviewer could further clarify. We hope our response addresses your concerns and look forward to your reply. [1] Xu, Zhiqiu, et al. "Initializing models with larger ones." ICLR 2024 (spotlight).
Response to Authors
In short, it is unacceptable to compare a pre-trained model (including models initialized from a pre-trained model) directly with a model trained from scratch, especially on small datasets. This is why I maintain that the comparison in Table 1 and Figure 5 is unfair. Consider the results on Pets in Table 1: there are only 7,349 samples in total in the Pets dataset. Training a Soft MoE (ViT-T) on this dataset yields poor accuracy due to insufficient data leading to underfitting. In contrast, MoE Jetpack performs well because its initialization weights are trained on ImageNet-21k. The performance discrepancy between Soft MoE and MoE Jetpack on Pets cannot be attributed solely to checkpoint recycling and SpheroMoE, as a significant portion of this discrepancy stems from the dense checkpoint. To make the comparison more persuasive and appropriate, I suggest using the following baselines in Table 1 instead of Dense, Dense (21k), and Soft MoE trained from scratch: 1. A dense model pre-trained on ImageNet-21k (the same model used for MoE Jetpack initialization) and fine-tuned on Pets. 2. An MoE model initialized from the dense checkpoint using Sparse Upcycling, and fine-tuned on Pets. Additionally, referring to the well-known Sparse Upcycling method[1], all their experiments compare the MoE model with a continuously trained dense model, not a model trained from scratch. [1] Komatsuzaki, Aran, et al. "Sparse Upcycling: Training Mixture-of-Experts from Dense Checkpoints." ICLR 2023.
Response to Reviewer Comments on Sparse Upcycling and Baseline Comparisons
Thank you for your insightful comments. I would like to address a few points regarding your review. You mentioned, "referring to the well-known Sparse Upcycling method [1], **all their experiments compare the MoE model with a continuously trained dense model, not a model trained from scratch.**" This statement appears to be problematic. Our experiments in Table 1 were indeed influenced by **Figure 4 in the Sparse Upcycling paper, which is titled: "Figure 4: Pretraining performance achieved by the upcycling method and a MoE model trained from scratch."** I kindly ask you to refer to Figure 4 in the Sparse Upcycling paper for confirmation. Regarding your suggested baselines: 1. A dense model pre-trained on ImageNet-21k (the same model used for MoE Jetpack initialization) and fine-tuned on Pets. 2. An MoE model initialized from the dense checkpoint using Sparse Upcycling and fine-tuned on Pets. For the first baseline, **a dense model pre-trained on ImageNet-21k and fine-tuned on Pets, we believe it does not directly relate to our MoE experiments.** Nonetheless, we will conduct this experiment and provide the corresponding results once completed (around 1 day). The second baseline is essentially an ablation study that compares Sparse Upcycling with our proposed checkpoint recycling method. **This ablation study is already presented in Table 3 of our paper.** To ensure fairness, **both Sparse Upcycling and checkpoint recycling utilize SpheroMoE layers and are compared on the larger ImageNet-1k dataset.** As per your request, we will conduct the experiments on the Pets dataset and include the results in Table 3. Thank you for your suggestions and thorough review. We will follow up with the additional experiments and provide the updated results accordingly.
Figure 4 in the Sparse Upcycling is indeed titled "Pretraining performance achieved by the upcycling method and a MoE model trained from scratch", but notice that the x axis in that figure only shows the "extra pretraining time". I think that what reviewer hQZu means is that: Suppose that you train the dense ImageNet-21k model for N steps, and then you use that to initialize the MoE for another M steps, you should compare the performance of the MoE with a dense model trained from scratch on N+M steps. Table 2 in your paper shows that MoE Jetpack is better than SoftMoE when training both for additional M steps, but it doesn't show if it's simply better to train the original dense model for another M steps. Of course, this is assuming that all models have the same cost per step, if not one should scale steps proportionally to FLOPs or wallclock time to fairly compare accuracy vs. cost.
Response to Reviewer 7BZ9
Thank you for your assistance in clarifying this matter. I understand your point regarding the comparison: "You should compare the performance of the MoE with a dense model trained from scratch on N+M steps." This experiment aims to demonstrate that **upgrading a dense checkpoint to an MoE model and training it yields better results than continuously training the dense model**, thereby validating the upcycling process. The third column in Table 1 (Dense(21k)) represents this process. It shows the performance of a model pre-trained on ImageNet-21K and then fine-tuned on downstream datasets using the same training strategy as MoE Jetpack. Therefore, I am a bit puzzled by the concern since this comparison appears to address the point raised. Additionally, Figure 4 in the sparse upcycling paper indeed compares sparse upcycling with MoE models trained from scratch. As stated in the paper, "Figure 4 compares sparse upcycling with sparse models trained from scratch. As training from scratch does not reuse the computation cost already sunk into the dense checkpoint..." The x-axis labeled "extra pretraining time" for the MoE curve represents its training time from scratch. If the MoE is pre-trained on the JFT dataset for M steps and then compared with sparse upcycling at N steps on the Extra Pretraining Time, the MoE will obviously achieve much better performance than sparse upcycling. I hope this explanation clarifies the intended comparisons and the significance of the upcycling process.
Response to Authors
1. The authors claim that the third column in Table 1 (Dense21k) is pre-trained on ImageNet-21k and fine-tuned on downstream datasets. However, I searched for the keyword 'dense' in the paper and did not find any clarification on this point. If I missed something, please point it out. 2. The MoE (green dots) in Figure 4 of the Sparse Upcycling paper is pre-trained on the JFT dataset for M steps (the cost to train the dense checkpoint) and then compared with sparse upcycling at N steps on the Extra Pretraining Time, unless the MoE can achieve 37% accuracy on JFT under random initialization.
Response to Reviewer hQZu comment on Dense21k
Thank you for your continued communication with us. We apologize for any confusion caused. Our paper mentions, "Tab. 1 compares the performance of the MoE Jetpack with **Dense ViT models (trained from scratch and with pre-trained weights on ImageNet-21k).**" The Dense ViT models with pre-trained weights on ImageNet-21k refer to models that use pre-trained weights from ImageNet-21k and are then fine-tuned on downstream tasks. **We will rewrite this section to: "Tab. 1 compares the performance of the MoE Jetpack with Dense ViT models trained from scratch, and Dense ViT with pre-trained weights on ImageNet-21k then fine-tuning on downstream datasets, represented as Dense (21k) in the table."** For the second point, **a randomly initialized MoE can achieve 37% accuracy on JFT after a short training period**, as shown in Figure 3 of softMoE [1] and Figure 2 of V-MoE [2]. We are unsure if we have misunderstood your point and look forward to your response. [1] Puigcerver, Joan, et al. "From sparse to soft mixtures of experts." ICLR 2024. [2] Riquelme, Carlos, et al. "Scaling vision with sparse mixture of experts." NeurIPS 2021.
Response to Authors on Dense21k
Your comment, "The Dense ViT models with pre-trained weights on ImageNet-21k refer to models that use pre-trained weights from ImageNet-21k and are then fine-tuned on downstream tasks," has added to my confusion. Previously, in your **Response to Reviewer Comments on Sparse Upcycling and Baseline Comparisons**, you made a contradictory statement: "For the first baseline, a dense model pre-trained on ImageNet-21k and fine-tuned on Pets, we believe it does not directly relate to our MoE experiments. Nonetheless, we will conduct this experiment and provide the corresponding results once completed (around 1 day)." Please clarify this contradiction, as it raises concerns about the soundness of your experimental results.
Response to Reviewer hQZu comment (Explain the weight selection process)
To clarify, it is important to distinguish between "using pre-trained weights from ImageNet-21k and then fine-tuning on downstream tasks" and "a dense model pre-trained on ImageNet-21k and fine-tuned on Pets." There are some differences between the two. The process of "using pre-trained weights from ImageNet-21k" involves the **ViT-S model, which, after weight selection, leads to the creation of the ViT-T model used for fine-tuning** on downstream tasks. This weight selection [1] step occurs between the pre-training and fine-tuning stages. For detailed information, please refer to the work by Xu et al. As we mentioned in our paper, **"We follow the implementation details outlined by Xu et al. [1] for comparisons of the dense models."** In contrast, "a dense model pre-trained on ImageNet-21k and fine-tuned on Pets" refers to directly fine-tuning the ViT-T model, which was pre-trained on ImageNet-21k, on the Pets dataset without the intermediate weight selection process from ViT-S to ViT-T. **Our MoE Jetpack follows the same procedure as Xu et al. [1], incorporating the weight selection process.** This enhances the flexibility of the generated MoE and ensures fair comparisons in Table 1. In contrast, the process of "a dense model pre-trained on ImageNet-21k and fine-tuned on Pets" bypasses the weight selection step from ViT-S to ViT-T and is not closely related to our MoE Jetpack. We will further modify the description of Table 1 as "Tab. 1 compares the performance of the MoE Jetpack with Dense ViT models trained from scratch, and Dense ViT with pre-trained weights on ImageNet-21k then fine-tuning on downstream datasets with weight selection [1], represented as Dense (21k) in the table." We hope this clarifies your question. Given that our work builds on numerous previous studies, many details may require referencing prior works for a clearer understanding. We apologize for any confusion caused, thank you for raising the question. [1] Xu, Zhiqiu, et al. "Initializing models with larger ones." ICLR 2024 (spotlight).
Response to Authors
Thank you to the authors and all the reviewers for the thoughtful discussions. After several rounds of back-and-forth, most of my concerns have been addressed, and the experimental results now appear sufficiently sound. I have revised my review and score accordingly. I recommend that the authors include the missing experimental setting details in the final version and reconsider how the results are presented to enhance clarity.
Reviewer hQZu has raised several insightful points. However, the discussion has heavily focused on experimental details, potentially diverting attention from the paper's contributions. The motivation of the paper is well-recognized, and the proposed methods are valuable. The main discussion centers on the fairness of the experimental comparisons. In my opinion, the most crucial aspect is how the experiments demonstrate the effectiveness of the proposed methods. The paper achieves this effectively in several ways: - Table 2 shows the SpheroMoE layer more effectively utilizes dense checkpoints than the SoftMoE layer, with improvements from checkpoint recycling. - Table 3 indicates that checkpoint recycling is more flexible and effective than sparse upcycling. - Table 5 demonstrates that the MoE Jetpack uses dense checkpoints to accelerate model convergence and enhance performance on downstream tasks. - Columns 3 and 4 in Table 1 show the combined effects of the SpheroMoE layer and checkpoint recycling, validating their effectiveness. While the authors are familiar with related work on MoE and have ensured experimental fairness, they need to provide clearer explanations of the experimental settings for each table. Including these in future revisions or supplementary materials will help reduce confusion for meticulous reviewers. I hope my comments help clarify the discussion.
Response to Reviewer ysa3
Thank you for further clarifying the discussion. We appreciate your insights into the effectiveness of our experiments and are pleased that the key aspects of our work are clear. We acknowledge your suggestion to provide clearer explanations of the experimental settings for each table. We will ensure that future revisions or supplementary materials include more detailed descriptions to make our methodologies and results fully transparent and comprehensible to all readers.
Summary
This paper introduces MoE Jetpack, a novel method to efficiently fine-tune dense model checkpoints into sparsely activated mixture of experts (MoE) models, reducing the need for extensive data and computational resources. Key techniques include checkpoint recycling and the hyperspherical adaptive MoE (SpheroMoE) layer, which effectively leverage the pre-training costs of existing dense models and optimize the MoE architecture. Experimental results demonstrate significant performance gains and improved convergence speed across various datasets and network architectures.
Strengths
1. The paper provides a thorough explanation of the methods, including various strategies for checkpoint recycling and detailed architecture of the SpheroMoE layer, offering a clear and replicable approach. The pseudo-code in the appendix is exceptionally clear and well-structured. 2. The idea of effectively decomposing pre-trained dense models into experts of varying sizes is very interesting. It significantly enhances the flexibility in utilizing the sunk pre-training costs of existing dense models. 3. The Adaptive Dual-path MoE employs experts of varying sizes within a single MoE model, effectively balancing speed and accuracy. This approach significantly enhances the flexibility and performance of MoE models. 4. The experiments cover a wide range of datasets and model architectures, ensuring that the results are robust and generalizable.
Weaknesses
1. The authors need to carefully check the formatting details of the paper: i) The order of Tables 2 and 3 is incorrect in lines 227 and 262. ii) There is a missing space before the parentheses in Table 1. iii) The numerical values within the parentheses in Table 2 need verification for correctness. 2. The paper misses some related work [2][3][4] and should discuss the proposed method in comparison with these approaches, such as the differences between the proposed Expert Regularization method and the load balancing loss [1] and router Z loss [2]. Are the methods proposed in this paper orthogonal to these existing methods? 3. The additional computational overhead introduced by checkpoint recycling should be detailed and quantified by the authors. [1] Fedus W, Zoph B, Shazeer N. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity[J]. Journal of Machine Learning Research, 2022, 23(120): 1-39. [2] Zoph B, Bello I, Kumar S, et al. St-moe: Designing stable and transferable sparse expert models[J]. arXiv preprint arXiv:2202.08906, 2022. [3] Xue F, Zheng Z, Fu Y, et al. OpenMoE: An Early Effort on Open Mixture-of-Experts Language Models[C]//Forty-first International Conference on Machine Learning. [4] Zhou Y, Du N, Huang Y, et al. Brainformers: Trading simplicity for efficiency[C]//International Conference on Machine Learning. PMLR, 2023: 42531-42542.
Questions
1. I am curious about the differences that might arise from initializing expert weights from models of different sizes. Is a larger model always better?
Rating
7
Confidence
5
Soundness
4
Presentation
4
Contribution
3
Limitations
N/A
Thank you for your detailed responses, which effectively address my concerns. Consequently, I maintain my positive rating. The expanded discussion on router-Z loss, alongside the comprehensive analysis of the computational overhead, will significantly enhance the manuscript's clarity and depth. I look forward to seeing the revised manuscript.
Thank you for your positive feedback and constructive comments. we are pleased that our rebuttal addressed your concerns. We appreciate your suggestions and will incorporate them into the revised version.
Summary
This paper introduces MoE Jetpack, a novel method for fine-tuning pre-trained dense model checkpoints into mixture of experts (MoE) models. It leverages checkpoint recycling to accelerate convergence and enhance accuracy, and incorporates a hyperspherical adaptive MoE (SpheroMoE) layer for optimized integration and performance. Extensive evaluations show that MoE Jetpack not only speeds up the fine-tuning process but also achieves higher accuracy across various datasets while maintaining computational efficiency.
Strengths
1. Tables and figures are informative, and the paper is generally well written. 2. The technical contributions of the paper are substantial. 3. The experiments comprehensively validate the effectiveness of the proposed methods.
Weaknesses
1. The work consumes extensive computational resources to validate a complex algorithm on small-scale datasets. Specifically, all experiments presented required 3300 GPU hours for training and a total of 8000 GPU hours for exploratory purposes. However, the datasets primarily showcased, aside from ImageNet, are of small scale and outdated. Given the current emphasis on large models, these small datasets are not suitable for substantiating the claims of this paper. It is recommended to conduct more extensive experiments on ImageNet or similar large-scale datasets to validate scalability and justify the extensive use of pre-trained checkpoints. 2. From an application perspective, the effectiveness of utilizing checkpoints in practical scenarios is questionable. For instance, the core result of the ViT-based models, where V-JetMoE-T with 264M parameters only achieves 79.9% accuracy, suggests that using pre-trained weights for MoE may not be as parameter-efficient as expected. Instead, using comparable computation models, like those on the ImageNet performance leaderboard, might be more advantageous. This critique underscores the need to demonstrate that utilizing pre-trained weights for MoE can surpass the performance of similarly scaled models trained conventionally.
Questions
Please refer to the concerns raised in the weaknesses section. It is essential to provide more experiments to demonstrate that MoE models using pre-trained weights deliver significantly better performance, and to verify that the proposed method has sufficient scalability.
Rating
5
Confidence
3
Soundness
3
Presentation
3
Contribution
3
Limitations
The paper has discussed the limitations in the conclusion, specifically noting that the method is only applicable to visual tasks. There is no need to consider potential negative societal impacts for this work.
Summary
The paper proposes several ideas to convert dense checkpoints into MoEs for vision tasks. In particular, four alternatives are presented to reuse checkpoints of dense models to be use as an initial checkpoint for an MoE model for a later fine-tuning phase. Of the four alternatives presented, Importance-Based Weight Sampling is the one that works the best. Most importantly, all of them are reported to work better than Sparse Upcycling, at least when evaluated on ImageNet (see Table 3). The second novelty that the paper presents is a variant of the SoftMoE router: the SpheroMoE router. The SpheroMoE has an additional projection on the inputs, and places the L2 normalization in different places, compared to SoftMoE. It also has a temperature parameter that is decayed over time. Similar to the SoftMoE router, the SpheroMoE computes slots as a weighted average of the input tokens (similar to attention). Last but not least, the paper introduces the Adaptive Dual-Path structure in the MoE blocks. The experts are divided in two sets: a set of light experts that process lots of slots, and a set of heavy experts that process fewer slots. The paper presents results on variety of common vision tasks, such as ImageNet-1k (a.k.a. ILSVRC2012), CIFAR-10/100, Food-101, etc. Sometimes the models are trained from scratch on these datasets, and others an initial (dense) checkpoint pre-trained on ImageNet-21k is used.
Strengths
- The different ideas presented to initialize a MoE from a dense checkpoint are new and are shown to work better as a previously proposed approach: Sparse Upcycling. The explanation of the methods is well-detailed in the paper. - A quite detailed implementation of the SpheroMoE router is given in the appendix, which will certainly widen the reproducibility and potential adoption of the paper. - Most of the presented ideas have been ablated on different experiments. In particular, the choice of Importance-Based Weight Sampling is justified in Table 3, experimenting on ImageNet, which is a quite rich dataset. The use of the SpheroMoE router is well justified in Table 2, and is showed to be superior to the SoftMoE router on 3 datasets (ImageNet, CIFAR-100, Flowers). - Overall, it's shown that using MoE Jetpack not only significantly accelerates the training on ImageNet-1k and CIFAR-100, compared to training from scratch, but also the models achieve a higher accuracy which would probably not be achieved from scratch, or at least not in a reasonable amount of time, in the case of CIFAR-100.
Weaknesses
- All experiments are done with quite modest backbones: Tiny (T) and Small (S) vision tranformer backbones. The Sparse Upcycling paper, that the authors use a baseline, observes that the benefits of initializing from a dense checkpoint are much smaller when using larger backbones (especially when evaluating transfer to small downstream tasks), see Figure 3 and 4 from https://arxiv.org/pdf/2212.05055. So, it's not clear at all wheather the proposed improvements will make a difference at much bigger scales. - SpheroMoE has far more experts that SoftMoE. It's not clear wheather the benefits come from the improvements or by using finer-grained & more experts, which has been proposed before (for instance, in https://arxiv.org/abs/2401.06066 and https://arxiv.org/abs/2402.07871). - The results for SoftMoE in Figure 1b and Table 1 are for models trained from scratch. Since a distinction is made for ViT on whether it's trained from scratch or pre-trained on ImageNet-21k, perhaps it would be fairer to present this distinction for SoftMoE as well in these main figure and table. This ablation is presented in Table 2, but only for 3 datasets. - Figure 2b and Equation 10 (and algorithm 1 in the appendix), do not quite match in one aspect: the L2-normalized X is also used to compute the content of the input slots of each expert, and not the original X as Figure 2b seems to depict. - SpheroMoE has a few differences with SoftMoE. It would be much better if the individual differences were ablated. For instance, what's the effect of the additional (linear?) projection on X used to compute the combine and dispatch weights? What's the difference between having a decayed temperature vs. having a learned temperature (or scale) as SoftMoE does?
Questions
- In Algorithm 1, the meaning of `K_project` and `inherit_layer_norm` are not clearly explained. Please, do so. In particular, what's the output size of the `K_project`? Is it the same as the embedding size of X? - In that regard, assuming that `inherit_layer_norm` is doing some sort of LayerNorm with shared parameters, LayerNorm and L2-normalization are very similar, except for a $\sqrt{n}$ factor and the additional $\alpha$ or scale parameters of LayerNorm (and potentially the biases too). How was the use and positioning of the additional L2-norm decided?
Rating
6
Confidence
5
Soundness
2
Presentation
3
Contribution
3
Limitations
No negative societal impact is particular to this work, in my opinion.
Comment after reading the response
First, I acknowledge the authors' contribution to checkpoint recycling, so I have raised my score accordingly. However, my concerns regarding the fairness of the experiments remain unresolved. I am not the only reviewer who doubts the fairness of the experiments (see R7BZ9 Soundness, Weakness 3). While the authors' proposed method can transfer a dense checkpoint to an MoE without additional computational cost, obtaining the dense checkpoint itself is not free. For instance, in Table 1, Row 8 for the "Pets" dataset, "Soft MoE" only utilizes the Pets data, "Dense(21k)" only utilizes the ImageNet-21k data, and "Dense" only utilizes the Pets data. However, the MoE Jetpack utilizes both ImageNet-21k and Pets data (since the dense checkpoint it initializes from is pre-trained on ImageNet-21k). The MoE Jetpack involves two phases: initialization from a dense checkpoint and fine-tuning (clearly described in lines 83-85). At the very least, the authors should include a baseline of a dense model pre-trained on ImageNet-21k and then fine-tuned on Pets. Reporting "84.3(+38.8)" in Table 1 is misleading. The increase of 38.8 is calculated as 84.3 (MoE Jetpack initialized from a dense model pre-trained on ImageNet-21k, then fine-tuned on Pets) - 45.5 (Soft MoE trained from scratch only on Pets), which exaggerates the improvement. If a stronger dense checkpoint (e.g., pre-trained on CC12M) were used, the gap would be even larger, but such a comparison would not be meaningful. For the same reasons, the comparison of convergence speed is also questionable. It appears the authors are comparing the convergence speed of a model trained from scratch (Random Init MoE) to a pre-trained model (V-JetMoE-T). I look forward to the authors' reply.
Acknowledgment of Reviewer hQZu's Extensive and Insightful Discussion
Thank you for dedicating your valuable time to providing thoughtful feedback on our paper. We are pleased that most of your concerns have been addressed, and we greatly appreciate your revised review. We acknowledge your recommendation to include the missing experimental setting details and to reconsider how the results are presented for enhanced clarity. We will carefully incorporate these suggestions in the final version to improve the overall quality and transparency of our work.
Decision
Accept (poster)