Tensor-Train LoRA MoE: Efficient Multi-Task Learning

• TT-LoRA MoE is a parameter-efficient framework that integrates tensor-train decomposed LoRA adapters with a sparse Mixture-of-Experts router to enable scalable multi-task and continual learning.
• It decouples expert specialization from expert selection, significantly reducing storage and computational overhead compared to dense or fused adapter methods.
• Empirical results demonstrate enhanced transfer performance with multi-task accuracy improvements and effective prevention of catastrophic forgetting.

Tensor-Train LoRA Mixture-of-Experts (TT-LoRA MoE) is a parameter-efficient fine-tuning (PEFT) framework for large pre-trained models, integrating tensor-train decomposed low-rank adapters with sparse mixture-of-experts (MoE) routing. Designed primarily for scalable multi-task and continual-learning scenarios, TT-LoRA MoE achieves substantial reductions in storage and computational overhead by decoupling expert specialization from expert selection, leveraging compact TT-LoRA adapter modules and a lightweight, trainable router (Kunwar et al., 29 Apr 2025). This combination automates expert selection per input, enables efficient multi-task inference without catastrophic forgetting, and outperforms dense and fused adapter approaches in both parameter efficiency and transfer performance.

1. Foundational Components and Principles

TT-LoRA MoE unifies two principal innovations: (i) parameter-efficient adapters constructed using tensor-train (TT) decompositions (TT-LoRA), and (ii) sparse dynamic expert selection via a MoE-style router.

TT-LoRA Adapter: Instead of a conventional LoRA update (i.e., $\Delta W = A B^\top$), TT-LoRA represents the weight update as a tensor-train decomposition. This entails expressing a weight matrix $W_0 \in \mathbb{R}^{m \times n}$ as a $(p+q)$-way tensor with factorized TT cores $\{G_k\}$: $$\Delta W \approx \sum_{i_1,\ldots,i_d=1}^r A^{(1)}_{i_1} \otimes A^{(2)}_{i_2} \otimes \cdots \otimes A^{(d)}_{i_d}$$ with the forward pass computed using sequential tensor contractions rather than reconstructing dense matrices, yielding significant memory and compute reductions (Kunwar et al., 29 Apr 2025).

Sparse MoE Router: The router is trained, after all adapters are frozen, to select exactly one expert per input using a deterministic or noisy top-1 gating mechanism. It operates on base model representations prior to the projection head and only affects the routing weights $W_{\rm gate}, W_{\rm noise}$, not the adapters themselves.

2. Architectural Overview and Workflow

TT-LoRA MoE employs a two-stage pipeline:

Stage 1: Expert Pretraining

• For each task $i \in \{1, \ldots, N\}$, a distinct TT-LoRA adapter $\phi_i$ is tuned with the base model $\Theta$ frozen.
• Independent adaptation prevents inter-task interference and catastrophic forgetting.
• Adapter parameters—TT cores $\{G_k^{(i)}\}$—are stored and then frozen post-convergence.

Stage 2: Router Training

• All frozen adapters and the base model are loaded.
• A lightweight router network is trained on a mixed multitask dataset, using only the last base hidden state as input.
• The router uses noisy top-1 gating during training ($W_0 \in \mathbb{R}^{m \times n}$0) and deterministic argmax routing at inference. Only the selected expert processes each input, ensuring constant inference cost irrespective of expert count.

This staged approach sharply contrasts with traditional MoE architectures, where all experts are co-trained, often resulting in capacity dilution and costly joint optimization (Kunwar et al., 29 Apr 2025).

3. Mathematical Formulation

TT-LoRA Expert Computation:

Given input $W_0 \in \mathbb{R}^{m \times n}$1, the adapter output is computed as: $W_0 \in \mathbb{R}^{m \times n}$2 with the final layer output: $W_0 \in \mathbb{R}^{m \times n}$3 where $W_0 \in \mathbb{R}^{m \times n}$4 is a LoRA scaling factor.

Sparse Routing:

