Parallel-Track Mixture-of-Experts (PT-MoE)

• Parallel-Track Mixture-of-Experts (PT-MoE) is a parameter-efficient fine-tuning framework that combines low-rank prompt matrix decomposition with dynamic MoE routing.
• It leverages a shared projection matrix and expert-specific factors to reduce trainable parameters while enhancing generalization on tasks like question answering and mathematical problem solving.
• Empirical results show PT-MoE outperforms standard prompt tuning and LoRA with significant improvements in QA F1 scores and math accuracy across multiple datasets.

Parallel-Track Mixture-of-Experts (PT-MoE) is a parameter-efficient fine-tuning (PEFT) framework that combines prompt matrix decomposition with mixture-of-experts (MoE) routing to enable effective and modular adaptation of LLMs. PT-MoE extends standard prompt tuning strategies by introducing a low-rank factorization of the prompt and dynamically routing inputs through multiple expert-specific factors, demonstrating cross-task consistency and improved generalization on a range of downstream tasks, particularly question answering (QA) and mathematical problem solving (Li et al., 14 May 2025).

1. Architectural Foundations

PT-MoE operates by integrating a pair of parallel structural components—termed "tracks"—into the prompt tuning workflow. The framework leverages:

• A shared, learnable projection matrix $B \in \mathbb{R}^{R \times H}$ that is common to all experts, where $H$ is the model's hidden size and $R$ is a low-rank dimension.
• $N$ expert-specific low-rank prompt factors $A_i \in \mathbb{R}^{T \times R}$ for $i=1,\ldots,N$, where $T$ is the prompt length.
• A sparse, input-dependent routing function $R(\cdot)$ (a small neural network) that assigns weights to each expert for every sample.

The final, input-adaptive prompt is a router-weighted sum of the expert-specific factors $A_i$, multiplied by the shared $B$, and prepended as soft tokens to the frozen base LLM. This architecture allows prompt adaptation that is both parameter-efficient and dynamically specialized.

2. Prompt Matrix Decomposition

Each expert's prompt matrix $P_i \in \mathbb{R}^{T \times H}$ is factorized into a low-rank product: $P_i = A_i B \quad\text{for }i=1,\dots,N.$ Under MoE routing, if $w_i$ is the router-assigned weight for expert $i$, the aggregated prompt becomes: $$P = \sum_{i=1}^N w_i P_i = \left(\sum_{i=1}^N w_i A_i\right)B.$$ This decomposition reduces the trainable parameter count from $O(N T H)$ (if every expert had a full prompt) to $N T R + R H$, where $R \ll H$.

3. Expert Routing Mechanism

Given a pooled input embedding $x \in \mathbb{R}^H$ (typically an average of token embeddings), the router computes: $$\ell = W x + b, \quad W \in \mathbb{R}^{N \times H},\, b \in \mathbb{R}^N.$$ During training, multiplicative Gaussian noise $\epsilon \sim \mathcal{N}(0, \sigma^2)$ is applied to promote robustness: $\ell' = \ell \odot (1 + \epsilon).$ Router weights are derived via softmax, followed by a selective + probationary gating scheme: $$\tilde{w}_i = \frac{\exp(\ell'_i)}{\sum_{j=1}^N \exp(\ell'_j)}$$ Only the top-$k$ experts per input retain their weights; others are zeroed. In the probationary variant, experts' outputs are multiplied by $w_i$ before being summed, yielding confidence-weighted aggregation.

4. Forward Pass and Training Algorithm

The PT-MoE forward pass alternates between embedding extraction, router computation, expert aggregation, and prompt concatenation. The algorithmic structure is:

for each batch: E = ℳ.embed(x) # [batch, seq_len, H] μ = mean(E, dim=1) # [batch, H] ℓ = W @ μ + b # [batch, N] ℓ' = ℓ * (1 + ε) # add noise soft = softmax(ℓ') # [batch, N] for each sample j, keep only top-k of soft: w_j A_weighted = sum_{i=1}^N w_{j,i} * A_i # [batch, T, R] P = A_weighted @ B # [batch, T, H] C = concat(P, E) → ℳ.forward(C)

5. Parameterization and Efficiency Comparison

PT-MoE achieves parameter efficiency through decomposed prompt representation and modular sharing. For comparison, let $T$ = prompt length; $H$ = hidden size; $R$ = low-rank dimension; $N$ = number of experts; $r$ = LoRA rank; $M$ = number of LoRA modules:

Method	Trainable Parameters	Experimental Value (k)
PT	$T \cdot H$	81
LoRA	$M \cdot 2Hr$	106
PT-MoE	$NTR + RH$	80

PT-MoE thus employs approximately 25% fewer parameters than LoRA for comparable performance (Li et al., 14 May 2025).

6. Empirical Performance and Ablation Studies

PT-MoE delivers state-of-the-art average results on both QA (F1) and math (accuracy) tasks across 17 datasets:

Method	Params	QA F1	Math Acc.
PT	81k	56.77	46.16
LoRA	106k	56.13	56.47
PT-MoE	80k	58.26	56.91
Δ vs PT	–	+1.49	+10.75
Δ vs LoRA	–	+2.13	+0.44

Ablation highlights:

• Prompt length: Performance peaks at $T=40$ for both in-domain and out-domain tasks.
• Number of experts ($N$): Single expert is suboptimal; $N=2$ yields best in-domain results, $N=4$ for out-domain.
• Trainable parameter count: Performance increases with more parameters (18k to 163k), saturating near 80k.
• Routing mechanism: Selective + probationary routing improves F1 by 1–2 points over alternatives. Probationary expert weighting enhances training stability.

7. Modular Design, Cross-Task Generalization, and Future Directions

PT-MoE’s modular structure, combining a shared output factor $B$ and specialized expert factors $A_i$ with dynamic MoE routing, enables efficient parameter sharing and input-dependent specialization. This arrangement supports consistent improvements in both QA and mathematical reasoning—tasks that favor conventional PT and LoRA, respectively. A plausible implication is that the cross-task consistency observed derives from the synergy between low-rank sharing (via $B$) and specialized adaptation (via $A_i$), realized through dynamic routing.

Recommended future extensions include:

• Hierarchical/multi-level router architectures to capture multi-granular task specialization.
• Application to continual or multi-task learning frameworks by modularly integrating new experts.
• Exploration of structured low-rank decompositions (e.g., block-diagonal forms) for further gains in efficiency (Li et al., 14 May 2025).