T-REX: Mixture-of-Rank-One-Experts with Semantic-aware Intuition for Multi-task Large Language Model Finetuning

1 Introduction

• Multitask MoE faces significant parameter and computational costs as expert counts scale with tasks, alongside suboptimal allocation by traditional routers.Source: §1 · pp. 1–2
• Rank-1 vectors pair freely into LoRA experts, expanding the covered weight subspace with sub-linear additional parameter growth.Source: §1 · p. 2
• Training-embedding clusters provide implicit intuition that conditions router decisions and guides experts toward smooth convergence and generalizability.Source: §1 · p. 2
• Up to 1.78% higher overall accuracy with around 30%-40% less parameter overhead than other LoRA-based methods.Source: §1 · p. 2
• Theoretical analysis shows that expert linear-combination subspace expansion enables approximate error reduction and smoother expert convergence.Source: §1 · p. 2

Multi-task learning

• Multitask methods use hard parameter sharing, adaptive loss weighting, or dynamically blended task representations, while T-REX allocates inputs among specialized experts.Source: §2 · p. 2

Mixture of Experts

• Transformer MoE work has scaled models through top-1/2 routing and explored training stability, variable expert assignment, and parameter-efficient experts.Fully differentiable sparse MoE also addresses training instability and token dropping.Source: §2 · pp. 2–3

Parameter-efficient finetuning

• LoRA-MoE methods improve flexibility but face a balance between parameter efficiency and scaling to more experts; T-REX uses rank-1 mix-and-match experts.Source: §2 · p. 3

3.1 Preliminary

• The multitask setup adapts a pretrained weight matrix across K distinct tasks, while LoRA-MoE dynamically selects and combines low-rank experts.The stated goal is efficient adaptation while avoiding catastrophic forgetting.Source: §3.1 · p. 3

3.2 Mix-and-Match for rank-1 experts

• Pairing row and column rank-1 vectors expands the expert pool quadratically, while additional parameters grow linearly.The routing matrix dynamically activates cross-combined experts.Source: §3.2 · pp. 3–4

3.3 Semantic-aware router with intuition clustering

• T-REX clusters training-token embeddings and uses similarity to cluster centroids as intuition scores fused with base routing weights.The fused routing signal informs the final expert combination.Source: §3.3 · p. 4

4.1 Rank-1 experts with subspace expansion

• Pairing any a_i with any b_j lets the Kronecker basis span up to I × J dimensions, although rank(∆W) ≤ min(I, J).The analysis states that the subspace dimensionality grows at O(IJ) as ranks I and J increase.Source: §4.1 · p. 5
• T-REX decomposes a target adaptation matrix into in-subspace and orthogonal residual components, with alignment error decreasing as IJ increases.The approximation account also identifies the residual outside the subspace and a routing-temperature contribution of O(1/τ).Source: §4.1 · p. 5
• The routing softmax temperature τ controls virtual-expert activation sharpness and contributes an approximate O(1/τ) error term.The passage states that increasing τ makes routing more selective and reduces error fluctuation.Source: §4.1 · p. 6

4.2 Intuition leads to low-loss region

• Under an orthonormal input basis, the cosine-based intuition vector is proportional to coordinates of the optimal adaptation, guiding the router toward smoother convergence.The paper presents the intuition vector as good initialization and guidance for decomposing each input’s optimal adaptation across mix-and-match rank-1 experts.Source: §4.2 · p. 6

5.1 Implementation Details

• The experiments use 20 multitask datasets spanning reasoning, examination, language, understanding, and knowledge, plus Alpaca’s mixed dimensions.The MoE module targets query, key, value, and feed-forward components; instance embeddings use nomic-embed-text-v1.Source: §5.2 · p. 7, §5.1 · p. 7

5.2 Multitask Datasets and Baselines

• Baselines include LoRA-32, MoLoRA, and SiRA with parameter-matched expert settings; T-REX uses 32 rank-1 experts formed from rank-4 and rank-8 LoRA weights.Experiments use batch size 64, a 5e-5 learning rate, cosine scheduling, Adam, and three epochs.Source: §5.2 · p. 7, §5.1 · p. 7

5.3 Quantitative Results

• On Mistral 7B, T-REX uses 39.9M additional parameters and 5.13G FLOPs, runs at 6.92 iterations per second, and raises accuracy from 83.08% to 84.86%.The reported training overhead is lower than LoRA and MoLoRA; SiRA has similar overhead to MoLoRA.Source: Table 2 · p. 7, §5.3 · p. 7
• T-REX achieves the highest average score across all models on 14 datasets and seven backbones, with fewer additional trainable parameters than the baselines.Reported examples of additional trainable parameters are +0.62% for LLaMA-2 13B and +0.76% for Mistral 7B.Source: §5.3 · p. 7, Table 1 · p. 8
• With five sub-datasets randomly excluded, T-REX improves both ID and OOD performance; in Drop-A, the gains are 2.48% and 1.55%, respectively.The excluded Drop-X datasets are treated as OOD and the remaining datasets as ID.Source: Table 3 · p. 8, §5.3 · p. 8
• In the reported BBH results, T-REX achieves accuracies of 40.97% on Gemma 2B and 47.59% on Yi 6B, consistently surpassing LoRA by over 1%.Figure 3 compares T-REX with LoRA, MoLoRA, and SiRA on BBH for Gemma 2B and Yi 6B.Source: §5.3 · p. 9, Figure 3 · p. 9

5.4 Ablation Studies

