PaperlayerView exact revision on arXiv

Source-linked AI summary

M$^3$ViT: Mixture-of-Experts Vision Transformer for Efficient Multi-task Learning with Model-Accelerator Co-design

Hanxue Liang, Zhiwen Fan, Rishov Sarkar, Ziyu Jiang, Tianlong Chen, Kai Zou, Yu Cheng, Cong Hao, Zhangyang Wang

arXiv:2210.14793v1cs.CV

TL;DR

Efficient on-device MTL must address both conflicting task optimization and the cost of activating nearly the entire model for one task. M3ViT combines task-dependent sparse MoE routing in a ViT with hardware co-design, achieving higher accuracy with 88% fewer inference FLOPs and improved FPGA efficiency.

  • Problem

    MTL faces conflicting task gradients during training and inefficient near-full-model activation during single-task inference in resource-constrained, latency-sensitive systems.

  • Method

    M3ViT replaces ViT feed-forward layers with task-dependent sparsely activated MoE experts and uses computation reordering for memory-efficient, zero-overhead task switching.

  • Results

    M3ViT surpasses encoder-focused MTL methods while reducing inference FLOPs by 88%, memory by 2.40×, and achieving up to 9.23× higher energy efficiency than comparable FPGA baselines.

  • Takeaways & Limitations

    Sparse task-specific pathways jointly improve MTL performance and single-task deployment efficiency, while the hardware design scales to any number of experts without task-switching overhead.

Abstract

from arXiv · show

Multi-task learning (MTL) encapsulates multiple learned tasks in a single model and often lets those tasks learn better jointly. However, when deploying MTL onto those real-world systems that are often resource-constrained or latency-sensitive, two prominent challenges arise: (i) during training, simultaneously optimizing all tasks is often difficult due to gradient conflicts across tasks; (ii) at inference, current MTL regimes have to activate nearly the entire model even to just execute a single task. Yet most real systems demand only one or two tasks at each moment, and switch between tasks as needed: therefore such all tasks activated inference is also highly inefficient and non-scalable. In this paper, we present a model-accelerator co-design framework to enable efficient on-device MTL. Our framework, dubbed M$^3$ViT, customizes mixture-of-experts (MoE) layers into a vision transformer (ViT) backbone for MTL, and sparsely activates task-specific experts during training. Then at inference with any task of interest, the same design allows for activating only the task-corresponding sparse expert pathway, instead of the full model. Our new model design is further enhanced by hardware-level innovations, in particular, a novel computation reordering scheme tailored for memory-constrained MTL that achieves zero-overhead switching between tasks and can scale to any number of experts. When executing single-task inference, M$^{3}$ViT achieves higher accuracies than encoder-focused MTL methods, while significantly reducing 88% inference FLOPs. When implemented on a hardware platform of one Xilinx ZCU104 FPGA, our co-design framework reduces the memory requirement by 2.4 times, while achieving energy efficiency up to 9.23 times higher than a comparable FPGA baseline. Code is available at: https://github.com/VITA-Group/M3ViT.

1 Introduction

M3ViT addresses conflicting task optimization and inefficient all-task inference by combining task-dependent sparse MoE routing in ViT with hardware co-design for efficient switching and deployment.

  • Motivation: MTL training can degrade when conflicting task gradients confuse shared weights, slowing convergence and biasing representations against some tasks.Negative cosine similarities between task gradients are reported as detrimental.
  • Method: M3ViT replaces ViT feed-forward layers with sparsely activated MoE experts selected by a task-dependent gating network.The routing is conditioned on tasks and selects a subset of experts for each input token.
  • Results: 88% inference FLOPs are reduced on PASCAL-Context while M3ViT achieves better performance than encoder-focused MTL methods.On NYUD-v2, it reduces 71% FLOPs for single-task execution with comparable results.
  • Results: 2.40× lower memory requirement and up to 9.23× higher energy efficiency are achieved on one Xilinx ZCU104 FPGA.The experiments also report up to 10.79× lower energy than the GPU baseline.
  • Method: M3ViT uses task-specific sparse pathways for single-task inference and targets zero-overhead switching between tasks.The framework adopts an inference setting that activates one task at a time while switching between tasks as needed.

2 Related Works

Related work spans multi-task learning, sparse mixture-of-experts models, vision transformers, and FPGA acceleration of transformer-based models.

  • Multi-task Learning: MTL methods have evolved from modeling common information with distance metrics or probabilistic priors to learning shared deep representations.ViTs also extend the task range toward modalities such as text and audio.
  • Mixture of Experts: MoE models use input-dependent routing across expert sub-models, while sparse MoE activates only a few experts during training and inference.Traditional dense MoEs select all experts and therefore incur intensive computational costs.
  • Vision Transformer: Vision Transformers have been applied to image generation, classification, segmentation, detection, 3D processing, and novel view synthesis.Their use in computer vision builds on transformer successes in natural language processing.
  • Hardware Acceleration: FPGA acceleration research addresses transformer computation and memory costs using compression methods including quantization, pruning, and structured weight representations.The passage characterizes these compression methods as lossy.

