Solving Token Gradient Conflict in Mixture-of-Experts for Large Vision-Language Model
Authors: Longrong Yang
Dong Shen
Chaoxiang Cai
Fan Yang
Tingting Gao
Di ZHANG
Xi Li
Venue: ICLR 2025 Poster (2025)
Keywords: Large Vision-Language Model (LVLM)
Mixture-of-Expert (MoE)
token-level gradient
conflicting token
Primary area: foundation or frontier models, including LLMs
TL;DR: We propose using token-level gradients to define conflicting tokens for modeling data interference in the LVLM, and design a novel loss to optimize the token routing to solve data interference.

Abstract: The Mixture-of-Experts (MoE) has gained increasing attention in studying Large Vision-Language Models (LVLMs). It uses a sparse model to replace the dense model, achieving comparable performance while activating fewer parameters during inference, thus significantly reducing the inference cost. Existing MoE methods in LVLM encourage different experts to specialize in different tokens, and they usually employ a router to predict the routing of each token. However, the router is not optimized concerning distinct parameter optimization directions generated from tokens within an expert. This may lead to severe interference between tokens within an expert. To address this problem, we propose to use the token-level gradient analysis to Solving Token Gradient Conflict (STGC) in this paper. Specifically, we first use token-level gradients to identify conflicting tokens in experts. After that, we add a regularization loss tailored to encourage conflicting tokens routing from their current experts to other experts, for reducing interference between tokens within an expert. Our method can serve as a plug-in for diverse LVLM methods, and extensive experimental results demonstrate its effectiveness. demonstrate its effectiveness. 
The code will be publicly available at https://github.com/longrongyang/STGC.

Published as a conference paper at ICLR 2025
SOLVING TOKEN GRADIENT CONFLICT IN MIXTURE- OF-EXPERTS FOR LARGE VISION-LANGUAGE MODEL
Longrong Yang'. Dong Shen?, Chaoxiang Cai², Fan Yang', Tingting Gao?, Di Zhang?, Xi Li!-t
CColleee of Computer Science and Technology. Zhejiang University
2 KKuaihhou Technology
School of Software Technology, Zhejiang University
ABSTRACT
The Mixture-of-Fxperts (MoE) has gained increasing attention in studying Large Vision-Language Models (LVLMs). uses a sparse model to replace the dense model, achieving comparable performance while activating fewer parameters during inference, thus significantly reducing the inference cost. Existing MoE methods in LVLM encourage different experts to specialize in different tokens, and they usually employ a router to predict the routing of each token. How- ever, the router is not optimized concerning distinct parameter optimization di- rections generated from tokens within an expert. This may lead to severe in- terference between tokens within an expert. To address this problem, we pro- pose to use the token-level gradient analysis to Solving Token Gradient Conflict (STGC) in this paper. Specifically, we first use token-level gradients identify conflicting tokens in experts. After that, we add regularization loss tailored to encourage conflicting tokens routing from their current experts to other ex- perts, for reducing interference between tokens within an expert. Our method can serve as a plug-in for diverse LVLM methods, and extensive experimen- tal results demonstrate its effectiveness. The code will be publicly available at https://gite www cow./onggongyyng///ann..
INTRODUCTION
Large Vision-Language Models (LVLMs) have recently demonstrated significant advancements by integrating visual processing modules into Large Language Models (LLMs). Many recent LVLMs (Zhang et al., 2023a; Bai et al., 2023b; Zhang et al., 2023b; Zhao et al., 2023; Chen et al., 2023b) show that large model size and large dataset size are significant enhance intelligence, i.e.. the scaling law. Even when the model size is sufficiently large, models exhibit "Emergent Abilities". Thus, a series of studies (Li al., 2022; Dai et al., 2023; Liu al., 2023b) have expanded the model size of LVLMs to 13 billion parameters, leading to state-of-the-ar performance on various tasks.
Under realistic applications, deploying such large models requires considerable computational re- sources, making inference extremely expensive. For reducing the inference cost, a popular solution using the Mfxture-of-Experts (MoE) architecture, replacing the FFN layer with multiple experts, which has been verified by many works (Fedus et al., 2022; Zoph et al., 2022; Komatsuzaki et al., 2022) to achieve comparable performance with dense models when activating fewer parameters.
With multiple experts in the MoE, fundamental problem is the routing of tokens. To route tokens to different experts, existing MoE works (Lin et al., 2024; Dai et al., 2024) typically train a router, such linear layer, to predict the probability of each token dispatched to different experts. The tokens are then dispatched to the experts with the Top-k predicted probability. However, a natural problem arises: What is the optimization goal of the router to dispatch tokens?
Given the wide variety of data used in LVLMs, we think that a critical goal of token routing to vari- ous experts is to reduce interference between diverse data. Some related LoRA-MoE. studies (Chen et al., 2023d; Gou et al., 2023; Shen et al., 2024; Liu & Luo, 2024; Zhou et al., 2024) have also con- ducted preliminary explorations from this perspective, usually modeling data interference through
"ooresppnniigg author is Xi L..

---

Published as a conference paper at ICLR 2025
Besellne STGC
0.3
Expert
Expert
Expart
Gradiest
0.2
MoE
Consistancy
oosine
0.1
ofthr bakcs
Router
withothritiom
simalatty)
Beselime
conflictnng Jolv
Token
Gradlennt
120
Trairing Step
(a) Our goel: reduce gradient conflicts of tokens within an expert
(b) Gradient consistency of tokens within an expat before and after using STGC
Figure 1: (a) In this work, we aim to solve data interference by adjusting token routing to reduce gradient conflicts. (b) We present statistics regarding gradient consistency (the mean cosine simi- larity between gradients of all tokens within an expert). In experiments, we fed one sample into the LVLM per device for each forward pass. The baseline LVLM MoB-LLaVA (Lin et al., 2024).
sample-level instruction features or embeddings. For instance, MoCLE (Gou et al., 2023) performs sample-level clustering instruction embedding, solving data interference by constraining samples of distinct cluster centers to pass different experts. However, even though samples have similar in- struction embeddings, they may generate distinct parameter optimization directions due to distinct targets, so embedding- baabedding-baeedd routing is still risky for optimization interference within an expert. Moreover, since the routing typically the token level (Lin al., 2024; Dai al., 2024), existing sample-level methods may struggle with interference between tokens (e.g., visual and text tokens) within sample. Fortunately. the gradient can directly indicate the direction of parameter optimiza- tion. Thus, this work aims to model data interference through the lens of token-level gradients. As shown in Figure (a), our basic idea to optimize the router to reduce gradient conflicts between tokens within an expert, for solving data interference under complex and real-world scenarios.
To this end, we propose employ the token-level gradient analysis to design novel regulariza- tion loss for Solving Token Gradient Conflict (STGC). The STGC is proposed to answer two key questions: (i) How to define conflicting tokens? After forwarding batch of data, we perform backward pass capture the token-level gradients each expert without updating any model pa- rameters. Within an expert, we compute the average gradient of all tokens, representing the holistic optimization direction of the expert. Then, tokens are identified conflicting tokens if their gradi- ents exhibit a negative cosine similarity to the average gradient. These conflicting tokens (outliers) harm the learning of the expert. (it) How solve conflicting tokens? After identifying confticting tokens, we design a conflict elimination loss to optimize the router to encourage conflicting tokens routing away from their current experts. As shown in Figure 1 (b), the STGC enhances gradient consistency between tokens within an expert, i.e., reduces gradient conflicts of tokens.
In conclusion, our contribution can be summarized
Beyond relying on sample-level embedding cues, we propose using token-level gradients define conflicting tokens for modeling data interference the LVLMs.
We propose novel conflict elimination loss to optimize token routing to solve gradient conflicts, making parameter optimization directions generated from tokens within an expert more consistent. This also prompts the further specialization of experts in the MoE.
Designed as plug-in, our method can be seamlessly integrated into existing MoE-based LVLMs. Extensive experiments have confirmed its effectiveness.
RELATED WORKS
2.1
LARGE VISION-LANGUAGGE MODEL
Large Language Models (LLMs) have demonstrated strong instruction following and generaliza- tion capabilities. LLMSs can only process textual information, while real-world applications require
2

---

Published as a conference paper at ICLR 2025
models to process visual information, e.g., object detection (Yang al., 2025) or instance segmen- tation (Yang et al., 2020; 2022). To incorporate visual information, Large Vision-Language Models (LVLMs) such as GPT-4 and LLaVA utilize frozen visual encoders and trainable visual projectors to integrate visual data into LLMs. Recent works have focused on improving performance through two types of methods. The first type optimizes training strategies, e.g., (Bai al., 2023b; Chen et al., 2023a). Most works belong to the second type, focusing on enhancing visual components, including expanding visual instruction-tuning datasets (Liu et al., 2023a; Zhang et al., 2023b), im- proving image encoders (Chen et al., 2023e; Bai et al., 2023b), and aligning the input and projection layers (Lin al., 2023; Cha al., 2023; Alayrac al., 2022; Dai et al., 2023; Ye et al., 2023; Zhao et al., 2023). These efforts, particularly the expansion of visual instruction-tuning datasets and the increase in model scales, have significantly enhanced the visual understanding abilities of LVLMs.
2.2 MIXTURE-OF EX-PERT (MoE)
Efficient training and inference, e.g., (Yang et al., 2023; 2024), are crucial for the deployment of large models. The Mixture-of-Experts (MoE) is hybrid model consisting of multiple sub-models known experts and has shown potential in reducing the inference cost (Shazeer et al., 2017). The critical concept of MoE lies in using a router to determine the token set that each expert handles, aiming reduce interference between tokens from diverse data. Early MoE. works have utilized the hard routing mode, where each expert is typically assigned a specific role. For example, a series of works (Bao al., 2022; Satar et al., 2022: Long et al., 2023; Wang et al., 2022; Shen al., 2023b) consider vision and language gaps in multi-modal data (Liang et al., 2022), decoupling experts by modal type and assigning a specific role each expert. The critical property of hard routers is that they eliminate the need to learn the routing. The hard routing has also been widely applied in task-specific MoEs (Kudugunta et al., 2021; Zhu et al., 2022; Li et al., 2023c; Ma et al., 2023).
Then, soft routers enable a dynamic allocation of tokens between different experts. Recent works have mainly focused on soft routers egg., (Lepikhin et al., 2020; Fedus et al., 2022; Zoph et al., 2022: Komatsuzaki et al., 2022; Shen et al., 2023a; Zadouri et al., 2023: Puigcerver et al., 2023; Chen et al., 2023c; Chalapathi et al., 2024; Zhong et al., 2024). For LVLMs, MoE-LLaVA (Lin et al., 2024). Uni-MoE (Li et al., 2024b), and MoAI (Lee et al., 2024) propose to employ MoE to empower LVLMs. Moe (Dai et al., 2024) and QwenMoE (Bai et al., 2023a) further segment experts by splitting the FFN hidden dimension to achieve further specialization. CuMo (Li et al., 2024a) designs MoE. for both the vision encoder and the MLP connector. DYNMOE (Guo al.,, 2024) and AdaMoLEE (Liu & Luo, 2024) enable each token to determine the number of experts to activate dynamically. Some recent works have claimed that the MoF structure itself is suitable for handling data interference, so they address data interference by adding LoRA-MoE on fixed FFN (Chen et al., 2023d; Gou et al., 2023; Wu et al., 2024; Chen et al., 2024; Shen et al., 2024; Liu & Luo, 2024: Zhou et al., 2024). LLaVA-MoLE (Chen et al., 2024) resembles MoE-LLaVA, using token representation to predict the routing scores. LoRA-MoE (Chen et al., 2023d) uses instance- level instruction token average representation to predict the routing scores. MoME (Shen et al., 2024) uses instruction embeddings to infer different visual representations' weightings and compute their weighted sum. Then, MoCLE (Gou et al., 2023) clusters instruction embeddings of samples and use the cluster-related learnable embeddings to predict the routing. MoLA (Zhou et al., 2024) constraints the router based on the sample task. These methods usually operate at the sample level, which makes it challenging to address interference between different tokens within a sample. This work utilizes token-level gradients to solve data interference in the MoE.
3 METHODOLOGY
3.1 OVERVIEW
Large Vision-Language Model: A Large Vision-Language Model (LVLM) aims to effectively integrate the capabilities of the pre-trained LLM and a visual model. Specifically, the input of the vision encoder is an image v E H x3 where H and W are its height and width, and its output
is a visual token sequence 2- = 21, 22, ................................ _2p] E IR. PxC, where P is the sequence length of visual tokens. Then, visual projection layer is used to map Z IR PxC to V PxD, where D represents the hidden size of Large Language Model (LLM). Besides, the instruction text projected
as instruction text tokens T = [4, 12 [N]] E NDD where N represents the sequence length
3

---

