TLDR: This research introduces Adaptive Shared Experts (ASE) within a LoRA-based Mixture-of-Experts framework for multi-task learning. ASE uses router-computed, jointly normalized gating weights for shared experts, enabling dynamic influence that aids early knowledge transfer and reduces gradient conflicts. Combined with efficient LoRA experts and fine-grained designs, it significantly improves multi-task performance and expert specialization on benchmarks like PASCAL-Context.
Multi-Task Learning (MTL) is a powerful approach in artificial intelligence where a single model learns to perform several tasks simultaneously. This method is inspired by how humans can learn multiple related skills, allowing models to leverage shared knowledge across tasks and achieve better performance and efficiency. A key architectural innovation enabling this is the Mixture-of-Experts (MoE), which expands a model’s capacity by activating only a small subset of specialized “experts” for each input, keeping computation manageable.
However, current MoE-based MTL methods often face challenges. They typically start with backbones trained for single tasks, leading to inefficiencies when adapting to multiple tasks. Experts might redundantly relearn similar information, and the model struggles to capture common features effectively. This can compromise both efficiency and the ability of experts to specialize in their respective tasks.
A natural idea to overcome this is to introduce a “shared expert” dedicated to capturing common knowledge, thereby freeing task-specific experts from redundant learning. Yet, existing implementations of shared experts in MoE-MTL have their own drawbacks. Often, the shared expert’s contribution is simply added to the sparse experts, implicitly giving it a fixed and dominant influence. This creates an imbalance, where the shared expert’s impact can overshadow the specialized experts. Furthermore, this disproportionate influence can worsen “gradient conflicts” – a common problem in MTL where optimizing for one task might negatively affect another – especially in later training stages.
To address these critical limitations, researchers from Hokkaido University have introduced a novel framework called Adaptive Shared Experts (ASE) within a Low-Rank Adaptation (LoRA) based Mixture-of-Experts. This innovative approach rethinks how shared experts operate, making them more dynamic and balanced. The core idea is to assign the shared experts gating weights that are computed by a router and then normalized jointly with the sparse experts. This means the shared expert’s contribution is no longer fixed but adapts based on the input and the training stage.
This adaptive design offers several advantages. It ensures stable and balanced outputs by normalizing contributions from both shared and sparse experts. Crucially, during the initial phase of adapting from single-task to multi-task learning, the shared experts can take a leading role, providing a strong foundation of common knowledge. As training progresses, their influence gradually reduces, which helps to alleviate the aforementioned gradient conflicts and allows task-specific experts to specialize further. This dynamic adjustment facilitates a smoother and more effective transition to multi-task learning.
Efficiency is also a key consideration. The proposed framework implements each expert as a Low-Rank Adaptation (LoRA) module. LoRA is a technique that significantly reduces the number of trainable parameters and computational overhead while maintaining the model’s expressive power. This makes the MoE-MTL models more practical to train, especially for large-scale applications.
Adding another layer of innovation, the researchers also incorporated the concept of “fine-grained experts.” This involves increasing the total number of LoRA experts while proportionally reducing their individual rank. This strategy allows for a comparable parameter budget but enables even finer specialization among experts and promotes more effective cooperation across different tasks, leading to improved performance and stability.
The effectiveness of the Adaptive Shared Experts framework was rigorously evaluated on the PASCAL-Context benchmark, a dataset commonly used for multi-task learning. Experiments demonstrated that ASE consistently improved performance across various configurations, outperforming both vanilla LoRA-MoE baselines and other established MTL methods. For instance, under a specific configuration (32/6/2/2), ASE achieved the best overall results with 74.0 mIoU on semantic segmentation and 60.3 mIoU on human parts segmentation, with an average performance gain of +7.58%. These significant improvements were achieved with only a marginal increase in parameters, highlighting the efficiency of the LoRA-MoE design.
A fascinating insight from the research involved visualizing the activation frequency of experts. In the early stages of training, the adaptive shared expert dominated most layers, underscoring its role in facilitating the initial knowledge transfer. As training progressed, the sparse experts became more specialized, and the shared expert’s contribution gradually decreased, confirming its adaptive nature in mitigating gradient conflicts. This dynamic behavior ensures a stable balance between shared and specialized knowledge throughout the learning process.
Also Read:
- Smarter Expert Selection: How LD-MoLE Enhances Language Model Fine-Tuning
- Unlocking Large AI Models for Edge Devices Through Collaborative Compression
In conclusion, the Adaptive Shared Experts (ASE) framework, combined with LoRA-based Mixture-of-Experts and fine-grained expert designs, offers a robust and efficient solution for multi-task learning. By dynamically managing the contribution of shared experts and optimizing computational resources, this research paves the way for more effective and specialized AI models capable of handling diverse tasks. You can read the full research paper here.
