PaperlayerView exact revision on arXiv

Source-linked AI summary

GaP-MoE : Explicit Expert Specialization in Mixture-of-Experts via Orthogonal Gradient Accumulation and Projection

Donghyun Jeong, Yoon-Yeong Kim

arXiv:2610.29355v1cs.LGstat.ML

TL;DR

Sparse Mixture-of-Experts can scale model capacity, but specialization methods may modify routing indirectly or add objectives that conflict with the task loss. GaP-MoE instead projects follower output gradients relative to a selected leader, and the reported experiments show competitive or improved performance over SMoE.

  • Problem

    Specialization methods may act indirectly through routing or use auxiliary objectives that can conflict with the original task loss.

    Source: §1 · p. 2

  • Method

    GaP-MoE retains standard SMoE training and orthogonalizes already-computed follower output gradients relative to a dynamically selected leader’s recent gradients.

    Source: §1 · p. 2

  • Results

    GaP-MoE achieves competitive or improved performance over baseline SMoE and can be combined with existing specialization methods.

    Source: §5 · p. 10

  • Takeaways & Limitations

    The method reuses computed gradients without gradient recomputation and is compatible with different MoE architectures.

    Source: §1 · p. 2, §5 · p. 10

    The evaluation used relatively small-scale models, and GaP-MoE hyperparameters were not systematically explored.

    Source: §5 · p. 10

Abstract

from arXiv · show

Sparse Mixture-of-Experts (SMoE) models improve scalability by routing input tokens to only a subset of expert networks, but they often suffer from representation collapse, in which multiple experts converge toward similar functions and lose their intended diversity. Prior work has addressed this limitation either by improving router architectures or by introducing auxiliary loss terms that implicitly penalize similarity among experts. However, router-based approaches remain vulnerable to routing collapses, and loss-based approaches can conflict with the task's original objective during optimization. In this work, we propose GaP-MoE (Gradient Accumulation and Projection in MoE), which induces expert specialization by directly manipulating expert output gradients rather than modifying the loss function. At each training step, GaP-MoE stochastically selects one of the $TopK$ routed experts as a leader based on normalized routing scores, and projects the gradients of the remaining follower experts onto the orthogonal complement of the subspace spanned by the leader's recent gradients. This removes directional overlap between expert gradients without altering the task loss, allowing specialization to emerge as a direct consequence of the optimization dynamics rather than an externally imposed constraint. We evaluate GaP-MoE in both fine-tuning settings (BERT-based, on GLUE) and from-scratch pretraining settings (GPT-2-based, on WikiText-103, enwik8, and text8), and further confirm its compatibility as a drop-in module for existing SMoE architectures. Experimental results show that GaP-MoE achieves performance comparable to or better than standard SMoE.

1 Introduction

Scaling model and dataset sizes can improve performance but increases computational cost, motivating sparse expert models whose experts do not necessarily specialize. GaP-MoE addresses this by orthogonalizing expert output gradients without adding a loss term.

1 Introduction

  • SMoE uses TopK routing to select experts for each token and combines their outputs using routing-score weights.This approach increases model capacity through multiple independent expert networks.Source: §1 · p. 1
  • Experts may remain insufficiently specialized because of router collapse or because they receive similar inputs or share similar initialization conditions.The introduction frames specialization as important for fully using SMoE capacity.Source: §1 · p. 2
  • Router-based methods encourage specialization indirectly, while auxiliary-loss methods may conflict with the task objective.GaP-MoE instead directly modifies expert gradients after they are computed.Source: §1 · p. 2
  • GaP-MoE applies orthogonality constraints to already-computed expert output gradients while retaining standard SMoE training and avoiding an additional loss term.It modifies redundant gradient directions after task gradients are generated.Source: §1 · p. 2
  • GaP-MoE probabilistically selects a leader among TopK experts using normalized routing scores and projects follower gradients away from the leader’s accumulated gradient subspace.The leader keeps its task gradient; the method reuses computed gradients without extra forward or backward passes.Source: §1 · p. 2
  • GaP-MoE is described as compatible with existing router architectures and auxiliary loss functions.The paper also identifies gradient orthogonalization and reuse of computed output gradients as benefits.Source: §1 · p. 2

2 Related Works

The related work describes sparse routing through independent experts and reviews representation collapse and approaches to expert specialization. GaP-MoE is distinguished by directly projecting follower update directions away from a dynamically selected leader’s recent gradient subspace.