Published as a conference paper at ICLR 2025
(a) Conflicting Token Identification
(b) Conflict Elimination Loss LCEL
average gradient
Experts
N
Routing Soore
Add Fees to Decruase
Routing Soore
Omoon
mmaan
Router
Sh
Expert index
N Expert index
Conflicting
Input Tokens
En
En
Cu'
bn is conflleting token
From Curren Expert to Another Expert
Figure 2: Our pipeline. (a) Conflicting Token Identification. When the gradient of a token has a sufficiently low cosine similarity the average gradient of its assigned expert, this token is marked as a conflicting token (an outlier for the expert). (b) Conflict Elimination Loss. We propose a loss aimed encouraging the routing of conflicting tokens from their current experts other experts.
of instruction text tokens. This model consists of stacked multi-head self-attention (MSA) and feed- forward neural networks (FFN), with layer normalization (LN) and residual connections typically used within each block:
Xo = [U1, U2, .. up,, ,,, 12, ,,,)
X MSA(LN(x(-I)) + xx-1, LL, L}.
(2)
X FFF((N((X))) () + t
(3)
where is the layer number of LLM. The LVLM model generates an output text sequence y [y1, 1/,... .KK] E K by progressively generating each element, where K represents the total length of the output text sequence. Then, the outputs are optimized through a generative loss in an ruto-segressive manner (Liu et al., 2023c). The loss (the main loss) is formulated as:
K
maan (0)
Σ log (y (V.Tyec;)) 1=1
(4)
where Yes indicates output text tokens 1y1,82, -1 (i 22) and no output text tokens when ==1 indicates the trainable parameters. The main loss for token L, is abbreviated as (0).
32.
MoF: The Mixture-of-Fixpee -xxeettttft--ttttte---e (MoH) layer is used to replace the FFN layer in this work, similar to (Lin et al., 2024; Dai et al., 2024). MoE layer consists of multiple FFNs, each representing an expert, i.e., & = [e1, e2, egl. where E is the number of experts. For one token fn, the router is typically a linear layer that predicts its probability of being assigned to each expert:
(In)
Prnce (t.));
(5)
where 2moe (tn) = tn W and Pmoe (t.)) v the routing score of fn for the i-th expert. The matrix W I 1x F. represents the router parameters. We calculate weighted of the outputs from the Top-k experts with the highest softmax probabilities:
Wmoe (In)
== @omme
(6)
MoE(tn) Σ Wmoc ei(tn);
where (t,) represents the weight of the i-th expert for tn. and e,(t,) is the output of the i-th expert. We express (0) as C. (0,,, 8'), where 6,, denotes the i-th expert, and ø' represents all parameters except for 0zg- The visual token u, is the same l. when passing the MoF.
4

---

Published as a conference paper at ICLR 2025
In this work, aim to token-level gradients model and solve data in- terference. First, tokens act as the basic unit during the forward pass of the LLM, so we can calculate the gradient generated from each token the expert parameters, i.e,, token-level gradient. Then, as illustrated in Figure 2, our method consists of two steps: (i) we compute token-level gradients and identify conflicting tokens by token-level gradients. During this process, do not update any model parameters. (ii) We add a regularization loss to the main loss to eliminate conflicting tokens. The details of these modules will be introduced in the subsequent sections.
Our Method STGC:
3.2 CONFLICTING TOKEN IDENTIFICATION
Data interference in LVLM generated from interference between tokens within an expert. Some recent works model data interference through instruction embeddings (Chen et al., 2023d; Shen et al., 2024; Gou et al., 2023), i.e., interference oceurs when samples have distinct instruction em- beddings. Alternatively, decisions be made based the specific task associated with each sample (Zhou et al., 2024). However, these works have two main limitations: (i) features and labels jointly influence the parameter optimization direction, but these works rely only one of these two factors. (ii) These work operate at the sample level, whereas routing all tokens within a sample to the same expert does not solve interference between tokens (e-g. visual and text tokens may be in interference) within the sample. To address the issues, we propose using token-level gradients, which can accurately depict the directions of parameter optimization at the token level, to identify interference between tokens within an expert.
First, we introduce the negative impact of the gradient conflict. Without loss of generality, we discuss two distinct instruction tokens, tn and tny, as shown in Figure 2 (a). Assume that both denote the gradient of the
in and t,' are processed by the expert c5. Let gn = Vo (9c., f') token in with respect the expert 0e.. A small change in 0e, in the direction of -gn is given by Oc < Ozg ōgn, with a step size 8. The effect of this change on the loss of another token ex' is measured by ALw' = L,
ögn. 8¹) (9 6') = -ōg, g'. + o()), where the second equality is obtained by first-order Taylor approximation. Therefore. the model updating for t, is considered to negatively affect token n' when g. gn' < O, since it increases the loss of token n', and vice versa. Thus, similar to (Yu et al., 2020), we define g., and gn as conflicting gradients when their cosine similarity cos Onn' < T, where T is a threshold and Onn' is the angle between ga and ga' Gradient conflicts cause the optimizer to converge to a sub-optimal solution.
We then define the conflicting token. Our goal is to adjust token routing to reduce gradient conflicts, but suddenly changing the routing of most tokens during training could lead to training instability. To stabilize training, we consider the expert as a whole to identify outliers as confficting tokens for the expert. Let the tokens processed by the expert e; be denoted as "t ,tN,, }, the average gradient on the expert 6 is represented as:
E gm Ne.
gmean
The average gradient indicates the holistic expert parameter updating direction at each iteration. When the gradient token and the average gradient are confficting gradients, this token detri- mental to the learning of the expert e,, so this token should be considered for assignment to another expert. Using the average gradient to identify conflicting tokes can keep most tokens in their current experts. A formal definition of confficting token is provided follows:
The token In is said to be a conflicting token if gn and gmenn are conflicting gradients, where gmean is the average gradient of all tokens in the expert tw.
Definition 1 (Conflicting Token)
Lastly, we detail our method for identifying conflicting tokens, as illustrated in Figure 2 (a). Initially, we unfreeze only the expert layer in the MoE and compute the main loss. We then perform back- propagation to calculate the gradient produced by each token on the expert parameters (i.e., token- level gradient). Subsequently, we calculate the average gradient, as well as the cosine similarity between the gradient of each token and the average gradiens. Lastly, when the similarity is less than T, mark the token parameters. Moreover, as the parameter size of each expert is very large, which leads to a very large parameter gradient size, we use engineering trick, i.e., using the gradients on part parameters as
conflicting token. During this process,
do not update any model
5

---

