Source-linked AI summary
Uni-Perceiver-MoE: Learning Sparse Generalist Models with Conditional MoEs
Jinguo Zhu, Xizhou Zhu, Wenhai Wang, Xiaohua Wang, Hongsheng Li, Xiaogang Wang, Jifeng Dai
TL;DR
Generalist models can lose performance on some tasks compared with specialized models, and the paper identifies task and modality interference as a key factor. It introduces Conditional MoEs for Uni-Perceiver and reports competitive downstream performance with limited prompt-tuning data while retaining generalization to new tasks.
Problem
Generalist models can degrade on some tasks compared with specialized models, and task and modality interference is identified as a main factor.
Source: §1 · pp. 1–2Method
Conditional MoEs vary expert routing by token, context, modality, task, or predefined token attributes, and are incorporated into Uni-Perceiver.
Source: §1 · p. 2Results
With prompt tuning on 1% downstream data, the sparse generalist model achieves competitive performance with previous SOTAs using <5% training data and <10% training cost.
Source: §5 · p. 9Takeaways & Limitations
Conditional MoEs allow joint pretraining without performance degradation while retaining generalization to novel tasks.
Source: §5 · p. 9
Abstract
from arXiv · showhide
To build an artificial neural network like the biological intelligence system, recent works have unified numerous tasks into a generalist model, which can process various tasks with shared parameters and do not have any task-specific modules. While generalist models achieve promising results on various benchmarks, they have performance degradation on some tasks compared with task-specialized models. In this work, we find that interference among different tasks and modalities is the main factor to this phenomenon. To mitigate such interference, we introduce the Conditional Mixture-of-Experts (Conditional MoEs) to generalist models. Routing strategies under different levels of conditions are proposed to take both the training/inference cost and generalization ability into account. By incorporating the proposed Conditional MoEs, the recently proposed generalist model Uni-Perceiver can effectively mitigate the interference across tasks and modalities, and achieves state-of-the-art results on a series of downstream tasks via prompt tuning on 1% of downstream data. Moreover, the introduction of Conditional MoEs still holds the generalization ability of generalist models to conduct zero-shot inference on new tasks, e.g., video-text retrieval and video caption. Code and pre-trained generalist models shall be released.
1 Introduction
Generalist models share parameters across tasks and can generalize to novel tasks, but task interference can make their performance suboptimal. The paper proposes Conditional MoEs and applies them to Uni-Perceiver to mitigate interference while retaining generalization.
1 Introduction
- Earlier specialized models require task-specific components to be redesigned and fine-tuned on sufficient downstream data for new tasks, so model size grows with task diversity.The paper contrasts this with the goal of generalist models.Source: §1 · p. 1
- Inconsistent or opposing task-gradient directions make the optimal update for shared parameters uncertain, leading to suboptimal performance.The authors frame inconsistent multi-task optimization as the main source of task interference.Source: §1 · p. 2
- Vanilla MoE routing uses token representations, which the authors say can raise training and inference costs while neglecting task and modality information.The paper argues that routing strategies need special design for generalist models.Source: §1 · p. 2
- Conditional MoEs condition routing on token, context, modality, task, or predefined token attributes to mitigate interference.The predefined-attribute variant uses 8-bit attribute embeddings to describe the current task and modalities, with reported computational and memory efficiency and good generalization.Source: §1 · p. 2
- The authors incorporate Conditional MoEs into Uni-Perceiver by replacing linear projections in its self-attention and FFN blocks with conditional MoE layers.They report joint pretraining without performance degradation while maintaining generalization to other tasks.Source: §1 · p. 2
- With 1% downstream-data prompt tuning, the generalist model achieves competitive performance with previous SOTAs using <5% training data and <10% training cost.The work also reports that joint pretraining with Conditional MoEs avoids performance degradation while maintaining generalization to other tasks.Source: §1 · p. 2
2 Related Works
Related work contrasts task-specific designs with generalist models that share parameters across tasks, and surveys multi-task learning and MoE approaches relevant to interference. These approaches motivate the paper’s focus on conditional expert routing.
Specialized Models
- Specialized models use task-specific architectures or modules across vision, language, and vision-language tasks, but this task-specific paradigm conflicts with the generalist objective.Integrated specialized models combine several specialized models while retaining task-specific modules.Source: §2 · p. 2
Vanilla Generalist Models
- Figure 1 contrasts fully shared, task-specific, and efficient MoE encoder blocks as alternative parameterizations for multi-task models.The caption identifies the three architectures but does not report a performance comparison.Source: Figure 1 · p. 3
- Generalist models unify tasks with shared parameters and can perform zero-shot inference on new tasks, but potential interference across modalities and tasks is rarely investigated.Examples include Uni-Perceiver [93], OFA [79], Flamingo [3], SimVLM [84], UniCORN [86], and Gato.Source: §2 · pp. 2–3
Multi-Task Learning
- Multi-task learning can enable collaboration between tasks but may also introduce interference; existing responses include adjusting task-loss weights or using task-dedicated parameters.Task-specific parameters are described as difficult to generalize to new tasks and inconsistent with generalist models.Source: §2 · p. 3
Mixture of Experts (MoEs)
- Prior MoE work spans language models, Transformers, CNNs, and vision Transformers, while Task-MoE addresses interference among multilingual tasks.The paper positions its work as exploring MoEs under different conditions for general models.Source: §2 · p. 3
3 Methodology
The methodology analyzes task interference through gradient directions, then introduces Conditional MoEs, which route tokens to sparse expert subsets using conditions at different levels. Its variants trade off routing information, computational and memory costs, and generalization to new tasks.
3.1 Task Interference
- The interference metric estimates task i’s loss change under task j’s gradient direction, with sampled batches and measurements at the 4-th and 12-nd FFN blocks.The reported analysis finds positive effects from image captioning on image classification and masked language modeling in shallow layers, but little enhancement—and sometimes opposing gradient directions—between different objectives in deep layers.Source: §3.1 · pp. 3–4, Table 1 · p. 4
- Shared parameters can interfere across tasks, while task-specific parameters address interference but make generalization to unseen tasks difficult; MoEs sparsely activate parameter subsets.The paper focuses on conditioning expert gates on information from different levels.Source: §3.1 · p. 4
3.2 Conditional Mixture-of-Experts (Conditional MoEs)
- A gate assigns each token to a small expert subset, whose weighted computations form the output; experts with zero gate weight need not be computed.The gate uses a routing strategy that can vary with the selected condition.Source: §3.2 · p. 4
- Context-level routing combines a token representation with attention-pooled global context to guide expert decisions when similar tokens occur in conflicting tasks.The pooled context is formed from the tokens in the current sample.Source: §3.2 · p. 5
- Modality-level routing uses a token’s modality embedding, whereas task-level routing uses its task embedding and sends all tokens from that task to the same expert set.Source: §3.2 · p. 5
- Attribute routing uses 8-dimensional binary token attributes to describe task and modality information without relying on task-specific IDs that limit generalization.The attributes include modality, model causation type, and token source.Source: §3.2 · p. 5, Table 2 · p. 6
- Conditional MoEs replace linear projections in Uni-Perceiver’s self-attention and FFN blocks with conditionally routed expert layers.The method is intended to mitigate interference while retaining the generality of the pretrained model.Source: §3 · p. 3, §3.2 · p. 6
3.3 Comparison of Conditional-MoE Variants
- All variants reduce computation through sparse routing, but data-dependent token and context routing have higher memory costs; data-independent variants can be merged into dense projections and can retain new-task generalization.Task-level routing is tied to task IDs, while attribute routing is designed to transfer to new tasks without task-specific modifications.Source: §3.3 · p. 6
4 Experiments
The experiments test Conditional MoEs on routing strategies, pretraining tasks, and novel-task generalization. Across evaluations, they report improved performance over vanilla Uni-Perceiver while retaining generalization to tasks absent from pretraining.
4.1 Datasets
- The evaluation uses ImageNet-1k, Kinetics-400, Flickr30k, and COCO Caption for pretraining-task assessment, plus MSVD and GLUE for novel-task generalization.The passage also describes pretraining datasets, including ImageNet-21k, Kinetics-700, Moments in Time, Books&Wiki, and image-text-pair datasets.Source: §4.1 · p. 7
4.2 Implementation Details
- Experiments use Tiny, Base, and Large Uni-Perceiver variants; Base and Large add Conditional-MoE layers every other layer, while Tiny uses them in all layers.A top-2 gate is used unless otherwise specified.Source: §4.2 · p. 7
4.3 Ablation Studies
- Task- and attribute-level routing mitigate interference better than token-, context-, and modality-level variants; task-level routing is harder to generalize because it uses specialized task IDs.Token- and context-level MoEs have relatively higher training and inference costs; the other three use reparameterization techniques for efficiency, and attribute routing is selected.Source: §4.3 · p. 7, Table 3 · p. 7
4.4 Evaluation on Pre-training tasks
- Across four pretraining tasks, Conditional-MoE Uni-Perceiver outperforms vanilla by a large margin; 1% prompt tuning approaches counterparts using >50× data and >10× training cost.The evaluation compares image classification, video classification, and image-text retrieval in Table 4, and also includes image captioning in Table 5a.Source: §4.4 · p. 8, Table 4 · p. 8, Table 5 · p. 9
- With 100% downstream data, the model matches or exceeds SOTAs trained with similar data sizes across the evaluated tasks.The authors interpret this as evidence that Conditional-MoE generalist models learn high-quality representations.Source: §4.4 · p. 8, Table 4 · p. 8, Table 5 · p. 9
4.5 Generalization to Novel Tasks
- On novel video tasks, Uni-Perceiver with Conditional MoEs achieves reasonable zero-shot performance, beats vanilla by a great margin, and reaches comparable results to similar-cost SOTAs after 1% prompt tuning.Conditional MoEs also boost performance on GLUE benchmarks.Source: §4.5 · p. 9, Table 5 · p. 9
5 Conclusion
Conditional MoEs sparsely activate subnetworks without task-specific designs, addressing interference while preserving generalization to novel tasks. Uni-Perceiver-MoE reports competitive performance with prior state of the art using limited downstream data and training cost.
5 Conclusion
- Using prompt tuning on 1% of downstream data, the model achieves competitive performance with previous SOTAs using <5% training data and <10% training cost.The paper reports these as comparisons with previous state-of-the-art methods.Source: §5 · p. 9
- Conditional MoEs sparsely activate subnetworks without task-specific designs, enabling joint pre-training on multiple tasks without performance degradation.The conclusion also states that generalization to novel tasks is maintained.Source: §5 · p. 9
- The model retains generalization to novel tasks while using Conditional MoEs.The paper cites zero-shot inference on new tasks in the abstract, but this conclusion passage does not specify examples.Source: §5 · p. 9
6 Appendix
The appendix documents pre-training and downstream evaluation settings, examines Conditional-MoE placement and expert routing, compares training data scale, and lists dataset licenses.
6.1 Experimental Details
- Pre-training used Uni-Perceiver-Ti, B, and L variants for 400k iterations on 32 NVIDIA-A100 GPUs, with staged image resolutions and load-balancing settings for data-dependent MoEs.Training began at 160×160 resolution and continued for 10% of iterations at 224×224; the balance-loss coefficient was 10^-2.Source: §6.1 · p. 15
- Evaluation includes untuned, full-data fine-tuning, and 1%-data prompt tuning; Tables 8–9 cover reimplementation metrics and tuning hyperparameters, and Kinetics-400 validation overlaps were removed from Kinetics-700 training.Prompt tuning adds learnable prompts at each encoder layer; class labels and linear heads are added for classification tasks.Source: §6.1 · p. 16, Table 8 · p. 16, Table 9 · p. 16
6.2 The Placement of Conditional MoEs
- Replacing more layers with attribute-level MoEs mitigates task interference more, and using them in both self-attention and FFN layers works better than self-attention alone.Table 11 compares FFN and self-attention placement and every-n-layer settings.Source: §6.2 · p. 16, Table 11 · p. 19
6.3 Visualization of Gating decisions for Attribute-level MoEs
- Top-1 expert activations are shown across self-attention and FFN layers; routing shares experts across some same-modality inputs and targets, and uses causal-mask attributes.The passage identifies possible interference between auto-encoding and auto-regressive tasks from causal-mask routing.Source: §6.3 · p. 16, Figure 3 · p. 17
6.4 Training Statistics of Relevant Methods
- Uni-Perceiver-MoE's training data scale is smaller than most generalists and comparable to some specialized models; accelerator training time is converted to TPU-v3-core-days using TFLOPS.Table 10 categorizes methods as specialized, integrated specialized, or generalist models.Source: §6.4 · pp. 16, 18, Table 10 · p. 18
6.5 Licenses of Datasets
- ImageNet-21K follows its terms of use; COCO and SBU Caption images follow Flickr terms, while Kinetics datasets use the Creative Commons Attribution 4.0 International License.Source: §6.5 · p. 19
- BooksCorpus is licensed under GNU GPL version 3; most Wikipedia text is co-licensed under CC BY-SA 3.0 and the GNU Free Documentation License, with stated exceptions.Source: §6.5 · p. 19
- YFCC media use Creative Commons licenses; CC12M and CC3M follow their respective Terms of Use, and Visual Genome uses a Creative Commons Attribution 4.0 license.Source: §6.5 · p. 19