Given the hidden state $W_0 \in \mathbb{R}^{m \times n}$5, scores and noise scales are calculated as: $W_0 \in \mathbb{R}^{m \times n}$6 Noisy logits $W_0 \in \mathbb{R}^{m \times n}$7 provide robustness during training, while at inference, $W_0 \in \mathbb{R}^{m \times n}$8 is used for exact selection. The final multi-task prediction uses only the selected expert: $W_0 \in \mathbb{R}^{m \times n}$9 Router and task losses are combined: $(p+q)$0

4. Parameter and Compute Efficiency

TT-LoRA MoE delivers significant parameter and computational savings compared to alternative PEFT and MoE models:

Adapter/Fusion	Params (per expert)	Fusion/Router Params	Relative Params (%)
LoRA ($(p+q)$1)	1.7M	--	100
TT-LoRA (rank=8)	33.9K	--	2
Pfeiffer Adapter	12.6M	--	>700
AdapterFusion	1.7M	205M	12,000
TT-LoRA MoE	33.9K	69K	0.03 (router)

TT contraction avoids explicit dense matrix reconstruction, reducing both FLOPs and peak memory by approximately 1.5–2× compared to matrix-based LoRA updates. The router’s parameter count makes it negligible relative to fusion-based methods, and top-1 sparse gating ensures per-token cost is expert-invariant (Kunwar et al., 29 Apr 2025).

5. Training Algorithm and Hyperparameters

Stage 1 (Expert Training) Pseudocode:

• For each downstream task $(p+q)$2, initialize TT cores, freeze $(p+q)$3 and the classification head.
• Fine-tune on $(p+q)$4 using Adam ($(p+q)$5), batch size 32.
• Recommended TT-ranks: 8–16.
• $(p+q)$6.

Stage 2 (Router Training) Pseudocode:

• Construct a balanced multi-task dataset.
• For each minibatch, use only the router parameters ($(p+q)$7, Adam $(p+q)$8; $(p+q)$9), batch size 64.
• Recommended: noisy top-1 at train, deterministic argmax at inference.

Implementation Recommendation: Freeze experts and routers for deployment; precompute base model states if evaluating multiple experts.

6. Empirical Performance

On 17 datasets, TT-LoRA achieves 33.9K parameters per expert, with LoRA at 1.7M and Pfeiffer Adapter at 12.6M (Kunwar et al., 29 Apr 2025). TT-LoRA MoE demonstrates:

• Single-Task: LoRA 80.99% avg. accuracy; TT-LoRA 79.58%.
• Fusion Benchmark: AdapterFusion avg. 75.16% (205M params), TT-LoRA MoE 79.04% (69K router params).
• Multi-task (mixed, Top-10): AdapterFusion 81.45%; TT-LoRA MoE 85.91% (+4.46 pts).

A plausible implication is that TT-LoRA MoE, despite using only $\{G_k\}$0 of AdapterFusion’s parameters for expert fusion, yields higher multi-task accuracy and avoids interference between tasks during continued scaling.

7. Relation to Modular Tensor Product MoEs

Concurrent with TT-LoRA MoE, “Mixture of Latent Experts Using Tensor Products” (TensorPoly) develops a modular MoE framework where LoRA adapters are reparameterized as entangled higher-order tensors (Su et al., 2024). TensorPoly introduces two routing strategies—TensorPoly-I (rank-level) and TensorPoly-II (order-rank level)—that allow nuanced mixtures of latent tensor slices for flexible, parameter-efficient multi-task transfer. Ablation studies indicate that both tensorized modules and data-dependent routing are essential for maximizing transfer and efficiency. TensorPoly-I (N=2, R=20) achieves state-of-the-art few-shot transfer with $\{G_k\}$169.3 accuracy on T0-benchmarks, using 1.4M parameters per adapter—substantially less than dense alternatives.

Both TT-LoRA MoE and TensorPoly demonstrate that modularization and tensor-based decomposition of adaptation layers, coupled with advanced routing, yield favorable trade-offs between transfer performance, interference mitigation, and efficiency in the multi-task setting (Kunwar et al., 29 Apr 2025, Su et al., 2024).