Published as a conference paper at ICLR 2025
an indicator, for reducing the GPU memory overhead to store gradients. Please refer to Sec. A of the supplementary material for more engineering implementation details. Identifying conflicting token allows us design the regularization loss to reduce gradient conflicts in the next section.
3.3 CONFLICT ELIMINATION LOSS
The learning of a conflicting token increases the average loss of tokens within its current expert. Thus, after a conflicting token is identified, should be reassigned to another expert for processing. To achieve this goal, we propose a simple yet effective regularization loss by constraining the routing scores predicted by the router, as shown in Figure 2 (b).
Specifically, we first identify the conflicting tokens for each expert within every MoE layer using token-level gradients. Then, the router predicts the routing logits (t,) of each conflicting token t. We record the current expert ID momo of each conflicting token f,. For LVLM, we use the recorded expert ID inmme,, to calculate the loss:
(tn) Zmoc (tn),
(mmo (fr)s
(mme(tn) (tn)i = L (tn)s
(8)
Naii
Lca.
logeP (tn)s) (mwe(tn)s.
Naul
i=1
where Nan is the count of all conflicting tokens, E is the number of experts, and Pimce (fw) represents the inverted routing score for the conflicting token ln- The Amoe define one-hot vectors, with quoo (fx)idmec. = 1. This loss is designed to encourage the reassignment of confficting tokens from their current experts other experts.
3.4 TOTAL Loss
To encourage experts to handle tokens in balanced manner, the differentiable load balancing loss, as introduced in (Fedus al., 2022), typically defined for each MoE layer as follows:
Laux F. F..P.
where Fú indicates the fraction of tokens processed by each expert e,, and P; indicates the average routing score of tokens assigned each expert ex. We also Laux (9) denote the average load balance loss of all MoE layers for convenience.
In conclusion, the total loss to update the model parameters is defined as:
Ciotal Lmoe + 8- LCEL = (Emain + a Laux) + 3 LCEL.
(10)
where Lmoe indicates the loss used in existing MoE-based LVLMs. a and 8 are We use Luain denote Lumin (0) for convenience, and the same applies to other losses.
EXPERIMENTS
4.1 EXPERIMENTAL. SETUP
Benchmark: Some academic-task- and instruction-following benchmarks are collected for evaluating the L.LMM. For academic-task-orientee benchmarks, VQA-v2 (Goyal et al., 2017b) and GQA (Hudson & Manning, 2019) assess the visual perception capabilities of models through open-ended short answers. VizWiz (Gurari et al., 2018) evaluates the zero-shot generalization of models visual questions asked by visually impaired people. ScienceQA (Lu et al., 2022), multiple-choice benchmark, evaluates the zero-shot generalization of models on scientific ques- tion answering. TextVQA (Singh al., 2019a) focuses text-rich visual question answering tasks. ChanQA (Masry et al., 2022) focuses on visual and logical reasoning tasks over charts. DoA (Mathew et al., 2021) focuses reading comprehension tasks document images.

---

Published as a conference paper at ICLR 2025
Table I: Comparison between diffferent LVLMS on image understanding benchmarks. "Act.", "V", "Q", "M". and "S" represent activated parameters, Vicuna (Chiang et al., 2023), Qwen (Bai et al., 2023a), Phi-2 (Microsoft, 2023), MobileLLaMA (Chu et al., 2023), and Sta- bleLM (Team), respectively. Main evaluation Benchmarks include VQA (Goyal et al., 2017a); GQA (Hudson Manning, 2019); VisWiz (Gurari et al., 2018); SQA': ScienceQA-IMG (Lu al., 2022); VQAT: TexIVQA (Singh et al., 2019b); POPE (Li et al., 2023a); MME (Fu et al., 2023);
MMB: MMBench (Liu et al., 2023d); MM-Ve Vet (Yu et al., 2023a). indicates that there is some Sapp--00000--aaaa0 384). All "Sparse Model" methods use the configure 4Top2. We calculate the average performance across all datasets except for MME, naming it "Avg". In the table below, report the best perfor- mance that we achieve when activating only 3.6B parameters during inference.
overlap in the training data. denotes the use of a stronger visual encoder
Method
LLM
Act.
/QA
GQA VieWiz SQA! VQA POPE
MM-Ver ARD ChanQA DoeVQA Avg
Dense Model
LIaVA-155
V.13R
13B
80.0"
63.3"
53.6
71.6
61.3
85.9
15313
67.7
35.4
49.6
18.1
24.0
55.5
Qwen- VL
Q-7B
6.78
78.8"
59.3"
35.2
67.1
63.8
38.2
115VA-15
V.7B
6.78
50.0
66.8
58.2
1510.7
30.5
TinyGPT-V
P-2.78
2.7B
33.6"
33.4
MobileVLM
M-2.7B
2.7B
59.0"
61.0
47.5
84.9
1288.9
59.6
LLaVA-Phi
P-2.7B
2.7B
71,4*
35,9
68.4
48.6
85.0
1335,1
59.8
28.9
Sporse Model
LLLVA
S-L6B
2.0
76.7' 7.
60.3"
36.2
62.6
50.1
85.7
1318.2
60.2
26.9
48.8
15.3
18.4
49.2
McE.LLaVA
P.2.7B
3.6
43.9
68.5
51.4
1423.0
34.3
19.9
21.5
DYNMOE-LLaVA
P-2.78
3.48
77.9"
61.6"
45.1
68.0
51.8
86.0
1429.6
66.6
33.6
MoE. LLaVA
P.2.7B
3.6B
79.9*
62.6"
43.7
70.3
57.0
85.7
1431.3
68.0
35.9
59.5
15.4
25.6
54.9
Our Method
P-2.7B 3.6B|
48.6
70.9
58.8
1481.7
40.7
42.1
For instruction followiin benchmarks, POPE (Li et al., 2023b) evaluates the degree hallucination in model responses on three sampled subsets of COCO (Lin et al., 2014): Random, Common, and Adversarial. MME (Fu et al., 2023) assesses the visual perception of models with yes/no questions. MMBBnch (Liu al., 2023d) evaluates the robustness of model answers with all-round shuffling on multiple choice answers. MM-Vet (Yu et al., 2023b) evaluates the model capabilities in engaging in visual conversations on a diverse range of tasks and evaluates the correctness and helpfuiness of the responses using the GPT-4 evaluation framework. A12D (Kembhavi et al., 2016), multiple-choice benchmark, evaluates the model capabilities for science diagram comprehension.
(Lin et al., 2024). MoE-LLaVA incorporates MoE into LVLMs and proposes a stage training scheme. It trains only the MoE in the third stage, i.e., the instruction tuning stage. MoF-LLLaVA has four experts and selects the Top-2 experts to handle tokens, and we refer to this configuration as 4Top2. Building on MoF- VA, we add novel regularization loss LCEL during the instruction tuning stage to enhance the MoE. For the language model backbone, we follow MoE-LLaVA to use StableLM- 1.6B and Phi2-2.7B. The visual encoder is usually set as clip-vit-large prnge-pptch114--1 13-336. a=0001, following MoE-LLaVA. We also compare with DYNMÓE-LLaVA. (Guo et al., 2024), which improves the MoE-LLaVA by dynamically setting the expert count. For more implementation details, please refer to Sec. of supplementary material.
Baseline: Our main baseline is
4.2 IMAGE UNDERSTANDING EVALUATION
We evaluate the performance of our method on five image question- answering benchmarks, as shown in Table 1, and report the number of activated parameters as a measure of efficiency. Our method demonstrates superior image understanding capabilities, achiev- ing the 80.0%, 63.0%, 48.6%, 70.9%, and 58.8% performance on VQA GQA, VisWiz, SQA!, and VQAT, respectively. Compared to LLaVA-1. with 7B activated parameters, our method brings 1.5%, 1.0%, 4.1% %, and 0.6% performance increase on VQA*2, GQA, SQA', and VQAT, respec- tively, when only activating 3.6B parameters.
Image Question Answering:
Benchmark Toolkit: To comprehensively evaluate the multi-modal understanding capabilities of method, assess its performance across four benchmark toolkits in Table 1. These toolkits typically serve as tools to verify the model ability to engage in natural language questioning. As shown in Table our method achieves 86.5%, 1481.7, 71. 11.0% and 40.7% performance on POPE,
7

---

Published as a conference paper at ICLR 2025
Table STGC as a plug-in. We set different baselines and add the proposed STGC. 4Topl means four experts are set, and the Top-1 expert is selected to handle tokens. indicates that a stronger visual encoder (silipp-s400mm 1at000mmpach11-338880 is used. We calculate the average performance across all datasets except for MME, naming it "Avg".
Method
LLM Act. VQA GQA WisWiz SQA!
VQAT
POPE
MME
MMB
MM-VE Vet
Avg
NOE-LLLVA-4Topp
S-1.6B 1.6B
74.5*
58.6*
25.7
55.8
45.0
85.2
1245.3
56.2
27.2
53.5
+STGC
S-1.6B 1.6
74.9*
59.4*
27.4
57.5
46.5
85.8
1276.8
56.8
28.5
54.6
MoE-LLaVA-4Top2
S-1.6B 2.0B
76.7*
60.3*
36.2
62.6
50.1
85.7
1318.2
60.2
26.9
57.3
+STGC
S-1.6B 2.0B
76.9*
60.9*
37.7
62.6
50.7
85.9
1355.1
60.7
28.2
58.0
MoE-LLaVA-4Top2
P-2.7B 3.6B
77.6"
61.4"
43,9
68.5
51.4
86.3
1423.0
65.2
34.3
61.1
+STGC
P-2.7B 3.6B
78.0"
62.1*
47.2
68.1
52.3
86.9
1429.2
66.7
33.3
61.8
MoF-LLaVA--Toop2
P-2.7B 3.6B
79.9"
62.6*
43.7
70.3
57.0
85.7
1431.3
68.0
35.9
62.9
+STGC P-2.7B 3.6B
80.3* 63.2"
45.1
70.3
57.4
86.1
1447.6
69.7
35.7
63.5
Table 3: Scalability of STGC. We consider employing the larger scale of data. With more data, the STGC can bring a more significant performance increase.
Method
LLM Data
VQA
GQA
WisWiz SQA'
VQA
POPE
MME
MMB MM-Vet Avg
MoE-LLaVA-4Top2
P-2.7B 665K
79.9"
62.6"
43.7
70.3
57.0
85.7
1431.3
68.0
35.9
62.9
+STGC
P.2.7B 665K
80.3"
63.2"
45.1
70.3
57.4
86.1
1447.6
69.7
35.7
63.5
MoF-LLaVA-4Top2""
P.2.7B 1021K
79.7"
63.0*
42.7
71.1
56.9
84.3
1439.9
70.4
42.2
163.8
+STGC
P-2.7B 1021K
80.0"
63.0"
48.6
70.9
58.8
86.5
1481.7
71.0
40.7
64.9
Table 4: Study about routing strategies. embedding- baaeed, feature-based. and task-based corre- spond to routing strategies similar to MoCLE (Gou et al., 2023), LoRA-MoE (Chen et al., 2023d), and MoCLE (Zhou al., 2024), respectively. The configure MoE-LLaVA-4To0) with StableLM- 1.6B is set as the baseline.
Method
VQA
GQA
VisWiz
SQA'
VQAT
POPE
MME
MMB
MM- Vet
Avg
MoE-LLaVA 76.7*
60.3"
36.2
62,6
50.1
85.7
1318.2
60.2
26.9
57.3
tembedding- hasse
75.8*
57.0"
34.0
63.7
50.3
86.1
1312.8
61.3
27.3
56.9
-feature-bsed
75.7"
58.1"
36.9
63.2
50.0
85.9
1338.8
61.5
26.6
57.2
+task-buse based
73.6"
58.2"
29.0
63.7
49.2
81.5
1306.3
59.5
25.2
48.9
+STGC
76.9*
60.9"
37.7
62.6
50.7
85.9
1355.1
60.7
28.2
58.0
MME, MMB, and MM-Vet, respectively. Compared LLaVA-1.5 with 7B activated parameters, our method brings 0.6%, 7.6%, and 10.2% performance increase on POPE, MMB, and MM-Vet, respectively, when only activating 3.6B parameters.
4.3 ABLATION STUDY
STGC as a plug-in: We design the STGC as plug-in for existing MoE methods. To verify the robustness and effectiveness of the proposed STGC as plug-in, we select different baselines, in- cluding different MoE configures (4Top2 & 4Top1). LLM (S-1.6B & P-2.7b), and visual encoders (clip& siglip). Then, we add the STGC onto these baselines. As shown in Table 2, adding STGC always brings stable and convincing performance increase under diverse baselines.
Scalability STGC: One necessary hypothesis this work that there severe interference between diverse training data, so STGC should bring a more significant performance increase when using a more complex dataset to train. To verify this, we expand the size of the public dataset we are using for further experiments. Specifically. NeXT project (Lin & Long, 2024). As shown Table 3, under more data, STGC brings a more significant average performance increase (1.1% % vs. 0.6%)
use the training data provided in the Open-LLaVA-
8

---

Published as a conference paper at ICLR 2025
Table 5:
Study about hyper-parameter sensitivity. Settings for results in Table 2 are highlighted The configure MoE-LLaVA-4Top2 with StableLM--1.6B set as the baseline.
blue
(a) The threshold for identifying confiticting tokens.
(b) The weighting of our proposed loss.
GQA
WisWiz
VQA*
MMB
MM-Vet
GQA
VisWiz
VQA*
MMB
MM-Vet
0.1
35.1
50.9
61.3
25.6
60.5
35.6
50.5
61.4
26.9
0.0
60.9
37.7
50.7
60.7
28.2
1.0
60.9
37.7
50.7
60.7
28.2
-0.1
60.6
34.9
50.5
61.4
25.9
2.0
60.6
35.9
50.9
60.6
27.2
(a)
1.54
Baselane
Buscline SIGC
1.5
Conflict Elimination
Gradient Consistency
0.2
1.46
Loss
Baseline
1.42
1.38
40
80
120
160
200
120
200
Training Step
Training Step
(b)
Layer
Layer
2.5
033
Heseline + STGC
Biseline + STGC
Baseline STEC
Gradient Consistency
0.1
Baseline
Baseline
5.3
Baseline
80
82
120
160
200
40
80
120
140
200
82
120
1.00
200
Training Step
Training Step
Training Step
Figure 3: Statistical verification. We conduct a deep analysis of the role of STGC. "Baseline" indicates MoE-LLaVA. "Baseline + STGC" indicates our method. (a) We compute novel metric, gradient consistency (the mean cosine similarity between gradients of all tokens within an expert), for verifying that the decrease of the proposed loss leads to the more consistent token gradients within an expert. (b) We further analyze the gradient consistency different layers.
Study about Different Routing Strategies: We compare with different routing strategies in Table 4. Some works (Chen et al., 2023d; Gou et al., 2023; Zhou et al., 2024) propose using instruction embeddings (features) sample task labels design routing strategies. Since they designed for LoRA-MoEE, we directly replace the FFN with the MoE structure, making it difficult to compare the performance of STGC with that in their papers. To this end, we re-implemented these methods in MoE-LLaVA. Specifically, (i) embedding- bbaeed. Similar to MoCLE (Gou et al., 2023), we use the K-means algorithm to cluster instruction embeddings into 128 groups before training and conduct the routing based on the cluster results. (ii) feature-based. Similar to LoRA-MoE (Chen et al., 2023d), we use instance-level instruction token average representation to predict routing scores. (ii) task-based. Similar to (Zhou et al., 2024), we categorize samples into four experts based on their tasks. For more details, please refer Sec. B of supplementary material. As shown in Table 4, the performance does not increase most datasets when using sample-level embedding or task information. Using token-level gradients achieves significantly higher average performance improvement.
Sensitivity Study of Different Hyper-parameters: We conduct the sensitivity study of STGC in Table for two main pappr--aramtters, the threshold of identifying conflicting tokens and the loss weighting 3. First, when the gradient ga of the token f,, and and the average gradient gmean satisfy the condition Onmean T, we flag the token conflicting loken. We discuss different

---

Published as a conference paper at ICLR 2025
thresholds of identifying conflicting tokens: E 00.1,0.0,-011;, Second, we discuss different loss weightings 3 E {0.5, 1.0, 2.0). As shown in Table 5, we find: (i) When T = 0, the performance on most datasets is the best. The results are consistent with the common belief, i.e., gradients considered conflicting when their cosine similarity is less than O. (ii) The proposed loss is relatively robust different loss weightings B, with the highest performance on most datasets when ,3=100.
Statistical Verification: One main claim of this work is that STGC can effectively reduce the gra- dient conflicts between tokens within an expert, i.e., making the gradient directions of tokens within an expert more consistent, thus decreasing data interference and improving model performance. To verify this, we design a statistical verification experiment. In this experiment, the training consists of two steps: The first step is to gather gradients of all tokens within an expert e; and compute their cosine similarity to form a similarity matrix. Then, we define the mean of the similarity matrix as sim;, We define the average of sim; on all experts as gradient consistency, serving novel metric to evaluate whether the gradient directions are consistent within the expert. The second step is to update models. In this experiment, undisturbedly observing its impact on the gradient consistency. In experiments, we fed one sample into the LVLM per device for each forward pass, and the gradient accumulation step is 16.
only use the proposed loss to update parameters, for
We first present the curve graph of the proposed loss and the gradient consistency obtained from the TensorBoard dashboard in Figure 3 (a). Then, we present the gradient consistency on different layers in Figure 3 (b); The total layer number is 32, and we analyze the layers {8, 16, 24). We find: (i) By decreasing the proposed conffict elimination loss, the gradient consistency significantly improves, indicating that the token gradient directions within an expert become more consistent. This verifies the role of STGC in reducing gradient conflicts. (ii) For different layers. STGC always increases the gradient consistency. The deeper layers seem to have higher gradient consistency after adding STGC, while the shallower layers have lower gradient consistency.
CONCLUSION AND LIMITATIONS
Our study reveals that there is the interference between tokens within an expert for the MoE, leading to sub-optimal learning for the expert. To reduce the interference between tokens, we propose em- ploying token-level gradients to identify conflicting tokens, and then adding novel conffict clim- ination loss to optimize token routing based on confficting tokens. Our method STGC acts as a plug-in, which can be easily integrated into existing MoE-based LVLMs. Extensive experiments demonstrate the effectiveness of STGC across diverse datasets. Especially when the data diversity larger, our method brings more significant performance increase.
Limitation. Since in each iteration, a token only passes one expert and not the others in each MoE layer, its gradient only reflects whether it conflicts with the holistic optimization direction of its current expert, but it is difficult to define its relationship with other experts. Thus. although it has been confirmed that the proposed solution can solve token gradient conflicts by optimizing token routing, the solution still has room for improvement, as it is challenging to determine the optimal expert of token. In the future, exploring whether is possible to determine the optimal expert of a token from the optimization perspective may be an intriguing direction.
Acknowledgements. This work supported part by National Science Foundation for Distin- guished Young Scholars under Grant 62225605, Zhejiang Provincial Natural Science Foundation of China under Grant I.D24F020016, "Pioneer" and "Leaddig Goose" R&D Program of Zhejiang (No. 2024C01020), Project 12326608 supported by NSFC, the Ningbo Science and Technology Innovation Project (No.2024Z294), and is supported by Kuaishou Technology.
REFERENCES
Jean-Baptiste Alayrac. Jeff Donahue, Pauline Luc, Antoine Miech, Iain Bart, Yana Hasson, Karel Lene, Arthur Mensch, Katherine Millican, Malcolm Reynolds, et al. Flamingo: a visual language model for few-shot learning. Advances in Neural Information Processing Systems, 35:23716- 23736, 2022.
Jinze Bai, Shuai Bai, Yunfei Chu, Zeyu Cui, Kai Dang, Xiaodong Deng, Yang Fan, Wenbin Ge, Yu Han, Fei Huang, et al. Qwen technical report. arXiv preprint arXiv:2309. 16609, 2023a.
10

---

Published as a conference paper at ICLR 2025
Jinze Bai, Shuai Bai, Shusheng Yang, Shijie Wang, Sinan Tan, Peng Wang, Junyang Lin, Chang Zhou, and Jingren Zhou. Qwen-vl: A frontier large lainule model with versatile abilities. arXiv preprint arXiv:2308. 2023b.
Hangbo Bao, Wenhui Wang, Li Dong, Qiang Liu, Owais Khan Mohammed, Kriti Aggarwal, Sub- hojit Som, Songhao Piao, and Furu Wei. Vlmo: Unified vision-language pre-training with mixture-of-modait -ooiu----eo-odllll-eeppepeere Advances in Neural Information Processing Systems, 35:32897- 32912, 2022.
Atilim Güneş Baydin, Barak A Pearlmutter, Don Syme, Frank Wood, and Philip Tor. Gradients without backpropagation. arXiv preprint arXiv:2202.08587, 2022.
Luisa Bentivogli, Peter Clark, Ido Dagan, and Danilo Giampiccolo. The fifth pascal recognizing textual entailment challenge. TAC, 7(8):1, 2009.
Stephen Boyd and Lieven Vandenberghe. Convex optimization. Cambridge university press, 2004.
Junbum Cha, Wooyoung Kang. Jonghwan Mun, and Byungseok Roh. Honeybee: Locality-enhanced projector for multimodal llm. arXiv preprint arXiv:2312.06742 2023.
Nithin Chalapathi, Yiheng Du, and Aditi Krishnapriyan. Scaling tehysies-infomme hard constraints with mf arXiv preprint arXiv:2402. 124 2024.
Jun Chen, Deyao Zhu, Xiaoqian Shen, Xiang Li, Zechun Liu, Pengchuan Zhang, Raghuraman Krishnamoorthi, Vikas Chandra, Yunyang Xiong, and Mohamed Elhoseiny. Minigpt-v2: large language model as a unified interface for vision-language multi-task learning. arXiv preprint arXiv:2310.09478, 2023a.
Lin Chen, Jisong Li, Xiaoyi Dong, Pan Zhang, Conghui He, Jiaqi Wang, Feng Zhao, and Dahua Lin. Sharegpt4v: Improving large multi-modal models with better captions. arXiv preprint arXiv:2311.12793 2023b.
Shaoxiang Chen, Zequn Jie, and Lin Ma. Llava-mole: Sparse mixture of lora experts for mitigating data conflicts in instruction finetuning milms. arXiv preprint arXiv:2401.16116, 2024.
Wuyang Chen, Yanqi Zhou, Nan Du, Yanping Huang, James Laudon, Zhifeng Chen, and Claire Cui. Lifelong language pretraining with distribationnsspeciaiiee experts. In International Conference on Machine Learning, pp. 5383-5395. PMLR, 2023c.
Zeren Chen, Ziqin Wang, Zhen Wang, Huayang Liu, Zhenfei Yin, Si Liu, Lu Sheng, Wanli Ouyang, Yu Qiao, and Jing Shao. Octavius: Mitigating task interference in milms via moe. arXiv preprint arXiv:2311.02684, 2023d.
Zhe Chen, Jiannan Wu, Wenhai Wang. Weijie Su, Guo Chen, Sen Xing, Zhong Muyan, Qinglong Zhang, Xizhou Zhu, Lewei Lu, et al. Internvl: Scaling up vision foundation models and aligning for generic visual- vinual-linguistic tasks. arXiv preprint arXiv:2312.14238. 2023c.
Wei-Lin Chiang, Zhuohan Li, Zi Lin, Ying Sheng, Zhanghao Wu, Hao Zhang, Lianmin Zheng, Siyuan Zhuang, Yonghao Zhuang, Joseph E Gonzalez, et al. Vieuna: An open-source chatbot impressing gpt-4 with 90%* chatgpt quality. See https://vicuna Imsys. org (accessed 14 April 2023), 2023.
Xiangxiang Chu, Limeng Qiao, Xinyang Lin, Shuang Xu, Yang Yang, Yiming Hu, Fei Wei, Xinyu Zhang, Bo Zhang, Xiaolin Wei, et al. Mobilevim: A fast, reproducible and strong vision language assistant for mobile devices. arXiv preprint arXiv:2312.168886 2023.
Damai Dai, Chengqi Deng, Chenggang Zhao, RX Xu, Huazuo Gao, Deli Chen, Jiashi Li, Wangding Zeng, Xingkai Yu, Y Wu, et al. Deepseekmoe: Towards ultimate expert specialization in mixture- of-experts language models. arXiv preprint arXiv:2401.0600 2024.
Wenliang Dai. Junnan Li, Dongxu Li, Anthony Meng Huat Tiong, Junqi Zhao, Weisheng Wang, Boyang Li. Pascale Fung, and Steven Hoi. Instructblip: Towards general-purpose purpose vision-language models with instruction tuning, 2023.
11

---

Published as a conference paper at ICLR 2025
Jacob Devlin. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04900 2018.
Bill Dolan and Chris Brockett. Automatically constructing corpus of sentential paraphrases. In Third international workshop on paraphrasing (IWP200S). 2005.
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. The Journal of Machine Learning Research, 23(1): 5232-5270, 2022.
Chaoyou Fu, Peixian Chen, Yunhang Shen, Yulei Qin, Mengdan Zhang, Xu Lin, Jinrui Yang, Xiawu Zheng, Ke Li, Xing Sun, Yunsheng Wu, and Rongrong Ji. Mme: A comprehensive evaluation benchmark for multimodal large language models. arXiv preprint arXiv:2306. 113223066.3344 13394, 2023.
Yunhao Gou, Zhili Liu, Kai Chen, Lanqing Hong, Hang Xu, Aoxue Li, Dit-Yan Yeung, James T Kwok, and Yu Zhang. Mixuure of cluster-conditional lora experts for vision-language instruction tuning, arXiv preprint arXiv:2312. 2023.
Yash Goyal, Tejas Khot, Douglas Summers-Stre -Stner Dhruv Batra, and Devi Parikh. Making the v in vqa matter: Elevating the role of image understanding in visual question answering. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 6904-6913, 2017a.
Yash Goyal, Tejas Khot, Douglas Summers-Stay, Dhruv Batra, and Devi Parikh. Making the v in vqa matter: Elevating the role of image understanding in visual question answering. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 6904-6913, 2017b.
Yongxin Guo, Zhenglin Cheng, Xiaoying Tang, and Tao Lin. Dynamic mixture of experts: An auto-tuning approach for efficient transformer models. arXiv preprint arXiv:2405. 2rXXi::405..4297, 2024.
Danna Gurari, Qing Li. Abigale Stangl, Anhong Guo, Chi Lin, Kristen Grauman, Jiebo Luo, and Jeffrey P Bigham. Vizwiz grand challenge: Answering visual questions from blind people. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 2018.
Drew Hudson and Christopher Manning. Gqa: A new dataset for real-world visual reasoning and compositional question answering. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 6700-6709, 2019.
Gagan Jain, Nidhi Hegde, Aditya Kusupati, Arsha Nagrani, Shyamal Buch, Prateek Jain, Anurag Arnab, and Sujoy Paul. Mixture of nested experts: Adaptive processing of visual tokens. arXiv preprint a40: 19985, 2024.
Albert Q Jiang, Alexandre Sablayrolles. Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bam- ford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, et al. Mixtral of experts. arXiv preprint arXiv:2401.04088 2024.
Aniruddha Kembhavi, Mike Salvato, Eric Kolve, Minjoon Seo, Hannaneh Hajishirzi, and Ali Farhadi. A diagram is worth dozen images. In Computer Visi 2016: 14th Euro- pean Conference, Amsterdam, The Netherlands, October 11-14, 2016, Proceedings, Part IV 14, pp. 235-251. Springer, 2016.
Aran Komatsuzaki, Joan Puigcerver, James Lee-Thorp, Carlos Riquelme Ruiz, Basil Mustafa, Joshua Ainslie, Yi Tay, Mostafa Dehghani, and Neil Houlsby. Sparse upcycling: Training mixture-of-experts from dense checkpoints. arXiv preprint 0rX1v:2212.05055, 2022.
Sneha Kudugunta, Yanping Huang, Ankur Bapna, Maxim Krikun, Dmitry Lepikhin, Minh-Thang Luong, and Orhan Firat. Beyond distillation: Task-level mixture-of-experts for efficient inference. arXiv preprint arXiv:21
2021.
Byung-Knn Kwan Lee, Beomchan Park, Chae Won Kim, and Yong Man Ro. Moai: Mixture of all intelligence for large language and vision models. arXiv preprint arXiv:2403.0753 0750088 2024.
12

---

Published as a conference paper at ICLR 2025
Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu. Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. Gshard: Scaling giant models with conditional computation and automatic sharding. arXiv preprint arXiv:2006. 16668, 2020.
Jiachen Li, Xinyao Wang, Sijie Zhu, Chia-Wen Kuo, Lu Xu, Fan Chen, Jitesh Jain, Humphrey Shi, and Longyin Wen. Cumo: Scaling multimodal IIm with co-upeycled mixture-of-experts arXiv preprint arXiv:2405. 05949, 2024a.
Junnan Li, Dongxu Li, Caiming Xiong, and Steven Hoi. Blip: Bootstrapping language-iman image pre- training for unified vision-language understanding and generation. In International Conference on Machine Learning, pp.
PMLR, 2022.
Yifan Li, Yifan Du, Kun Zhou, Jinpeng Wang, Wayne Xin Zhao, and Ji-Rong Wen. Evaluating object hallucination in large vision-language models. arXiv preprint arXiv:2305. 10355, 2023a.
Yifan Li, Yifan Du, Kun Zhou, Jinpeng Wang, Wayne Xin Zhao, and Ji-Rong Wen. Evaluating object hallucination in large vision-language models. arXiv preprint arXiv:2305 2023b.
Yunshui Li, Binyuan Hui, ZhiChao Yin, Min Yang, Fei Huang, and Yongbin Li. Pace: Unified modal dialogue pre-training with progressive and compositional experts. arXiv preprint arXiv:2305. 14839, 2023c.
Yunxin Li, Shenyuan Jiang, Baotian Hu, Longyue Wang, Wanqi Zhong, Wenhan Luo, Lin Ma, and Min Zhang. Uni-moe: Scaling unified multimodal Ilms with mixture of experts. arXiv preprint arXiv:2405.11223. 2024b.
Victor Weixin Liang, Yuhui Zhang, Yongchan Kwon, Serena Yeung, and James Y Zou. Mind the gap: Understanding the modality gap in multi-modal contrastive representation learning. Ad- vances in Neural Information Processing Systems, 35: 75511116627 2022.
Bin Lin, Bin Zhu, Yang Ye, Munan Ning, Peng Jin, and Li Yuan. Video-llava: Learning united visual representation by alignment before projection. arXiv preprint arXiv:2311.10122, 2023.
Bin Lin, Zhenyu Tang, Yang Ye, Jiaxi Cui, Bin Zhu, Peng Jin, Junwu Zhang, Munan Ning, and Li Yuan. Moe-Ilava: Mixture of experts for large vision-language models. arXiv preprint arXiv:2401.15947 2024.
Chen Lin and Xing Long. Open-llava-next: An open-source implementation of llava-next ries for facilitating the large multi-modal model community. https://githuu. com/ xiacachen98/ (atpeenn///pen--aiee LLaVA NeXT,, 2024.
Tsung-Y Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C Lawrence Zitnick. Microsoft COCO: Common objects in context. In ECCV, 2014.
Fuxiao Liu, Kevin Lin, Linjie Li, Jianfeng Wang, Yaser Yacoob, and Lijuan Wang. Aligning large multi-modal model with robust instruction tuning. arXiv preprint arXiv:2306. /4565, 2023a.
Haotian Liu, Chunyuan Li, Yuheng Li, and Yong Jae Lee. Improved baselines with visual instruction tuning. arXiv preprint arXiv:2310.03744 2023b.
Haotian Liu, Chunyuan Li. Qingyang Wu, and Yong Jae Lee. Visual instruction tuning. arXiv preprint arXiv:2304 08448, 2023c.
Yuan Liu, Haodong Duan, Yuanhan Zhang, Bo Li, Songyang Zhang, Wangbo Zhao, Yike Yuan, Jiaqi Wang, Conghui He, Ziwei Liu, al. Mmbench: your multi-modal model an all-around player? arXiv preprint arXiv:2307.06288 2023d.
Zefang Liu and Jiahua Luo. Adamole: Fine-tuning large language models with adaptive mixture of low-rank adaptation experts. arXiv preprint 20X0::2405.003611 2024.
Zijun Long, George Killick, Richard McCreadie, and Gerardo Aragon Camarasa. Multiway- adapater: Adapting large-scale multi-modal models for scalable image-text retrieval. arXiv preprint arXiv:2309.015/6, 2023.
13

---

Published as a conference paper at ICLR 2025
Pan Lu, Swaroop Mishra, Tanglin Xia, Liang Qiu, Kai-Wei Chang. Song-Chun Zhu, Oyvind Tafjord, Peter Clark, and Ashwin Kalyan. Learn to explain: Multimodal reasoning via thought chains for science question answering. Advances in Neural Information Processing Systems. 35:2507-2521, 2022.
Guangyuan Ma, Xing Wu, Peng Wang, and Songlin Hu. Cot-mote: Exploring contextual masked auto-encoder pre-training with mixture-ofttxtuallexper for passage retrieval. arXiv preprint arXiv:2304. 10/95, 2023.
Ahmed Masry, Do Xuan Long, Jia Qing Tan, Shafiq Joty, and Enamul Hoque. Chartqa: bench- mark for question answering about charts with visual and logical reasoning. arXiv preprint arXiv:2203. 10244, 2022.
Minesh Mathew, Dimosthenis Karatzas, and CV Jawahar. Docvqa: A dataset for vqa on document images. In Proceedings of the IEEE/CVF winter conference on applications of computer vision, pp. 2200-2209, 2021.
Microsoft. cls.
Phi-2:
The surprising
power of small language mod-
https://www.iirrre
2023.
Fangzhou Mu, Yingyu Liang, and Yin Li. Gradients as features for deep representation learning. arXiv preprint arXiv:2004.05529, 2020.
Max Paulus, Dami Choi, Daniel Tarlow, Andreas Krause, and Chris Maddison. Gradient estimation with stochastic softmax tricks. Advances in neural information processing systems, 33:5691- 5704, 2020.
Joan Puigcerver, Carlos Riquelme, Basil Mustafa, and Neil Houlsby. From sparse to soft mixtures of experts. arXiv preprint arXiv:2308. 00951, 2023.
Zihan Qiu, Zeyu Huang, and Jie Fu. Emergent mixxure-of-experree Can dense pre-trained transform- ers benefit from emergent modular structures? arXiv preprint arXiv:2310. 109008 2023.
N Reimers, Sentence-bert Sentence embeddings using siamese bert-networks. 1rXiv:1908.10080 2019.
arXiv preprint
Carlos Riquelme, Joan Puigcerver, Basil Mustafa, Maxim Neumann, Rodolphe Jenatton, André Susano Pinto, Daniel Keysers, and Neil Houlsby. Scaling vision with sparse mixture of experts. Advances in Neural Information Processing Systems, 34:8583-8595, 2021.
Burak Satar, Hongyuan Zhu, Hanwang Zhang, and Joo Hwee Lim. Rome: Role-aware mixture-of- expert transformer for text-to-video retrieval. arXiv preprint (tti:2206.128445 2022.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. arXiv preprint arXiv: 707700001588. 2017.
Leyang Shen, Gongwei Chen, Rui Shao, Weili Guan, and Liqiang Nie. Mome: Mixture of multi- modal experts for generalist multimodal large language models. arXiv preprint arXiv:2407. 12709, 2024.
Sheng Shen, Le Hou, Yanqi Zhou, Nan Du, Shayne Longpre, Jason Wei, Hyung Won Chung, Bar- ret Zoph, William Fedus, Xinyun Chen, et al. Mixture-of ex-eeerss meets instruction tuning: A winning combination for large language models. arXiv preprint 27Xi 2023a.
Sheng Shen, Zhewei Yao, Chunyuan Li, Trevor Darrell, Kurt Keutzer, and Yuxiong He. Scaling vision-language models with sparse mixture of experts. arXiv preprint arXiv:2303.07226, 2023b.
Jiaxin Shi, Yuhao Zhou, Jessica Hwang, Michalis Titsias, and Lester Mackey. Gradient estimation with discrete stein operators. Advances in neural information processing systems, 35:25829- 25841, 2022.
14

---

Published as a conference paper at ICLR 2025
Amanpreet Singh, Vivek Natarajan, Meet Shah, Yu Jiang, Xinlei Chen, Dhruv Batra, Devi Parikh, and Marcus Rohrbach. Towards vqa models that can read. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 8317-8326, 2019a.
Amanpreet Singh, Vivek Natarajan, Meet Shah, Yu Jiang, Xinlei Chen, Dhruv Batra, Devi Parikh, and Mareus Rohrbach. Towards vqa models that can read. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 8317-8326, 2019b.
Stability Al Language Team. staailityai/staalllee-t stableim-2-1.00 6b).
Stable Im 1.6b. URL (https: htttpss//huggingaac..o
6b] (https: htte:///hggggngggeeeeeee wtteaaaccccecoeeeeeeeeeeeeeeeeee co/stabilityai.
Alex Wang. Glue: multi-task benchmark and analysis platform for natural language understand- ing. arXiv preprint arXiv:1804.07461, 2018.
Wenhui Wang, Hangbo Bao, Li Dong, Johan Bjorck, Zhiliang Peng. Qiang Liu, Kriti Aggarwal, Owais Khan Mohammed, Saksham Singhal, Subhojit Som, et al. Image as a foreign language: Beit pretraining for all vision and vision-language tasks. arXiv preprint arXiv:2208. 10442, 2022.
Yisen Wang, Xingjun Ma, Zaiyi Chen, Yuan Luo, Jinfeng Yi, and James Bailey. Symmetric cross entropy for robust learning with noisy labels. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 322-330, 2019.
A Warstadt. Neural network acceptability judgments. arXiv preprint 1805. 12471, 2019.
Xiao-Ming Wu, Dian Zheng, Zuhao Liu, and Wei-Shi Zheng. Estimator meets equilibrium perspec- tive: A rectified straight through estimator for binary neural networks training. In Proceedings of the JEEE/CVF International Conference on Computer Vision, pp. 17055-17064, 2023.
Xun Wu, Shaohan Huang, and Furu Wei. Mixture of lora experts. arXiv preprint arXiv:240..13688, 2024.
Liang Xu, Hai Hu, Xuanwei Zhang, Lu Li, Chenjie Cao, Yudong Li, Yechen Xu, Kai Sun, Dian Yu, Cong Yu, et al. Clue: A chinese language understanding evaluation benchmark. arXiv preprint arXiv:2004.05980 05986,, 2020.
Longrong Yang, Fanman Meng, Hongliang Li, Qingbo Wu, and Qishang Cheng. Learning with noisy class labels for instance segmentation. In Computer Vision-ECCV 2020: 16th European Conference, Glasgow, UK, August 23-28, 2020, Proceedings, Part XIV 16, pp. 38 13.53. Springer, 2020.
Longrong Yang, Hongliang Li, Qingbo Wu, Fanman Meng, Heqian Qiu, and Linfeng Xu. Bias- correction feature learner for supe-supe instance segmentation. IEEE Transactions Mul- timedia, 25:5852-5863, 2022.
Longrong Yang, Xiampan Zhou, Xuewei Li, Liang Qiao, Zheyang Li, Ziwei Yang, Gaoang Wang, and Xi Li. Bridging cross-task protocol inconsistency for distillation in dense object detection. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 17175-17184, 2023.
Longrong Yang, Hanbin Zhao, Yunlong Yu, Xiaodong Zeng, and Xi Li. Res-prompt: Learning prompt to rearrange class space for prompt-based continual learning. In European Conference Computer Vision, pp. 1-20. Springer, 2024.
Longrong Yang. Hanbin Zhao, Hongliang Li, Liang Qiao, Ziwei Yang, and Xi Li. Gestg: Generating class-confusion-awwe aware samples with tree tree-structure graph for few-shot object detection. IEEE Transactions on Image Processing, 2025.
Qinghao Ye, Haiyang Xu, Guohai Xu, Jiabo Ye, Ming Yan, Yiyang Zhou, Junyang Wang, Anwen Hu, Pengcheng Shi, Yaya Shi, et al. mplug-owl: Modularization empowers large language models with multimodality. arXiv preprint 2304 14/78, 2023.
15

---

Published as a conference paper at ICLR 2025
Tianhe Yu, Saurabh Kumar, Abhishek Gupta, Sergey Levine, Karol Hausman, and Chelsea Finn. Gradient surgery for multi-task learning. Advances Neural Information Processing Systems, 33:5824-5836, 2020.
Weihao Yu, Zhengyuan Yang, Linjie Li, Jianfeng Wang, Kevin Lin, Zicheng Liu, Xinchao Wang. and Lijuan Wang. Mm-vet: Evaluating large multimodal models for integrated capabilities. arXiv preprint arXiv:2308. 02490, 2023a.
Weihao Yu, Zhengyuan Yang, Linjie Li, Jianfeng Wang, Kevin Lin, Zicheng Liu, Xinchao Wang, and Lijuan Wang. Mm-vet: Evaluating large multimodal models for integrated capabilities. arXiv preprint arXiv:2308 02490,, 2023b.
Ted Zadouri, Ahmet Üstün, Arash Ahmadian, Beyza Ermiş, Acyr Locatelli, and Sara Hooker. Push- ing mixture of experts to the limit: Extremely parameter efficient moe for instruction tuning. arXiv preprint arXiv:2309.088 054444, 2023.
Pan Zhang, Xiaoyi Dong Bin Wang, Yuhang Cao, Chao Xu, Linke Ouyang, Zhiyuan Zhao, Shuan- grui Ding, Songyang Zhang, Haodong Duan, Hang Yan, et al. Internlm-xcome xcooposr... A vision- language large model for advanced text-image comprehension and composition. arXiv preprint arXiv:2309. 15112, 2023a.
Yanzhe Zhang, Ruiyi Zhang, Jiuxiang Gu, Yufan Zhou, Nedim Lipka, Diyi Yang, and Tong Sun. Llavar: Enhanced visual instruction tuning for text-rich image understanding. arXiv preprint arXiv:2306. 17107. 2023b.
Zhengyan Zhang, Yankai Lin, Zhiyuan Liu, Peng Li, Maosong Sun, and Jie Zhou. Moefication: Transformer feed-forward layers are mixtures of experts. arXiv preprint arXiv:2110.01788 2021.
Zhilu Zhang and Mert Sabuncu. Generalized cross entropy loss for training deep neural networks with noisy labels. Advances in neural information processing systems, 31, 2018
Bo Zhao, Boya Wu, and Tiejun Huang. Svit: Scaling up visual instruction tuning. arXiv preprint wwXVv22307.04087. 2023.
Shanshan Zhong, Shanghua Gao, Zhongzhan Huang, Wushao Wen, Marinka Zitnik, and Pan Zhou. Moextend: Tuning new experts for modality and task extension. arXiv preprint arXiv:2408.03511, 2024.
Yuhang Zhou, Zihua Zhao, Haolin Li, Siyuan Du Jiangehao Yao, Ya Zhang, and Yanfeng Wang. Exploring training on heterogeneous data with mixture of low-rank adapters. arXiv preprint arXiv:2406.09677 2024.
Jinguo Zhu, Xizhou Zhu, Wenhai Wang, Xiaohua Wang. Hongsheng Li, Xiaogang Wang, and Jifeng Dai. Uni-perceiver-moe Learning sparse generalist models with conditional moes. Advances in Neural Information Processing Systems, 35:2664-2678, 2022.
Barret Zoph, Irwan Bello, Sameer Kumar. Nan Du, Yanping Huang, Jeff Dean, Noam Shazeer, and William Fedus. St-moe: Designing stable and transferable sparse expert models. arXiv preprint 0ww:22022.89066 2022.
16

---

Published as a conference paper at ICLR 2025
6 SUPPLEMENTARY MATERIAL
In our supplementary material, we provide the following details and experiments:
section 6.1: We provide more engineering implementation details about training.
section 6.2: We provide more implementation details about sample-level routing.
section 6.3: We provide more experimental results about expert loading, loss design, token gradient statistic, and different routing mechanisms.
section 6.4: We provide a analysis about computational overhead.
section 6.5: We provide experimental results about language tasks.
section 6.6: We provide a brief theoretical analysis.
section 6.7: We provide a deeper analysis about gradient conflicting phenomenon.
6.1 IMPLEMENTATION DETAILS ABOUT TRAINING
Token-level gradient: Each expert is an FFN containing multiple linear layers. For instance, the FFN in Phi-2 (Microsoft, 2023) includes two linear layers, fc' and fc². Assuming there are E x L experts, [e1: C2, CEEL], where E is the number of experts in each MoE and L is the MoE layer number of LLM. expert eG, the weights corresponding to the two linear layers are it DD and w E D'DD and the biases are b E RD' and by E IR. respectively, where is the hidden size of LLM and D' is the intermediate size.
First, during the forward pass of a batch, we freeze the parameters except for the biases and compute the main loss. Then, we perform backward pass, using the Operator sail-foo-perssmplle seppleeggads" provided by PyTorch capture the gradients g IR 2 and g E R of the token f, the biases. Next, we calculate the average gradients gmean and 22 gmean. Let the tokens processed by the expert e, be denoted as (t,, ,tN., }. the average gradients represented as:
No.
g
gmean
N
(11)
Σ Ena g
Omean
No:
gl where gmeen E R.D and gmaan E DRD Following that, we compute the cosine similarity between gi and gmean as 81 and the cosine similarity between g and gmeon as s Thus, for the token In, the similarity metric 8, defined as:
8n
(12)
Finally, we identify the conflicting token: when 8n is lower than the threshold the token t, is conflicting token.
Why to the gradients on the biases? This engineering trick brings two advantages. (i) Assume the gradients of the token in on the weights be gl. and g2." respectively. gl." DDD' and g2.w DDD, Thus, we need significant GPU memory overhead to store gradients. storing only the gradients on the biases, the GPU memory overhead significantly reduces. (ii) When the parameter size in the computational graph is larger, the backward pass longer. We compute only the gradients on the biases, so it is very fast to capture gradients. Whether does this operation work? Similar to the computation of 8n, we use g and 22. to compute the similarity metric s of the token ty We then compute the Pearson correlation coefficient between s and s, (n {1,-..) }) and find the value is usually larger than 0.9. Thus, we believe that s, reveals the relationship between the token gradient and the average gradient in the expert well. Our experiments also verify that the proposed STGC can increase performance.
After identifying conftict tokens, add the parameters that need to be updated into the optimizer (we follow the training configure of MoE-LLaVA). We add the conflict elimination loss to optimize the router based the identification of conflict tokens.
17

---

Published as a conference paper at ICLR 2025
Table 6: saper--prrmeters in training.
Epoch
Learning rate
Learning rate schedule Weight decay Cosine
Instruction Tuning
2c-5
0.0
Text max length Batch size per GPU 2048
GPU 8xA800-80G
Precision Bfl6
Instruction Tuning
16
Data
Size
Response formatting prompts
LLaVA ShareGPT
158K
40K
VQAv2
83K
Answer the question using a single word or phrase.
GQA
72K
OKVQA
9K
OCRVQA
80K
A-OKVQA
66K
Answer with the option's 's letter from the given choices directly.
TextCaps
22K
Provide one-sentence caption for the provided image.
RefCOCO
48K
Note: randomly choose between the two formats Provide short description for this region.
VG
86K
Provide the bounding box coordinate of the region this sentence describes.
Total
665K
Table 7: Instruction-follooinng data mixture. The data from LLaVA-1.5 (Liu et al., 2023b).
Training Scheme: Our training scheme follows MoE-LLaVA (Lin et al., 2024). The details are presented in Table 6. During instruction finettnning we use batch size of 128 and learning rate of 2e-5. We directly use the pre-trained models from MoE-LLaVA (Lin et al., 2024) to conduct instruction tuning.
We use LLaVA-mix-6661 (Liu et al., 2023b) as instruction tuning training data to conduct most experiments. The data structure is presented in Table To verify the scalability of the STGC model, we conducted experiments using (021K data from the Open-LLaVA-NV dataset (Lin & Long, 2024). In Table I of the main paper, we report the best performance that we achieve when activating only 3.6B parameters during inference, using 1021K data for training.
Training Datasets:
6.2 IMPLEMENTATION DETAILS ABOUT SAMPLE-LEVEL ROUTING SCHEMES
baabeedddng-bssedd Similar to MoCLE (Gou et al., 2023), we encode all the instructions of different datasets using the all-Minii.M-L66v22 variant of the Sentence Transformer model (Reimers, 2019) and cluster their embeddings via K-means clustering algorithm. After clustering, following the practice of MoCLE, we initialize K learnable embeddings, and each embedding corresponds cluster center. When sample belongs to the k-th cluster center, the k-th learnable embedding extracted and fed into the router to predict routing scores. our experiments, we set K=128. Following the practice of MoCLE, we do not add the load balance loss.
Srature-based. Similar to LoRA-MoE (Chen al., 2023d), we take the average of instruction token representations of each instance as input to predict its routing scores in each expert. Then, the Top-k experts are selected based on routing scores for each sample to generate the prediction. Following the practice of LoRA-MoE, we do not add the load balance loss.
Task-based: Similar to MoLA (Zhou et al., 2024), we promote similarity in routing between data from the same task while emphasizing distinctiveness in routing between data from different tasks. (Liu et al., 2023b) to conduct experiments, significantly different from the used data in MoLA (Zhou et al., 2024). Therefore, empirically divide the data into four types of tasks.: (i) Caption. For example, the instruction "Provide a one-sentence caption for the provided image.". (ii) VQA. For example, the instruction is "Answer the question using single
We employ LLaVA-mix-66l
18

---

Published as a conference paper at ICLR 2025
0.23
MoE-LLaVA
0.21
Load Balancing
0.19
Loss
0.17
0.15
0 500 1k 1.5k 2k 2.5k 3k 3.5k 4k 4.5k 5k
Training Step
Figure 4: Load balance loss. "Baseline" indicates MoE-LLaVA. "Baseline+STGC" indicates our method. We present the load balancing loss curve before and after adding STGC. The results obtained from the regular training. The total training step count is 5194 for an epoch. When the load balancing loss lower, the expert load is more balanced.
Table 8: Study about the load balance loss weighting a. The configure MoE-LLaVA-4Topp with StableL.M-1.6B is set for experiments.
VQA
GQA
VisWiz
SQA'
VQA"
POPE
MME
MMB
MM-Vet
Avg
MoB-LLaVA 0.01
76.7*
60.3*
36.2
62.6
50.1
85.7
1318.2
60.2
26.9
57.3
MoE-LLaVA 0.1
75.7*
59.7*
37.7
61.3
49.9
85.6
1338.0
60.4
27.2
57.2
word phrase.". (tii) OCR, including all data in OCRVQA. (iv) Region-aware. For example, the instruction is "Provide a short description for this region.". The expert label of Caption, VQA, OCR, Region-aware data is 0, or 3,, respectively. The ratio of Caption, VQA, OCR, Region-aware data is 3.5%, 61. 61.6%, 12.8%, and 22.1 %, respectively.
6.3
MORE EXPERIMENTAL RESULTS
6.3.1 EXPERT LOADING
Loss Curve: We present the the load balancing loss curve Figure 4. As shown in Figure 4, the proposed STGC benefits the decrease of the load balancing loss. 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 moved from the "crowded"" expert (many tokens) to the "emptt" expert (few tokens). This may be also new perspective on why the load balance is important the MoE system.
Additional Ablations on a: is the weighting of the load balance loss. We discuss different loss weightings E {0.01, 0.1] (0.01 is the standard value set in MoB-LLaVA (Lin et al., 2024)). As shown in Table 8, we find that increasing the weighting of the load balance loss degenerates the model performance.
Visualization of Expert Loading: we follow MoE-LLaVA (Lin et al., 2024) to obtain the distri- bution of expert loading and the visualization of the activated pathways. The distribution of expert loading examines the expert use frequency for all tokens (Lin et al., 2024). Activated pathways examine the behavior of experts at the token level (Lin et al., 2024): this visualization tool tracts validation datasets; given all activated pathways, the visu-
the activated pathways of all tokens alization tool employs PCA to obtain the top-10 pathways. As shown Figure 5, 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 gra- dient conflicts. After adding the STGC, some tokens moved from the "crowded" expert (many
19

---

Published as a conference paper at ICLR 2025
SQA
TexVVQA
MMBennc
Exeel
Expel
309%
100%
MoB-LLaVA
15%
All emporrs
All deperts
A1 eoperrts
Pryio1
Pages13
Papust
Peperi
137h
13.7%
MoE-LLaVA
+STGC
276
25%
35%
13
12 23
MoE-LLaVA STTGC
Figure 5: Expert Loading and activated pathways. The configure MoE-LLaVA-4Topp with StableLM-116BB is set for experiments We select three validation datasets, i.e., SQA (Lu al., 2022), TextVQA (Singh et al., 2019a), and MMBench (Liu et al., 2023d), to analyze expert loading and activated pathways. In activated pathways, the colorful paths represent the top-2 paths for and image, respectively, while the gray paths represent the remaining paths.
tokens) to the "empty" expert (few tokens). This further validates that STGC would effectively uti- lize each expert, instead of collapsing into only using one expert, which is also a new perspective on why the load balance important the MoE system. (ti) The activated pathways cantly different for SQA (Lu et al., 2022), TextVQA (Singh et al., 201 2a)9a), and MMBench (Liu et al.,
signifi-
2023d). 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 various experts.
6.3.2 LOss DESIGN
The goal of the conflict elimination loss is to reduce the routing score (t,) of the confticting token i, on current expert. We discuss different designs for the conflict elimination loss: (i) MSE-like: Simply setting the routing score Pmoc (fx) to the minimum.
e (fmelt...)
Pince (tn);
Σ 1= eftmee (tn)
(13)
Nat
MSE CEL
Σ
Proce (In) jid
Nall
where Nau is the count of all confticting tokens. E is the number of experts, and (tx) represents the routing score for the conflicting token t,. idmoe, is the current expert ID of the confficting token fn- (ii) CE-like: Utilizing the inverted routing score along with cross-entropy loss. Our motivation
20

---

Published as a conference paper at ICLR 2025
Table 9: Study about different conffict elimination loss designs. 4Top2 with StableI.M-1.68 is set for experiments.
The configure MoE-LLaVA-
VQA
GQA
VisWiz
SQA'
VQA
POPE
MME
MMB
MM-Vet
Avg
MSF-like
76.7*
60.7"
37.0
62.8
50.6
85.7
1346.5
60.6
27.8
57.7
CF-like
76.9*
60.9*
37.7
62.6
50.7
85.9
1355.1
60.7
28.2
58.0
1.46
Conflict Elimination Loss
1.44
1.42
1.4
0
500 1k 1.5k 2k 2.5k 3k 3.5k 4k 4.5k 5k
Training Step
Figure 6: Loss curve of conflict elimination loss. The above graph shows the loss curve of conflict elimination loss during the normal training process when Phi2-2.7B is used as the LLM. Since the total number of sampling points limited 1000 in TensorBoard, the sampling interval is set to 7.
for taking the inverted routing score minimize the routing score token on its current expert.
2/mose (tm))
Zmoo 2moo((n)) (tn),
Pmoc (tn)it
E/= 6 t (fimeett)))) (tn)))
(14)
Nau
LCEL
log (tn)t) (tno(
Natl
As shown in Table 9, the CE-like loss performs better. The reason may be that, although the opti- mization direction of the MSE-like loss is consistent with the CE-like loss, the optimization speed of the CE loss superior to the MSE loss (Zhang Sabuncu, 2018; Wang et al., 2019; Yang al., 2020). An analysis of the gradient of CE (Yang et al., 2020) reveals that when the probability of the sample on the ground- class is small, CE will produce a significantly larger gradient than MSE.
Loss Curve: Different from the statistical verification in Figure that only uses conflict elimination loss Lmoes LCEL. Lmoe, and Laux are used during the normal training process. We present the loss curve of the conflict elimination loss during training. As shown in Figure 6, the loss is convergent.
6.3.3 TOKEN GRADIENT STATISTIC
Similarity statistic: We compute the distribution of cosine similarity between the gradient of each token and the averaged gradient. Specifically, we randomly sample some data. We extract token- level gradients, calculate the average gradient gmean for each expert, and compute the cosine sim- ilarity between the gradient of each token ga and the average gradient of its current expert gmean- We perform the statistics for both initial and fully-trained models. As shown in Figure 7, we find that: (i) In the initial model, there is a conflict between token gradients and average gradients; (ii) In the traieet model, the conflict between token gradients and average gradients is reduced. This verifies that STGC can effectively increase the cosine similarity between the gradients of tokens within expert and the average gradient.
21

---

Published as a conference paper at ICLR 2025
Similarity distribution
60.0000%
50.0000%
40.0000%
Token
Ratio
30.0000%
+ Initial Model
20.0000%
Fully-trained Model
10.0000%
0.0000%
(40.7,-0.6)
(-0.6, -0.5]
(-0.5, -0.4]
(-0.4,-0.3)
(-0.4, -0.3)
(-0.3, -0.2]
(-0.3,-0.2)
(4.22 -0.1]
(40.1.0) -0.2,-0.11
(0, (0.0.1]
[0.1.02]
(0.3 (0.2, 0.3]
(0.4.0.5) 0.5]
(0.5.0.6) 0.6]
(0.6, (..6
(0.7. (0.7.088
Similarity Interval
Figure 7: Gradient similarity distribution. We compute the distribution of cosine similarity be- tween the gradient of each token and the averaged gradient. The configure StableLM-1.66 set for experiments. We add the proposed STGC to train models, to generate fully-trained model from the initial model.
with
Gradient consistency std and conflicting ratio: We further explore the std deviation of gradient consistency and the ratio of confficting tokens. Expanding on the statistical verification in the main part, we conduct a statistical verification experiment and define three metries: The first step is to gather gradients of all tokens within an expert e; and compute their cosine similarity to form similarity matrix. Then, we define the mean of the similarity matrix sim,. We define the mean of sim; on all experts as gradient consistency, serving as novel metric to evaluate whether the gradient directions are consistent within the expert. We define the std deviation of sim; on all experts as gradient consistency std, serving as metric to evaluate the degree of discreteness in the gradient consistency among different experts. Additionally, tokens within all experts as Nj. and the total number of tokens as N, defining the conflicting ratio Ni. The second step is to update models. We only use the proposed loss to update parameters, to as N. undisturbedly observe its impact on three metrics.
calculate the number of confticting
We present the curve graph the three metrics obtained from the TensorBoard dashboard in Fig- ure 8. We can observe: (i) as the conflict elimination loss decreases, the gradient consistency increases and the conflicting ratio 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 among different experts enlarges. We speculate that this is due the different rates which gradient consistency increases across var- ious layers. For example, Figure (b) indicates that deeper layers have faster gradient consistency increase rate after adding STGC. Figure 9 also shows that STGC mainly reduces conflicting tokens at deep layers, and most conflicting tokens emerge in the shallow layers after learning.
Conflicting token after learning: In Figure 7, although we observe significant reduction in con- flicting tokens for the fully-trained model, we find that there are still some conflicting tokens. As shown in Figure 9, we analyze the layers in which they appear and find that most conflicting tokens emerge in the shallow layers after learning.
6.3.4 ROUTING
SMoE (Jiang et al., 2024) claims that "Surprisingly, we do not observe obvious patterns in the as- signment of experts based the topic." As shown in Figure 5, we also observe that the distribution of expert loading across some different datasets is similar, but we notice diversity in token-level activated pathways for different datasets. We suspect that the distribution of expert loading may not
22

---

Published as a conference paper at ICLR 2025
Baseline
1.5
0.3
Baseline - STGC
Conflict
1.46
0.2
Elimination
Gradient Consitency
1.42
Randime STOC
0.1
Baeclin:
200
500
800
200
6.00
800
1k
Step
Training Step
0.1
Baselize - STOC
0.16
0.08
Contlicting
Gradient Consisinccy vid
Ratio
0.12
0.06
0.08
0.04
0.04
0.02
200
400
1k
600
800
Training Step
Training
Figure 8: Gradient consistency and conflicting ratio analysis. The configure MoE-LLaVA-44opp with StableLM-1166 set for experiments. We finish the statistic on one GPU, so the total step number is 41581. The sampling interval is set to I in TensorBoard for the above graph and the sampling interval in TensorBoard for Figure 3, so the above graph appears to have a slight difference with Figure 3.
Conflicting ratio distribution
Initial Model
Fully-trained Model
20.00%
15.00%
Conflicting Ratio
10.00%
3.00%4
0
10
12
14
16
18
20
22
Layer
Figure 9: Conflicting ratio distribution on different MoE layers after learning.
be sufficiently accurate to reflect the routing of diverse data. Then, V-MoE (Riquelme et al., 2021) and MoNE (Jain et al., 2024) focus leveraging the token importance difference to further accel- erate MoE. V-MoE proposes Batch Prioritized Routing to discard unimportant tokens and MoNE proposes Expert Preferred Router allocate more tokens to experts with larger volume. We focus on avoiding token interference during training to enhance performance, which is parallel to the focus of V-MoE or MoNE. Thus, theoretically, STGC could integrated with V-MoE or MoNE, Since MoNE does not have the official open-source code, attempt to use Batch Prioritized Routing from V-MoE for further inference acceleration.
During inference, the eval capacity of MoE is set 2.0 Lin al. (2024). As shown in Table 10, when we reduce the capacity from 2.0 to 0.5, the model performance both MoE-LLaVA and MoE- LLa SL declines significantly because many tokens are discarded. When Batch Prioritized Routing is added, there is noticeable performance improvement. We find that when the capacity is reduced 0.5, regardless of whether Batch Prioritized Routing is added or not, MoE-LLaVA--STEE
23

---

Published as a conference paper at ICLR 2025
Table 10: Study about Batch Prioritized Routing (BPR). The configure MoE-LLaVA-4Top2 with StableLM-1166B is set for experiments. BPR refers to Batch Prioritized Routing.
eval capacity
BPR
GQA
SQA!
VQA"
POPE
MME
MMB
MM-Vet
Avg
MoF-LLLVA +STGC
2.0
60.3"
62.6
50.1
85.7
1318.2
60.2
26.9
57.6
2.0
60.9"
62.6
50.7
85.9
1355.1
60.7
28.2
58.2
MoE-LLaVA +STGC
0.5
10.2"
fail
15.2
69.8
fail
6.0
16.4
0.5
21.1"
18.0
13.6
71.9
fail
7.0
20.7
MoE-LLaVA +STGC
0.5
51.0*
fail
33.4
83.8
1032,4
26,8
22.0
0.5
58.0"
57.7
44.3
85.3
1234.3
48.5
22.8
52.8
shows a significant performance improvement compared to MoE-LLaVA. A possible reason is that STGC can prevent a single expert from handling too many tokens, thereby reducing the number of discarded tokens.
6.4 COMPUTATIONAL OVERHEAD
STGC does not increase the inference overhead, while it may need the memory and time overhead during training. The memory overhead mainly results from storing gradients. The time overhead mainly results from computing gradients. We analyze the additional overhead during training from these two aspects.
Training memory overhead. We begin our analysis from StableLM-11.6. The FFN layer in StableLM-1.66 contains three linear layers, fedate E R2048x5632, fc" E R2048x5631 and fedown E R5632x2018 As stated in section 6. for each token, we only store the gradient it pro- duces on the bias within the experts. The token count per layer is about 1000, and each token only goes through one expert, so the gradient matrix stored per layer G E 11000x (0000 (66324 We need to store a gradient matrix G for each MoE layer, and the MoE layer in StableLM-1.., 1.6.-.6BBBBB is {0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22} (the total layer number is 24). Thus, the theoretical memory overhead is 12x 1000 x (5632+ 5632+ 2048). The parameter is bfloat (2 bytes), so the theoretical memory overhead is about 0.29 GB, which significantly less than the memory cost of LLM and data (usually larger than 10 GB during training). The FFN layer in Phi2-2.7B contains two linear layers, fc E 2560x 10240 and fc² E g10240X2560 and Phi2-2.7B has 16 MoE layers. Thus, the theoretical memory overhead is 16 x 1000 (2560 + 10240), i.e., 0.38 GB.
If storing the token-level gradients on each weight the required storage overhead is 12 x 1000 X (2048 5632 2048 x 5632 + 5632 x 2048) (773 GB) for StableLM-1.66 and 16 1000 x (2560 x 10240 + 10240 x 2560) (1562 GB) for Phi2-2.7B, which amazingly large. Thus, it impossible use the token-level gradient each weight conduct experiments.
Training time overhead. MoF-LLLVA, performs a forward pass, followed by backward pass to up- date parameters. As Section section 6.1 mentioned, MoE-LLaVA+STGE 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+STOO to update parameters. Thus, the main time overhead results from "the computation of token-level gradients". We directly report the train samples per second and train steps per second recorded in "trainer state.json" after training. As shown in Table 11, we have reduced the additional time over- head 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 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.
unfreeze parameters and performs backward pass
Besides, the overhead in gradient computation a common issue faced by existing gradient-basee methods. We believe that it is hopeful to address this common issue in the future, such as through gradient estimation methods (Mu et al., 2020; Paulus et al., 2020; Baydin et al., 2022; Shi et al., 2022; Wu et al., 2023).
24

---

Published as a conference paper at ICLR 2025
Table 11: Study about computational overhead. We MoE-LLaVA- with StableLM-1.6B
and Phi2-2.7B to conduct the analysis. One step (s) means the time needed for one step during training. STGC-full means computing token-level gradients on the weight.
Memory
Time train_steps -per_second
LLM
train samples .perseeoond
theoretical overhcad
one step (s)
MoE-LLaVA
19.796
0.154
6.494
StableLM-1.6B
+STGC
0.29 GB 773 GB
16.283
0.127
7.874 (+21.3%)
+STGC-full
fail
fail
fail
MoE-LLaVA
10.765
0.084
11,905
Phi2-2.7B
+STGC
0.38 GB
8.851
0.069
14.493 (421.7%) fail
+STGC-full
1562 GB
fail
fail
Table 12: Study about language tasks. Our study investigates the MoE integrated with STGC on language tasks using the GLUE benchmark (Wang, 2018), with BERT-large as the backbone model. MoE-8Top2 means traditional MoE, configured with 8 experts, of which the Top-2 are activated, i.e., 8Top2. means the re- re-implemented results.
COLA
MRPC
QNLI
MNLI
RTE
Avg
MoE-8Top2 (Guo et al., 2024)
64.5
90.2
92.4
86.7
74.9
81.7
DYNMOE (Guo et al., 2024)
65.2
90.6
92.6
86.4
73.4
81.6
MoF-8Top2" +STGC
64.5
90.0
93.4
86.9
72.9
81.5
66.8
91.2
93.8
87.6
74.7
82.8
6.5 LANGUAGE TASKS
We study the use of STGC on language tasks. Theoretically, the deployment of STGC not con- strained to specific task types. To validate the generalization of STGC, extend its to lan- guage tasks. Specifically, we follow DYNMOE Guo et al. (2024) to apply the MoE framework for language tasks. Specifically, the language tasks adhere to the same settings those in MoEffica- tion Zhang et al. (2021) and EMoE (Qiu et al., 2023). The MoE is built upon the BERT-large Devlin (2018) architecture, employing the MoEfication method, and is fine-tuned benchmark, which encompasses COLA Warstadt (2019), QNLI Wang (2018), RTE Bentivogli et al. (2009), MNLI Xu et al. (2020), and MRPC Dolan & Brockett (2005). For MoE configuration, we set the total count of experts to 8, with the Top-2 experts being activated, and we refer to this config- uration as 8Top2. We add the proposed STGC the MoE. As shown in Table 14, the experimental results show the significant effectiveness of STGC on language tasks.
GLUE Wang (2018)
6.6 THEORETICAL ANALYSIS
We conduct brief theoretical analysis based on the theory of PCGrad (Yu et al., 2020). Suppose there are tokens C, and E which pass through the same expert and generate gradients gn and ga' on that expert. Let gn V.n, En V Ln's and the parameter). cos(@nn)) the cosine similarity between gradients gu and g' Token gradient conflicts mean the cosine similarity cos(@un) < O, cos(ónn') < O potentially leads to an increase in the loss. When cos()nn))) > O, the loss decreases strictly, i.e., VL < O, which can reach the optimal value. Then, different from PCGrad, STGC does not alter the gradients of the tokens but changes the routing of the tokens to avoid conflicting gradients, with the function of increasing the cosine similarity cos(ónn)) to satisfy the condition cos(ónn') > 0.
V L = gn + ga' (VL VoL. C, where
The notation || represents the norno During each iteration, cos(ónn') > 0, a standard gradient descent step with learning rate < is employed. This results in strict reduction in the value the objective function L(0), given that the function is convex, unless the gradient (L()) is zero, which happens exclusively when 0 equals the optimal value 0* (Boyd & Vandenberghe, 2004).
25

---

Published as a conference paper at ICLR 2025
We further analyze the loss: Assuming that the gradient of the loss V is Lipschitz continuous with a Lipschitz constant L, implies that the Hessian matrix (2 - LI is negative semi-definite. Leveraging this property, we perform quadratic expansion of L around [(0), leading to the subsequent inequality:
((0 - [()) + V=(0) -0)+ [2 - 011
- [ + VL(0)" - 0) + - 0012
Then, given -+ =0- - g. the inequality can be expressed as:
[(0+) - [()) - - g" +
(Expanding, using the identity g - -n gn')
1 Lt gal 12+ +1/gor 112 + gn 'nn' = -(() (t
(Using the identity cos(dn)))
Bn gn'
ggn
gn'
1 = = (1 -1--2) L( (11ga
gn
gn/
gm II)
+ 2 cos(dnn))))
gn'
gn
(15)
By setting t < 4,, we can obtain -(1 Lt) = Incorporating the bound into the previous expression, we can deduce the following conclusion:
- and Lt <t.
[(O+) S c(0) 2
t(llgn ll (
gn
ggn
+ 2
gn
|gn
If coo(( < - |Brll
|1 (1
[|s.r]
Ig. [g.r1|
+ cocs(@nn))) will be negative. The bound of
18n
Tgall
C()+) is larger than (()), so the loss may increase. cos(@nn)) > 0, Tg.. g
|| || + cos( cos(dmnn)
will always be positive. This positivity ensures that the objective function value decreases strietly with each iteration, suggesting that by repeatedly applying this process, we can converge to the optimal value. Note that this result only holds when we select sufficiently small learning rate, specifically t 1
The goal of STGC is to increase cos( to satisfy the condition cos(Øn.)) > 0. Figure 3 and Figure 8 have verified that STGC can increase vos()nn)). Consequently, STGC is beneficial to the convergence of the objective function towards its optimal value, thereby improving the overall effectiveness of the model.
6.7 GRADIENT CONFLICTING PHENOMENON
6.7.1
RELATIONSHIP
Similar features do not necessarily imply similar gradients (Mu al., 2020). When tokens have highly similar features (similar tokens), the router can assign them to one expert. However, they may have dissimilar gradients. STGC can 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 different experts. verify this, we sample some data for analysis. Suppose there are N tokens and the cosine similarity is S NNXN between features of these tokens. We flatten S NNN to S R Meanwhile, we calculate the cosine similarity S, NNN between gradients of these tokens, flattening it to Sg R.V We compute
26

---

Published as a conference paper at ICLR 2025
Table 13: Study about feature-gradient relationship. The configure MoE-LLaVA--Top: with StableLM--1.6f is set for experiments.
Pearson correlation coefficient
MoE-LLaVA + STGC
0.2654 0.3563
Table 14: Study about gradient conflicting. zLy, denotes the average logit of tokens for their corresponding labels when calculating the main loss. zy,y, means the average logit of tokens on the labels of their conflicting tokens (meaning tokens having similar feature but divergent gradients). zi,0 denotes the average logit of tokens for all other labels, excluding the aforementioned labels.
K,yy
2,,0
Value 23.44855
14.5673
-0.6793
the Pearson correlation coefficient between S and S, to measure whether feature relationships reveal gradient relationships. We select MoE-LLaVA and MoE-LLaVA STGC for our experiments. As shown in Table 13, the experimental results indicate that token features and gradients have higher correlation after using STGC.
6.7.2 POSSIBLE REASON FOR TOKEN GRADIENT CONFLICTING
In theory, similar tokens will be assigned to the expert by a feature-based router. We focus on token gradient conflicting within an expert, so we want to further conduct study about the phe- nomenon "similar tokens have divergent gradients". In specific, we compute the cosine similarities between token features in an expert as S. We then compute the cosine similarities between token gradients an expert as Sg. We directly compute the difference S' = Sg. When S is large, token i; and token tj has high feature similarity but a low gradient similarity, meaning that the phenomenon "similar tokens have divergent gradients" significant. We find that Si is large when two tokens are confusing. Suppose labels of two tokens l; on yi are yE and Vj respectively. is the logit of each token for computing the main loss. 21,y, means the logit of token t, on yt. We define that t; and ty are confusing when 2i,y, and 2j.y, are high.
For example, in one expert, S" has the maximum at S75,51 S78,51 is 0.9316 and S/ is -0.6875. 178 has the label 13 and f51 has the label 11. We find that 17s (1s1) has second highest logit 11 (13). This is somewhat similar to the class confusion phenomenon in traditional visual classification tasks. For example, visual classification, dog and are visually similar and they have similar features, but the learning of cat may lead to the misclassification of dog.
We further sample some pairs of tokens conduct statistical verification. Each pair of tokens (ti, } has the significant feature and gradient cosine similarity difference S 21,w means the logit of token & on yı. In addition, removing Zi.y, and si, we record the mean logit 2; as 2,,0- a is the average of zi.g. As shown in the below table, 2,y, is significantly higher than 2(,0 and close to o.y. which validates the above hypothesis.
27


Reviews:
[{"review_id": "0PpeyObxvj", "reviewer": "Reviewer_U9YJ", "summary": "This paper introduces an approach to improve the Mixture-of-Expert(MOE) by eliminating token-level gradient conflict during training. The gradient of the token is compared with the average gradient to decide the conflict tokens. A Conflict Elimination Loss (CEL) is proposed as a regularization term to encourage the reassignment of conflicting tokens. The results on various model setting and tasks validate the effectiveness of proposed method.", "questions": "- Will it introduce additional training costs for the computation of the gradient and similarity?\n- It would be better to provide some analysis on the distribution of expert loadings for better understanding.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is well-written and easy to follow. The motivation for solving token-level gradient conflict during training MOE is sound to me.\n- The proposed CEL can serve as a direct plug-in to improve the training of MOEs.\n- A consistent improvement can be observed in Table 2 and 3, and the proposed STGC can be applicable to a wide range of model and data settings.\n- The evaluation of gradient consistency in Figure 3 also provides further evidence to account for the improvement.", "weaknesses": "- It would be better to illustrate the distribution of cosine similarity between gradients of all tokens and the averaged gradient for better understanding.\n- 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? \n- Why only the Large Vision-Language Model (LVLM) setting is considered? \n\n- Some minor points:\n  - Figure 3b: better label the x-axis with \"Training Step\"\n  - In TL;DR, a novel loss.\n  - Figure A in supplementary: part of the figure is not visible", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "lxtme0LVI8", "reviewer": "Reviewer_EECo", "summary": "This paper identifies conflicting token as one of the major problem in MoE-based LVLM training. Furthermore, the author propose a novel Conflicting-Elimination-Loss to resolve the conflict. By applying the proposed loss, it provides performance improvements over baseline.", "questions": "See weakness section.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "(1) The proposed approach is novel and is effective in improving the model performance.\n(2) Resolving conflicting-token demonstrates a strong correlation with improvement of the performance, which makes it a seemingly good metric to study MoE.\n(3) Thorough empirical results are provided to illustrate the effectiveness of the method. Sensitivity analysis are provided to show the robustness of the method.", "weaknesses": "(1) Visualization on the routing mechanism of the trained model should be studied to provide insights into the trained model. By conducting such visualization, we might be able to know better why the proposed method can yield performance improvements. In addition, this could also validate the proposed method would still effectively utilize every one of the expert, instead of collapsing into only using one expert.\n(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.\n(3) Computational overhead should be provided to ensure the proposed method does not incur too much overhead. This added study can show whether the proposed method is promising in terms of scaling up to large scale.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "OvF0vhUDPJ", "reviewer": "Reviewer_LbXo", "summary": "The authors make a very interesting and useful observation regarding the inconsistency/randomness in gradient direction for tokens routed to a particular expert in the Mixture of Experts framework. They then propose adding a loss term that aims to align the expert-wise token gradients to the expert-wise mean gradient over all tokens. The results presented by the authors suggest that the gradient alignment correlates with better performance, and also provide ablations showing that their method actually gets the token gradients to align.", "questions": "Please see the weaknesses section for questions as well.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "**Relevance**\n\nRouting has been a difficult task, which is pertinent to the Mixture of Experts framework. However, the understanding and analysis of the router predictions has been lacking in the literature. I believe this paper takes a significant step towards identifying and addressing a potential issue that leads to hindrance in learning in the MOE framework. The authors define the notion of a conflicting token, and then addresses it well using their proposed method, comprehensively justifying the performance in a variety of scenarios.\n\n**Presentation**\n\nI enjoyed reading the paper, it was easy to follow (except some parts listed in weaknesses), introduces the problem well, and presents detailed results and ablations to justify their proposed method. I specifically appreciate the statistical verification in Figure 3 along with the results which indicates that token conflict is actually a problem and resolving it leads to better performance.\n\n**Practicality**\n\nThe section on memory considerations and the engineering tricks for reducing memory overheads is very helpful and leads to a much practical implementation of the proposed method.", "weaknesses": "I will combine and list the weaknesses and questions here. \n\n**Practical Implementation**: While sec A in the supplementary compares the practical implementation with the heavy one in terms of Pearson correlation coefficient between the similarity metrics, I believe additional considerations \n- a quantitative measure of the memory overhead reduction (if not possible then a qualitative measure)\n- performance difference on an actual dataset\ncould make the claim of practicality even stronger. Specific experiments or analysis in the main paper that quantifies the memory savings and any potential trade-offs in performance when using their practical implementation versus the full gradient computation would be helpful.\n\n**Token Conflicts after STGC**: Fig 3 indicates the increase in gradient alignment over the course of the training process, but it would be interesting to understand the tokens that remain conflicting even after STGC. What proportion of the total tokens stay conflicting? Can something be reasoned about this behaviour? 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. The authors' insights and justifications on this consideration would be useful.\n\n**Conflict Elimination Loss**: Eqn 8 proposes the loss term, and the results do indicate its efficacy, but a more detailed motivation for the particular form of the loss could be useful. Specifically, why is $z_{moe}'(t_n)$ set to $-z_{moe}(t_n)$. Are there other forms that were considered, or insights into why this particular expression works would help.\n\n**MOE Motivation**: The authors work with the assumption that MOE frameworks reduce interference between tokens from diverse data. However, there have been previous works that mention that it is not always the case. Particularly, [1] and [2] suggest that MoE does not necessarily lead to diverse data going to different experts, or observe weak correlation of router decisions with diversity. Additionally, [3] presents a MoE variant that works against this assumption, and uses nested experts. It would be helpful to suggest how STGC could be utilized in such scenarios. It might be useful to include a discussion section in the paper that explicitly addresses these alternative perspectives on MoE behavior \n- how STGC relates to or differs from these other findings, and \n- how STGC could be adapted or extended to work with different MoE variants like nested experts.\n\n[1] Mixtral of Experts, Jiang et. al,, https://arxiv.org/abs/2401.04088\n\n[2] Scaling Vision with Sparse Mixture of Experts, Riquelme et. al, https://arxiv.org/abs/2106.05974\n\n[3] Mixture of Nested Experts: Adaptive Processing of Visual Tokens, Jain et. al, https://arxiv.org/abs/2407.19985", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "iBNmDCe2aQ", "reviewer": "Reviewer_5uJQ", "summary": "The paper introduces Solving Token Gradient Conflict (STGC), a method to address interference between tokens within experts in Mixture-of-Experts (MoE) models for Large Vision-Language Models (LVLMs). STGC uses token-level gradient analysis to identify conflicting tokens and a regularization loss to optimize token routing, reducing interference and improving model performance. Experiments show STGC's effectiveness across diverse datasets, especially with larger data diversity. The method serves as a plug-in for existing MoE-based LVLMs. The study highlights the importance of managing token interference in MoE architectures and provides a novel approach to enhance their performance in vision-language tasks.", "questions": "Can you provide an analysis of the training costs?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. Routing is the key problem when training MoE and this paper targets this important problem to reduce interference between tokens within an expert in MoE models by using token-level gradient analysis.\n2. The method demonstrated significant performance improvements on various vision-language benchmarks, indicating its effectiveness in enhancing model capabilities.\n3. STGC is designed as a plug-in, which means it can be easily integrated into existing MoE-based LVLMs without the need for fundamental architecture changes.", "weaknesses": "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.\n2. 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.\n3. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "N3fDJcAXUH", "reviewer": "Reviewer_BTYK", "summary": "This paper proposes a novel approach to address token interference in Mixture-of-Experts (MoE) architectures for Large Vision-Language Models (LVLMs). While MoE improves computational efficiency by activating only a subset of experts during inference, token interference within the same expert remains a challenge. The authors introduce Solving Token Gradient Conflict (STGC), which leverages token-level gradient selection to detect conflicting tokens and implements a conflict elimination loss to optimize token routing, thereby mitigating interference. The proposed method serves as a plug-and-play solution for various LVLM architectures. Empirical results substantiate the effectiveness of the approach, demonstrating significant improvements across multiple datasets.", "questions": "All relevant questions are included under the “Weaknesses” section. If the authors can address these concerns adequately, I will consider increasing my rating.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* **Novel Perspective on Token-Level Interference:** The paper introduces an insightful perspective on token interference by analyzing conflicting optimization directions among tokens assigned to the same expert, which is an important and under-explored issue.\n* **Direct and Effective Methodology:** The method efficiently identifies conflicting tokens through comparisons between individual token gradients and the global average gradient within experts. The regularization loss improves routing by discouraging conflicting tokens from remaining with the same expert.\n* **Comprehensive Experiments:** The experimental results are extensive and validate the proposed method’s effectiveness. The evaluations across diverse datasets highlight the robustness and general applicability of the approach.", "weaknesses": "* **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. \n* **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. \n* **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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)