Solving Token Gradient Conflict in Mixture-of-Experts for Large Vision-Language Model

Thank you for your review and your feedback on our paper. We address your questions and concerns below.

Comment 1: Clarification on Gradient Conflict Hypothesis. The authors assert that the learning of conflicting tokens increases the loss of most other tokens within the same expert. However, it is counterintuitive since conflicting tokens should be relatively few, making the directions of parameter optimization align with the global average gradient of all tokens in the expert. Clarifying this assumption would improve the paper's clarity.

Response: Thank you for pointing out this. "The learning of conflicting tokens increases the loss of most other tokens within the same expert" is not precise enough. We have modified it as "The learning of conflicting tokens increases the average loss of tokens within their current experts" in the manuscript. We further explain this claim: the gradient generated by conflicting token is in the opposite direction of the average gradient. Updating parameters along the direction of the conflicting token gradient is equivalent to updating parameters in the direction of the average gradient's ascent, which leads to an increase in the average loss. For example, when an expert has conflicting tokens, the average loss decreases from 0.3 to 0.25 after updating the parameters; after removing conflicting tokens, the average loss may decrease from 0.3 to 0.2 after updating the parameters.

Comment 2: Handling Conflicting Tokens Across Experts: The paper suggests redistributing conflicting tokens to different experts. However, it is unclear how the method ensures that these conflicting tokens do not cause conflicts with tokens in other experts.

Response: It is challenging to determine the optimal expert of a conflicting token during redistributing. We have recognized this limitation and presented it in the Limitation part. However, importantly, even if some tokens remain conflicting after redistributing, it is unlikely to be the case for all; STGC cannot guarantee the total elimination of conflicting tokens, but it can reduce the number of conflicting tokens. In Figure 3, Figure 7, and Figure 8 of the manuscript, it has been confirmed that STGC is beneficial for improving gradient consistency and reducing conflicting tokens. In the future, exploring whether it is possible to determine the optimal expert of a token from the optimization perspective may be an intriguing direction.

Comment 3: Questioning the Gating Network’s Learning Capabilities: It remains unclear why the gating network fails to learn an optimal distribution for the conflicting tokens. Given that each expert should ideally handle tokens with similar characteristics, why would the gating network assign such divergent tokens to the same expert?

Response: Similar features do not necessarily imply similar gradients. When tokens have highly similar features (similar tokens), the router will assign them to one expert. However, they may have dissimilar gradients. STGC can be understood as learning token features based on token-level gradient relationships. After using STGC, the features of tokens become less similar when they have dissimilar (conflicting) gradients, for being routed to different experts. To verify this, we compute the Pearson correlation coefficient between feature similarities and gradient similarities to measure whether feature relationships reveal gradient relationships. As shown in the below table, the experimental results indicate that token features and gradients have a higher correlation after using STGC. For more details please refer to Section 6.8 and Table 12 of the revised manuscript.

Thank you for your review and your feedback on our paper. We address your questions and concerns below.

Comment 1: The paper primarily focuses on the empirical validation of the STGC method. While the experimental results are promising, there could be a more in-depth theoretical analysis to understand the underlying principles and limitations of token gradient conflicts and how STGC addresses them.

Response: Thanks for your advice. We have conducted a brief theoretical analysis in Section 6.9 of the revised manuscript.