• Among the tested mix-and-match configurations, 8×8 reaches 70.75% accuracy but doubles FLOPs; 2×16 reaches 67.52%, so 4×8 is the default efficiency setting.The authors report better performance when the low-rank matrices’ ranks are closer.Source: Table 4 · p. 9, §5.4 · p. 9
• Implicit intuition and rank-1 experts improve LLaMA-2 7B by up to 1.12% and 0.70%, respectively; their combination gains 1.35% on Tiny-LLaMA.For Mistral 7B, adding mix-and-match provides a further 1.02% performance boost while reducing parameter overhead by 40.6%.Source: Table 5 · p. 9, §5.4 · p. 9

6 Conclusions

• T-REX combines semantic-aware implicit intuition with Mix-and-Match rank-1 experts, with experiments reporting superior performance in both in-distribution and out-of-distribution scenarios.The supplementary material describes evaluations of dropped training tasks, different Mix-and-Match strategies, component effects, and BBH results.Source: §6 · p. 9

Appendix for Mixture-of-Rank-One-Experts

• The supplementary material includes a detailed approximation-bound derivation and analyses of intuition, dropped training tasks, Mix-and-Match strategies, component effects, and BBH results.It also identifies Appendix G as the location of the limitations discussion.Source: Appendix for Mixture-of-Rank-One-Experts · p. 13

A.1 Subspace Decomposition of ∆W ⋆

• The target adaptation matrix decomposes into an in-subspace component spanned by rank-1 components and an orthogonal residual.The decomposition follows from the projection theorem.Source: §A.1 · p. 13

A.2 Approximation of ∆Win by T-REX

• Increasing IJ improves subspace coverage and reduces the in-subspace approximation term compared with traditional LoRA's N-dimensional subspace at rank N.The T-REX adaptation matrix uses a dynamic routing matrix and learned basis vectors; its residual error within the subspace depends on alignment with that subspace.Source: §A.2 · pp. 13–14

A.3 Residual Error ∆Wout

• The out-of-subspace residual cannot be captured by rank-1 combinations, so increasing IJ cannot reduce its contribution to total error.That contribution may be negligible in practice when the target matrix aligns well with the learned subspace.Source: §A.3 · p. 14

A.4 Impact of Routing Matrix G(x)

• The routing matrix dynamically activates rank-1 component combinations, while softmax temperature controls activation smoothness and contributes an O(1/τ) term to approximation error.Higher temperature yields smoother activations across components, whereas lower temperature sharpens focus on specific components.Source: §A.4 · p. 14

A.5 Final Error Bound

• The final bound combines residual error outside the subspace, an in-subspace approximation term that decreases as IJ increases, and an O(1/τ) routing contribution.The routing contribution reflects dynamic activation of rank-1 components.Source: §A.5 · p. 14

B Motivations

• The paper frames useful AI intuition as capturing instance-level semantic variability, being autonomously derived, and integrating into LLMs with low computational overhead.These criteria are stated as the motivation for integrating implicit intuition into multitask learning.Source: B Motivations · p. 14
• Embedding visualizations of 20 datasets show that semantic clusters do not align with predefined task groupings, and human-annotated task categories underperformed baseline fine-tuning.The authors use this mismatch to motivate embedding multitask knowledge implicitly through an intuition score.Source: B Motivations · p. 15, Figure 4 · p. 15

C Dropping out training data

• Across four data configurations, T-REX consistently improves average accuracy for both in-distribution and out-of-distribution tasks on Tiny-LLaMA 1B.Each configuration randomly excludes five training datasets; the table caption indicates excluded datasets are gray.Source: C Dropping out training data · p. 15, Table 6 · p. 15
• The four configurations exclude five datasets apiece, representing 24% to 44% of total training samples.The excluded sets differ across experiments A–D.Source: C Dropping out training data · p. 15

D Analysis of different Mix-and-Match strategies

• 8×8 achieves the highest overall accuracy and ranks first on 9 of 14 tasks, while 4×8 leads on RTE, COPA, and CSQA and is selected for subsequent experiments.Although 8×8 is marginally stronger on average, it has about 50% more parameter overhead and correspondingly higher compute cost than 4×8.Source: D Analysis of different Mix-and-Match… · p. 16

E Synergistic Effects of intuition and rank-1 experts

• Across three LLMs ranging from 1 billion to 7 billion parameters, implicit intuition and rank-1 experts individually boost performance, and their combination consistently delivers optimal results on average accuracy across 14 datasets.Source: E Synergistic Effects of intuition and… · p. 16

• 1.52% enhancement over competing methods is reported for T-REX on the BBH dataset.Source: F Big-Bench Hard results · p. 16
• T-REX surpasses LoRA, MoLoRA, and SiRA on BBH in the reported OOD comparison.Source: F Big-Bench Hard results · p. 16
• The BBH evaluation uses Yi 6B, Gemma 2B, Mistral 7B, Llama 2 7B, Tiny-LLaMA 1B, and Phi-2 without finetuning on BBH samples.Source: F Big-Bench Hard results · p. 16

G Limitations

• Experiments are limited to LLMs, although the authors believe T-REX could extend to VLMs with minimal modifications.The paper does not report experiments on other modalities.Source: G Limitations · p. 17
• Mix-and-Match increases the number of experts and router size, but low-rank router matrices can reduce parameter overhead without significantly affecting performance.Source: G Limitations · p. 17

H Broader impacts

• By reducing computational and parameter overheads, T-REX may broaden access to high-quality language models in resource-constrained domains such as education, healthcare, accessibility, and multilingual applications.Source: H Broader impacts · p. 17
• Lower resource requirements may accelerate LLM proliferation and exacerbate misinformation, malicious content generation, or training-data biases; flexible models may also reach high-stakes contexts requiring robustness and transparency.Source: H Broader impacts · p. 17