3 Method

M3ViT adapts sparse mixture-of-experts layers to a ViT backbone for multi-task learning and co-designs hardware to execute selected experts efficiently. Task-dependent routing supports sparse training and inference, while computation reordering addresses FPGA memory limits and task switching.

  • MoE ViT backbone: M3ViT replaces the ViT feed-forward network with sparsely activated MoE experts and uses routers to select relevant experts for each token.The standard ViT backbone processes image patches as tokens through transformer layers, while MoE replaces its dense MLP component.
  • Sparse expert selection: The router uses top-K gating to retain only the largest K routing values; experiments choose K = 4 from N = 16 experts.Expert sizes are reduced fourfold relative to standard ViT MLPs to keep computation FLOPs equivalent, with balancing loss discouraging repeated selection of the same experts.
  • Task-dependent routing: The multi-gate design assigns each task its own router while sharing expert candidates across tasks.Each task-dependent router receives shared token embeddings and selects its experts independently.
  • Task-dependent routing: The task-conditioned design shares one router by concatenating token embeddings with a task-specific embedding derived from a one-hot task vector.A two-layer MLP maps the task vector to a 64-dimensional embedding before routing.
  • Hardware constraints: Naively storing all expert weights on-chip scales memory as O(N), while cache-based loading incurs delays and cache-unfriendly access patterns.These constraints motivate reorganizing computation around experts rather than processing tokens in sequence.
  • Computation reordering: M3ViT queues tokens for their selected experts, computes each expert over its entire queue, and overlaps parameter loading through double buffering.The hardware reorders MoE computation expert-by-expert and swaps buffers between iterations.
  • Hardware efficiency: The FPGA design uses O(1) on-chip memory with respect to K and N, scales to any number of experts, and introduces zero-overhead task and frame switching.The implementation shares hardware across ViT blocks and hides nearly all latency from off-chip expert-weight accesses.

4 Experiments

Experiments evaluate M3ViT on dense-labeling MTL benchmarks using ViT backbones, compare encoder-focused and task-dependent MoE variants, and measure software and hardware efficiency. M3ViT improves multi-task performance while reducing inference computation, memory, latency, and energy under single-task execution.

  • Experiment setup: Experiments use PASCAL-Context and NYUD-v2 for dense-labeling multi-task learning evaluation.PASCAL-Context covers five tasks, while NYUD-v2 covers semantic segmentation and monocular depth estimation.
  • Experiment setup: The evaluation reports task-specific accuracy metrics, average per-task performance change, latency, energy, and on-chip memory for batch-one single-task inference.Software metrics include mIoU, mErr, rmse, odsF, and Δm; hardware evaluation measures latency, energy, and memory.
  • Comparison with state-of-the-art dense prediction MTL: On PASCAL-Context, M2ViT achieves +2.71% MTL performance and reduces Cross-Stitch single-task inference FLOPs by 88%.The comparison includes encoder-focused methods and approaches targeting training conflicts.
  • Comparison with state-of-the-art dense prediction MTL: On NYUD-v2, M2ViT reduces previous SoTA inference FLOPs by 68% while achieving comparable MTL performance.Adding MoE improves MTL performance from −6.27% to +1.59% without increasing inference FLOPs.
  • Effect of task-dependent MoE design: Task-dependent MoE selection improves performance over a shared router, with multi-gate routing outperforming task-conditioned routing.The comparison covers single-router, multi-router, and task-conditioned variants on both datasets.
  • Hardware performance results: The computation-reordering design reduces on-chip memory from 11.610 MiB to 4.840 MiB, a 2.40× reduction.On PASCAL-Context, it also achieves 9.23× lower latency and energy than the memory-constrained FPGA comparison, and 10.79× better energy efficiency than MoE ViT on GPU.

5 Conclusion, Discussion of Limitation and Broader Impact

M3ViT combines task-specific sparse expert pathways with hardware co-design for efficient on-device multi-task learning. It reports improved accuracy and lower resource and energy use, while remaining mainly evaluated on academic datasets.

  • Conclusion and Contribution: M3ViT customizes MoE layers into a ViT backbone, sparsely activates task-specific experts during training, and uses task-relevant expert pathways during inference.The hardware co-design additionally enables zero-overhead switching between tasks.
  • Conclusion and Contribution: M3ViT surpasses encoder-focused MTL methods, reduces 88% FLOPs, and saves more than 8× energy over the baseline.
  • Limitation: The main limitation is that M3ViT has so far been evaluated primarily on academic datasets rather than real applications such as autonomous driving.
  • Broader Impact: The authors position the framework as reducing resource and energy consumption for MTL while maintaining SOTA performance, supporting Green AI goals.

A Implementation Details