Suppose there are tokens $t_n$ and $t_n'$, which pass through the same expert and generate gradients $\mathbf{g_n}$ and $\mathbf{g_n'}$ on that expert. $cos(\phi_{nn'})$ is the cosine similarity between gradients $\mathbf{g_n}$ and $\mathbf{g_n'}$. Token gradient conflicts mean the cosine similarity $cos(\phi_{nn'}) < 0$, We can prove that $cos(\phi_{nn'}) < 0$ potentially leads to an increase in the loss. We can prove that when $cos(\phi_{nn'})$>0, the loss decreases strictly, which can reach the optimal value. The goal of STGC is to increase $cos(\phi_{nn'})$ to satisfy the condition $cos(\phi_{nn'})$>0. Figure 3 and Figure 8 verify that STGC increases $cos(\phi_{nn'})$.

Comment 2: The paper mentions the reduction in inference cost but does not discuss the potential increase in training cost due to the additional computations required for token-level gradient analysis.

Response: We have conducted a thorough computational overhead analysis in Section 6.6 and Table 10 of the revised manuscript.

Training memory overhead. We compute the memory overhead by the size of the stored gradient. As stated in Section 6.1, for each token, we only store the gradient it produces on the bias within the experts. Thus, the theoretical memory overhead is merely about 0.29 GB on StableLM-1.6B and 0.38 GB on Phi2-2.7B, which is negligible during the training of MoE-based LVLMs. For the specific calculation please refer to Section 6.6.

Training time overhead. We directly report the train_samples_per_second and train_steps_per_second recorded in "trainer_state.json" after training. MoE-LLaVA performs a forward pass, followed by a backward pass to update parameters. As Section 6.1 mentioned, MoE-LLaVA+STGC freezes parameters except for the bias within the experts, performs a forward pass, followed by a backward pass to compute token-level gradients; then, MoE-LLaVA+STGC unfreeze parameters and performs a backward pass to update parameters. Thus, the main time overhead results from "the computation of token-level gradients". As shown in the below tables, we have reduced the additional time overhead to about 20% through some engineering tricks (e.g., only computing the token-level gradient on the bias and freezing parameters that do not require gradient computation). Some other methods may further speed up STGC, such as using STGC only on even iterations or applying STGC to only half of the MoE layers. We will further explore these experiments in the future.

Computational overhead analysis on StableLM-1.6B:

(STGC-full means computing token-level gradients on the weight. Meanwhile, STGC is the used scheme, only computing token-level gradients on the bias.)

Computational overhead analysis on Phi2-2.7B:

We would like to update the experimental result below. Sorry for the late.

Comment 3: The paper demonstrates the effectiveness of STGC on vision-language tasks. However, it is not clear how well these findings generalize to other domains or tasks outside of vision-language models. Further experiments in diverse domains could strengthen the paper's claims.

Response: Thanks for your advice. Theoretically, the deployment of STGC is not constrained to specific task types. Thus, to validate the generalization of STGC, we extend its use to language tasks. We have supplemented the study on language tasks on Section 6.10 and Table 13. We follow DYNMOE [1] to apply the MoE framework for language tasks and we add the proposed STGC to the MoE. As shown in the below table, the experimental results show the effectiveness of STGC on language tasks. For more details please refer to Section 6.10 and Table 13.

$*$ means the re-implemented results.

[1] Guo, Yongxin, et al. "Dynamic mixture of experts: An auto-tuning approach for efficient transformer models." arXiv preprint arXiv:2405.14297 (2024).

Comment 3: Conflict Elimination Loss: Are there other forms that were considered, or insights into why this particular expression works would help.

Response: Thanks for your advice. We have supplemented the discussion about conflict elimination loss design in Section 6.4 and Table 9 of the revised manuscript. In specific, the goal of the conflict elimination loss is to reduce the routing score $p_{\text{moe}}(t_n)$ of the conflicting token $t_n$ on its current expert. We discuss different designs for the conflict elimination loss: $(i)$ MSE-like: Simply setting the routing score $p_{\text{moe}}(t_n)$ to the minimum. $(ii)$ CE-like: Utilizing the inverted routing score along with cross-entropy loss. Our motivation for taking the inverted routing score is to minimize the routing score of the token on its current expert. As shown in the below table, the CE-like loss performs better. The reason may be that, although the optimization direction of the MSE-like loss is consistent with the CE-like loss, the optimization speed of the CE loss is superior to the MSE loss [1,2].

[1] Zhang, Zhilu, and Mert Sabuncu. "Generalized cross entropy loss for training deep neural networks with noisy labels." Advances in neural information processing systems 31 (2018).

[2] Wang, Yisen, et al. "Symmetric cross entropy for robust learning with noisy labels." Proceedings of the IEEE/CVF international conference on computer vision. 2019.

Thank you for your review and your feedback on our paper. We address your questions and concerns below.

Comment 1: Practical Implementation: A quantitative measure of the memory overhead reduction. Performance difference on an actual dataset could make the claim of practicality even stronger.

Response: Thanks for your advice. We have conducted a detailed memory overhead analysis in Section 6.6 and Table 10 of the revised manuscript. We compute the memory overhead by the size of the stored gradient. As stated in Section 6.1, for each token, we only store the gradient it produces on the bias within the experts. Thus, the theoretical memory overhead is merely about 0.29 GB on StableLM-1.6B and 0.38 GB on Phi2-2.7B, which is negligible during the training of MoE-based LVLMs.

If storing the token-level gradients on each weight, the required storage overhead is 773 GB for StableLM-1.6B and 1562 GB for Phi2-2.7B, which is amazingly large. Thus, it is impossible to use the token-level gradient on each weight to conduct experiments.

For the specific calculation please refer to Section 6.6.

Computational overhead analysis on StableLM-1.6B:

(STGC-full means computing token-level gradients on the weight. Meanwhile, STGC is the used scheme, only computing token-level gradients on the bias.)

Computational overhead analysis on Phi2-2.7B:

Comment 2: Token Conflicts after STGC: Along with the mean gradient consistency, it would be good to report the std deviations as well. It might also be worth tracking the proportion of total tokens that stay conflicting over the course of training.

Response: Thanks for your advice. We have reported the std deviation of gradient consistency and the ratio of conflicting tokens in Section 6.5 and Figure 8 of the revised manuscript. We can observe: $(i)$ as the conflict elimination loss decreases, the gradient consistency increases, and the ratio of conflicting tokens decreases. This means that STGC effectively reduces gradient conflicts within the expert and reduces the count of conflicting tokens. $(ii)$ The gradient consistency std increases, meaning that the difference of gradient consistency of different experts enlarges. We speculate that this is due to the different rates at which gradient consistency increases at various layers. For example, Figure 3 (b) indicates that deeper layers have a faster gradient consistency increase rate after adding STGC. Furthermore, we analyze the layers in which conflicting tokens appear in Figure 9 of the revised manuscript. We find that most conflicting tokens emerge in the shallow layers after learning, further verifying the finding in Figure 3 (b) and Figure 8.

Thank you for all your efforts in answering my concerns, the current version of the manuscript looks much more self-contained. I am happy to raise my score to 8.

It seems that the eq (8) is missing a negative sign. If I want to minimize the loss CEL, what I need to do is to minimize the logit of $z_{moe}(t_n)$ for conflicting token $t_n$ instead of minimizing $-z_{moe}(t_n)$.

The standard CE loss increases the score of a sample on the ground-truth class. In our case, we want to reduce the score of the sample on $id_{\text{moe},n}$. Therefore, we add a negative sign, making our loss increase with $z_{\text{moe}}^\prime(t_n)$, which in turn decreases $z_{\text{moe}}(t_n)$.

Thank you for your review and your feedback on our paper. We address your questions and concerns below.

Comment 1: It would be better to illustrate the distribution of cosine similarity between gradients of all tokens and the averaged gradient for better understanding.

Response: Thanks for your advice. We have supplemented the distribution of cosine similarity between the gradient of each token and the averaged gradient in Section 6.5 and Figure 7 of the revised manuscript. We find that STGC can effectively increase the cosine similarity between the gradients of tokens within an expert and the average gradient.

Comment 2: Since the average gradient is dynamically changing during training, is it possible that some tokens do not conflict with a specific expert but become conflicting after several steps of training?

Response: It is possible that conflicting tokens are dynamic. Hence, rather than identifying conflicting tokens before training, STGC dynamically identifies conflicting tokens during training and constrains their routing through conflict elimination loss. Experimental results in Figure 7 and Figure 8 of the manuscript have indicated that STGC effectively reduces the number of conflicting tokens.

Comment 3: Why only the Large Vision-Language Model (LVLM) setting is considered?

Response: Due to limited time, we did not apply STGC to other tasks in the original manuscript. Theoretically, the deployment of STGC is not restricted to task types. To verify this, we extend STGC to language tasks. We are still preparing the experimental results of this part, and we will report them within the next day or two.

Comment 4: Some minor points.

Response: Thanks for your advice. We have fixed all the points you mentioned except for TL;DR as we cannot change it now.

Thank you for your review and your feedback on our paper. We address your questions and concerns below.

Comment 1: Visualization on the routing mechanism of the trained model should be studied to provide insights into the trained model.

Response: Thanks for your advice. We have presented a visualization of the routing mechanism of the trained model in Section 6.3 and Figure 5 of the revised manuscript. Specifically, we follow MoE-LLaVA to obtain the distribution of expert loading and the visualization of the activated pathways. The distribution of expert loading examines the expert use frequency for all tokens. Activated pathways examine the behavior of experts at the token level: this visualization tool tracts the activated pathways of all tokens on validation datasets; given all activated pathways, the visualization tool employs PCA to obtain the top-10 pathways.

As shown in Figure 5 of the manuscript, we find that $(i)$ STGC benefits the expert load balance. A possible reason is that the expert load imbalance means that many tokens are routed to an expert, significantly increasing the possibility that tokens have gradient conflicts. After adding the STGC, some tokens are moved from the "crowded" expert (many tokens) to the "empty" expert (few tokens). This validates that STGC would effectively utilize each expert, instead of collapsing into only using one expert. $(ii)$ The activated pathways are significantly different for SQA, TextVQA, and MMBench. This implies that although the distribution of expert load across different datasets is similar, the token routing behavior is still significantly different among datasets, i.e., different tokens have been assigned to various experts.

Comment 2: It seems that the optimization process of the proposed CEL is not very stable (as in Figure 3(a)), the reviewer wonders whether this would be a problem when scaling up to larger scale.

Response: The CEL appears unstable possibly because too few steps are displayed. We present the actual CEL curve during the training process (5197 steps) in Figure 6 of the revised manuscript and find that the CEL is convergent.