2.1 Sparse Mixture-of-Experts

  • SMoE uses independent experts and a router that selects the K highest-scoring experts for each token, then combines their outputs using the retained routing scores.Only selected experts are activated for each token, yielding sparse computation.Source: §2.1 · pp. 2–3
  • Figure 1 depicts leader selection from the TopK set using normalized routing probabilities and aggregation of the leader’s output gradients.When the most recent basis-vector similarity exceeds γ, the follower gradient is projected onto the orthogonal complement of the leader subspace.Source: Figure 1 · p. 3, §2.1 · p. 3

2.2 Expert Specialization and Representation Collapse

  • Representation collapse is used for insufficient differentiation among experts, including similar expert representations or routing most tokens to only a few experts.Expert specialization refers to making individual experts learn distinct representations or functions.Source: §2.2 · p. 3
  • Prior approaches include low-dimensional router representations, perturbed cosine routing, expert-parameter or output-similarity losses, and SVD-based expert-specific subspaces.These approaches address specialization through routing, expert weights, outputs, or gradients.Source: §2.2 · p. 3
  • GaP-MoE differentiates expert update directions by projecting follower gradients away from a dynamically selected leader’s recent gradient subspace, without an additional optimization objective.The selected leader retains its original task gradient.Source: §2.2 · p. 3

3 Proposal of GaP-MoE

GaP-MoE selects a leader among co-routed experts, accumulates its recent output-gradient directions, and projects follower gradients away from their shared subspace. The method uses this gradient-level separation to reduce aligned expert updates without directly constraining output representations.

3.1 Dynamic Leader Selection

  • Normalized TopK routing scores determine the probability of selecting one expert as leader; the remaining selected experts become followers.The leader keeps its original task gradient, and follower gradients are differentiated from its gradient subspace.Source: §3.1 · p. 4

3.2 Gradient Projection

  • Projection activates according to the similarity threshold γ and removes the follower gradient’s component in the leader’s buffered-gradient subspace.With γ = −1, projection is activated for every leader-follower pair; α controls projection intensity.Source: §3.2 · p. 4

3.3 Gradient Accumulation

  • Each expert’s buffer stores recent gradient directions; mini-batch leader gradients are aggregated, orthogonalized against the existing basis, and appended.When the buffer reaches size h, its oldest vector is removed; projected follower gradients are excluded from buffer updates to avoid reflecting altered directions immediately.Source: §3.3 · p. 5

3.4 Theoretical Justification of GaP-MoE

  • GaP-MoE projects follower output gradients away from a dynamically selected leader’s recent gradient subspace, reducing shared update directions while preserving the leader’s task gradient.The method modifies already-computed gradients rather than adding an optimization objective.Source: §3.2 · p. 4, §3.4 · pp. 6–7
  • Theoretical analysis states that projection can reduce the post-update inner product between leader and follower outputs, subject to efficient MoE training.The method avoids directly minimizing similarity, which could push representations in opposite directions and potentially hinder optimization through cancellation in the weighted sum.Source: §3.4 · pp. 6–7

4 Experiments

The experiments test GaP-MoE against SMoE in fine-tuning, from-scratch training, router-free training, and integrations with other MoE methods. They also examine expert-output differentiation and short-term gradient consistency.

4 Experiments

  • Experiments compared GaP-MoE with SMoE across fine-tuning and from-scratch settings, using four-expert MoE layers with Top2 routing in the baseline models.The fine-tuning experiments used BERT on GLUE; the from-scratch experiments used a reduced GPT-2 configuration.Source: §4 · p. 7, §4.1 · p. 7, §4.2 · p. 8

4.1 Fine-tuning

  • Applying GaP-MoE to all 12 fine-tuning MoE layers yielded consistent but modest average improvements over SMoE; continuous projection performed better on average than conditional projection.The conditional variant projected only when the last basis similarity was high.Source: §4.1 · p. 7

4.2 From-scratch

  • GaP-MoE generally performed better than SMoE from scratch, with a less significant improvement on text8 than on WikiText-103 and enwik8.The text8 dataset retains only alphabetic characters and spaces from enwik8.Source: §4.2 · p. 9

4.3 Expert Specialization Without Router Training

  • With an untrained router and two randomly selected experts per token, performance fell overall, but the drop was considerably smaller for GaP-MoE than for SMoE.Table 3 reports WikiText-103 results averaged over three seeds; lower scores are better.Source: Table 3 · p. 8, §4.3 · p. 9

4.4 Generalizability with Different MoE

  • Across CoLA, STS-B, and RTE, GaP-MoE improved performance on most tasks with XMoE and SimSMoE, but results were more sensitive to γ than with GaP-MoE alone.The reported metrics are Matthews correlation for CoLA, Pearson correlation for STS-B, and accuracy for RTE; results are averaged over three seeds.Source: Table 4 · p. 8, §4.4 · p. 9