The implementation evaluates MoE ViT encoders across multiple DeiT-style backbones and uses task-conditioned routing with lightweight PUP decoders for multi-task predictions.

  • Encoder: The MTL encoder uses ViT-tiny, ViT-small, and ViT-base variants following DeiT, with MoE expert layers embedded once every two ViT blocks.
  • Task-Conditioned Routing: A single-layer MLP router maps token embeddings to expert-selection probabilities, while a two-layer task embedding network conditions routing on the task.
  • Decoder: The decoder uses progressive upsampling with five convolutional and four upsampling layers, providing a lighter alternative to Deeplab for ViT outputs.
  • Decoder: The final encoder-block output is fed into decoders for multi-task predictions, with intermediate decoder features also used for multi-task feature distillation.

A.2 Training Setup

Training uses ImageNet-pretrained encoder initialization, polynomial learning-rate decay, and SGD-based optimization, with experiments conducted on specified FPGA and GPU platforms.

  • Pre-training: The encoder backbones are pretrained on ImageNet, while decoders are randomly initialized during the MTL pre-training stage.MoE ViT encoders are pretrained following the strategy used for their DeiT counterparts.
  • Training Setup: MTL training uses polynomial learning-rate decay and SGD with initial learning rate 0.002, momentum 0.9, weight decay 0.0001, and batch size 16.
  • Platform Specifications: The target Xilinx ZCU104 FPGA provides 1,728 DSPs, 504K LUTs, 461K registers, 11 Mbit block RAM, and 27 Mbit UltraRAM.
  • Platform Specifications: The NVIDIA Quadro RTX 8000 baseline has 4,608 CUDA cores, 48 GB of GDDR6 memory, a 1,395 MHz clock, and 295 W power consumption.

B.1 Additional Experiments on ViT-tiny and ViT-base

Additional experiments show that M3ViT scales across ViT-tiny and ViT-base, improving accuracy while reducing computation and, with hardware co-design, energy consumption.

  • ViT-base: +4.00% on PASCAL-Context and +8.32% on NYUD-v2 are achieved by MoE ViT-base over the reported prior results.
  • ViT-tiny: MoE ViT-tiny uses less than 1/10 the FLOPs of the previous SoTA method Cross-Stitch.
  • Hardware Efficiency: The MoE ViT-tiny hardware co-design achieves energy consumption an order of magnitude lower than Cross-Stitch.
  • Experimental Comparison: The experiments compare M3ViT across ViT-tiny and ViT-base variants against STL-B, MTL-B, TAPS, and Cross-Stitch.

B.2 Additional Experiments on Different Numbers of Tasks

Across datasets and task counts, M3ViT’s performance advantage generally grows as more tasks are added, while retaining lower computational cost. Its strongest gains appear on settings with larger task counts, though some fine-detail tasks remain exceptions.

  • Different task counts: M3ViT consistently outperforms the MTL-B ViT-small baseline as the number of tasks increases across NYUD-v2 and PASCAL-Context.The comparison includes STL-B, MTL-B, Cross-Stitch, M-ViT, and M3ViT.
  • Different task counts: −0.91% versus −3.26% on four-task NYUD-v2 shows M3ViT surpassing Cross-Stitch as task count increases.M3ViT is slightly lower than Cross-Stitch on the two-task NYUD-v2 setting but performs better in other reported settings.
  • Different task counts: M3ViT achieves much better PASCAL-Context performance than Cross-Stitch, except on normal estimation and saliency detection.The authors speculate that those tasks benefit from smaller receptive fields and Cross-Stitch’s local information fusion.
  • Taskonomy evaluation: M3ViT’s relative superiority over MTL-ViT becomes stronger as Taskonomy task counts rise from three to nine.The evaluation reports relative performance improvement using the same preprocessing and evaluation method as the baseline.

B.3 Comparisons with Decoder-focused Methods

Decoder-focused architectures use predictions or intermediate features from all tasks during training and inference, conflicting with sparse task-specific activation for efficient MTL inference. Their high FLOP requirements also hinder deployment on resource- and latency-constrained edge devices.

  • Decoder-focused architectures typically require predictions or intermediate features from all tasks during training and inference.
  • Activating all tasks during inference conflicts with sparse network activation for efficient MTL inference.
  • High FLOP consumption makes decoder-focused models difficult to deploy on resource- and latency-constrained edge devices.

C Latency Breakdown of Our Model

The FPGA implementation measures M3ViT latency on NYUD-v2 across patch embedding, ViT layers, and MoE layers. Expert computation takes nearly the same time as ViT fully connected layers, indicating memory-efficient computation reordering with near-zero latency impact.

  • 84.538 ms is the total FPGA inference latency of ViT-small M3ViT on NYUD-v2.The latency is divided among patch embedding, ViT layers, and MoE layers.
  • 18.567 ms for all MoE experts nearly matches 18.447 ms for ViT fully connected layers.This comparison is reported for the NYUD-v2 latency breakdown.
  • The hardware computation reordering mechanism maintains memory efficiency with near-zero impact on latency.The conclusion follows from the near-equal expert and fully connected computation times.
  • Figure 6 partitions latency into patch embedding, six standard ViT layers, and six MoE layers.ViT and MoE layers are further divided into shared self-attention and their respective MLP or expert computations.
Loading 2210.14793v1…