4.5 Expert Output Differentiation

  • Setting γ to −1 substantially decreased expert-output similarity, while γ = 0.7 produced only a marginal reduction; the figures examine token-wise similarity and expert-output angles.Figure 3 covers tokens in a COLA evaluation sequence; Figure 4 shows sixth-layer angles for a QNLI token, distinguishing Top2 outputs from unselected experts.Source: Figure 3 · p. 8, Figure 4 · p. 8, §4.5 · p. 9

4.6 Short-term Gradient Consistency

  • Gradient consistency fluctuated early, then stabilized; with wider gaps, early-layer similarity approached zero while later layers remained consistently high.The study tracked cosine similarity to accumulated bases at gaps of 1, 3, 5, and 10 steps and recommends gradually increasing projection intensity during initial training.Source: §4.6 · p. 9, Figure 5 · p. 10

5 Conclusion and Future Work

GaP-MoE applies orthogonalization to precomputed expert output gradients while retaining standard sparse-MoE TopK training. The reported results are competitive or improved versus baseline SMoE, but hyperparameters were not systematically explored and evaluation used relatively small-scale models.

5 Conclusion and Future Work

  • GaP-MoE retains standard sparse-MoE TopK training while orthogonalizing precomputed expert output gradients.It uses dynamic leader selection and gradient projection without adding parameter constraints or requiring gradient recomputation.Source: §5 · p. 10
  • GaP-MoE achieves competitive or improved performance over baseline SMoE and is compatible with existing specialization methods and different MoE architectures.Source: §5 · p. 10
  • The evaluation used relatively small-scale models, and GaP-MoE hyperparameters were not systematically explored, so reported settings may not be optimal.The listed hyperparameters include buffer size h, fine-tuning layer choice, projection threshold γ, and projection intensity α.Source: §5 · p. 10

A Algorithm of GaP-MoE

The algorithm selects and weights a top-2 expert pair for each token, then samples one expert as leader and treats the other as follower. It computes the MoE output and gradients, projects follower gradients using prior buffers, and updates expert weights and buffers.

A Algorithm of GaP-MoE

  • For each token, the algorithm selects the top-2 experts, normalizes their routing scores, samples a leader, and assigns the other expert as follower.Source: A Algorithm of GaP-MoE · p. 13
  • The selected experts process the token, their outputs form the weighted MoE output, and backpropagation computes output gradients.Source: A Algorithm of GaP-MoE · p. 13
  • Follower gradients are projected using buffers from before the current step, and expert weights are updated using leader and projected-follower gradients.Source: A Algorithm of GaP-MoE · p. 13
  • For each expert, the algorithm performs accumulation and updates the buffer for the next step; it also updates Qi with bnew.Source: A Algorithm of GaP-MoE · p. 13

B Proof of Proposition 1

The proof expresses expert parameters through accumulated gradient updates beginning when short-term gradient consistency holds. It argues that orthogonalizing output gradients reduces alignment between expert outputs, while fully opposite directions can obstruct training with linearly combined outputs.

B Proof of Proposition 1

  • The parameter at time t is expressed by accumulating updates from t0, the earliest time short-term gradient consistency holds.Source: B Proof of Proposition 1 · p. 14
  • Under the stated assumption, accumulated updates yield a parameter expression in terms of the initial parameter and a scalar update magnitude times the current unit gradient.Source: B Proof of Proposition 1 · p. 14
  • Orthogonalizing output gradients reduces alignment between expert output vectors, but completely opposite directions can obstruct training when MoE linearly combines expert outputs.Source: B Proof of Proposition 1 · p. 14

C.2 Experimental Settings

The experimental settings report an NVIDIA RTX A6000 D6 with 48GB of memory. Fine-tuning used five epochs, a learning rate of 2e-5, basis rank 16, FFN intermediate size 3072, and evaluation batch size 64.

C.2 Experimental Settings

  • The reported hardware was an NVIDIA RTX A6000 D6 with 48GB of memory.Source: Fine-tuning · p. 15
  • Fine-tuning used learning rate 2e-5, five epochs, basis rank 16, FFN intermediate size 3072, and evaluation batch size 64.Source: Fine-tuning · p. 15

D Expert Output Angles Across All Layers

The section visualizes expert-output angles across layers for a QNLI test input, comparing SMoE with two GaP-MoE settings.

D Expert Output Angles Across All Layers

Related papers

Loading 2610.29355v1…