UniDense: Unleashing Diffusion Models with Meta-Routers for Universal Few-Shot Dense Prediction

Venue: acmmm_org_ACMMM_2024_Conference
Published: 2024-03-18 03:32:43

Abstract: Universal few-shot dense prediction requires a versatile model capable of learning any dense prediction task from limited labeled images, which necessitates the model to possess efficient adaptation abilities. Prevailing few-shot learning methods rely on efficient fine-tuning of model weights for few-shot adaptation, which carries the risk of disrupting the pre-trained knowledge and lacks the capability to extract task-specific knowledge contained in the pre-trained model. To overcome these limitations, our paper approaches universal few-shot dense prediction from a novel perspective. Unlike conventional fine-tuning techniques that directly use all parameters of the model and modify a specific set of weights for few-shot adaptation, our method focuses on selecting the task-relevant computation pathways of the pre-trained model while keeping the model weights frozen. Building upon this idea, we introduce a novel framework UniDense for universal few-shot dense prediction. First, we construct a versatile MoE architecture for dense prediction based on the Stable Diffusion model. We then utilize episodes-based meta-learning to train a set of routers for this MoE model, called Meta-Routers, which act as hyper-networks responsible for selecting computation blocks relevant to each task. We demonstrate that fine-tuning these meta-routers for novel tasks enables efficient adaptation of the entire model. Moreover, for each few-shot task, we leverage support samples to extract a task embedding, which serves as a conditioning factor for meta-routers. This strategy allows meta-routers to dynamically adapt themselves for different few-shot task, leading to improved adaptation performance. Experiments on a challenging variant of Taskonomy dataset with 10 dense prediction tasks demonstrate the superiority of our approach.

Main paper: 
 related work: 
 Universal few-shot dense prediction. Existing few-shot learning methods in computer vision are typically designed for specific tasks like classification and semantic segmentation [17, 23, 39] . However, these methods often rely on task-specific prior knowledge and assumptions in their model architecture and training procedures, making them less suitable for generalizing to arbitrary dense prediction tasks. VTM [18] is the pioneering work that addresses the challenge of few-shot learning for arbitrary dense prediction tasks in a universal manner, which utilizes non-parametric matching on tokenized image and label embeddings. However, VTM's computational demands and data inefficiency arise from the matching operation between all tokens of the query image and the support images. Furthermore, the bias tuning technique used for few-shot adaptation in VTM lacks the ability to excavate the task-specific knowledge embedded in the pre-trained model. In contrast, our UniDense framework offers rapid inference regardless of the number of support images, and the router fine-tuning technique allows our method to extract task-specific knowledge from the pre-trained model by selecting task-relevant computation pathways. Few-shot fine-tuning. Conventional fine-tuning strategies for few-shot learning achieve efficient adaptation by selectively finetuning specific parts of the model [18, 19, 35, 38] . Another set of methods involve inserting adapter modules into the frozen pretrained model and fine-tuning these adapters for few-shot adaptation [12, 46] . However, all these methods directly use all parameters of the model and overlook the extraction of task-specific knowledge hidden in the well-learned pre-trained model. In contrast, our approach achieves few-shot adaptation by selectively choosing computation pathways that excavate the task-specific knowledge. The closest relevant studies to our method involve training a masking module to select relevant features for novel few-shot tasks [4, 49] . However, these methods are not tailored for the model with multitask pre-training, which limits their ability for universal few-shot dense prediction. In contrast, our MoE design allows us to fully exploit the knowledge from multiple base training tasks. Unleashing diffusion models for visual perception. While diffusion models are primarily trained using generative loss, their features have shown impressive performance in specific visual perception tasks that demand a comprehensive understanding of pixel-level fine-grained information [1, 41, 48] , such as semantic segmentation and depth estimation. However, existing works either focus on leveraging diffusion models for a specific task [1, 41] or address multiple visual perception tasks without considering the few-shot setting [48] . Currently, there is no effective and elegant method to fully harness the pre-trained knowledge of diffusion models and adapt them to few-shot visual perception tasks in a universal manner. To the best of our knowledge, we are the first to make an attempt to bridge this gap. 
problem definition: 
 The problem of universal few-shot dense prediction [18] aims to build a universal few-shot learner F that, for any dense prediction task T : R 𝐻 ×𝑊 ×3 → R 𝐻 ×𝑊 ×𝐶 T , where 𝐶 T ∈ N, can produce a prediction Ŷ𝑄 for an unseen (query) image 𝑋 𝑄 given only a few labeled (support) examples S T :EQUATIONwhere 𝑁 𝑆 denotes the number of labeled images for each few-shot dense prediction task, namely the shot number.To capture the general knowledge for the domain of dense prediction, several base dense prediction tasks {T 𝑖 𝐵 } 𝑖 ≤𝑁 𝐵 are used for the multi-task pre-training of the few-shot learner F , where each base task does not have a limitation on the number of examples. Here, 𝑁 𝐵 represents the number of base tasks and the testing task is not included among the base tasks. 
methodology: 
 In this section, we present the details of our proposed UniDense framework. We begin by introducing the foundational concepts of MoE (Mixture of Experts) and Stable Diffusion model, upon which UniDense is built. Next, we delve into the architectural design details of our method. Subsequently, we proceed to introduce the three-stage pipeline of our UniDense framework.Mixture of Experts. A MoE layer typically contains a set of expert networks 𝐸 1 , 𝐸 2 , ..., 𝐸 𝑁 𝐸 along with a routing network (router) 𝐺. The output of a MoE layer is the weighted sum of the output from every expert. The router 𝐺 calculates the weight 𝐺 𝑘 (𝑥) for each expert given the input 𝑥. Formally, the output of a MoE layer is𝑦 = 𝑁 𝐸 ∑︁ 𝑘=1 𝐺 𝑘 (𝑥)𝐸 𝑘 (𝑥).(EQUATIONThe router 𝐺 is a noisy top-K routing network [34] with parameters 𝑊 𝐺 and 𝑊 𝑁 . It models 𝑃 (𝐸 𝑘 |𝑥) as the probability of using expert 𝐸 𝑘 and selects 𝑁 𝐾 experts with the highest 𝑃 (𝐸 𝑘 |𝑥) to contribute to the final output. The whole process is shown as follows:EQUATIONwhere 𝑇𝑜𝑝 (•, 𝑁 𝐾 ) sets all elements in the vector to zero except the elements with the largest 𝑁 𝐾 values. N (0, 1) represents sampling from a normal distribution. Softplus is the smooth approximation to the ReLU function:EQUATIONStable Diffusion. The Stable Diffusion model [29] is a text-toimage model trained on the large-scale image-text dataset LAION-5B [31] . This model has demonstrated impressive performance in generating images controlled by language descriptions. Initially, they train a VQGAN model [7] , which comprises an encoder E and a decoder D, enabling conversion between the pixel space and the latent space. Subsequently, a diffusion model is trained on this latent space. The noise predictor network 𝜖 𝜃 employed in the latent diffusion model is implemented as a UNet [30] and incorporates rich pre-trained knowledge, which has empirically shown its effectiveness in dense prediction tasks [1, 41, 48] . Hence, we opt to adopt the UNet architecture as the foundation for our model's backbone. Our approach does not utilize the diffusion process and language interface of the Stable Diffusion model. Thus, we leave the detail descriptions of the Stable Diffusion model in the appendix.Our model architecture consists of a MoE backbone based on the UNet-like denoising autoencoder from the Stable Diffusion model, along with a prediction head that produces predictions with varying dimensions for different tasks, as shown in Figure 2 . Backbone. To leverage the pre-trained knowledge of the Stable Diffusion model, our backbone is constructed based on the UNetlike denoising autoencoder 𝜖 𝜃 from Stable Diffusion. Specifically, given an input image 𝑋 ∈ R 𝐻 ×𝑊 ×3 , we firstly employ the encoder E of the VQGAN to encode the input image into the latent space:EQUATIONWe feed the latent feature map 𝑍 to 𝜖 𝜃 and then a set of hierarchical feature maps are extracted from the last layer of each output block in different resolutions:EQUATIONwhere𝐹 𝑖 ∈ R 𝐻 2 𝑖+2 × 𝑊 2 𝑖+2 ×𝐶 𝑖, with 𝑖 = 1, 2, 3, 4. For instance, if the input image size is 256×256, the output feature maps of 𝜖 𝜃 will have side lengths of 32, 16, 8, 4. Note that we only use the self-attention layers in 𝜖 𝜃 since we do not utilize the language interface of the Stable Diffusion model. Additionally, no extra noise is added to the input image, which is equivalent to using a timestep of 0 in the diffusion process.To enable efficient adaptation by selecting task-specific computation pathways for novel few-shot tasks, we propose to fine-tune the routers of a MoE backbone. Here, we show how to construct the MoE backbone based on the Stable Diffusion model. Inspired by Mod-Squad [3] , which transforms the vision transformer into a MoE model, we convert the multi-head attention layers and the feed-forward layers in the transformer blocks of the Stable Diffusion model's denoising autoencoder 𝜖 𝜃 into MoE attention layers and MoE MLP layers.We follow MoA [47] to construct the MoE attention layers. Given a query token 𝑞 𝑖 ∈ R 1×𝐶 𝑀 from the query sequence 𝑄, the objective of a MoE attention layer is to generate a new token 𝑦 𝑖 for 𝑞 𝑖 , with𝑁 𝑎 𝐸 attention experts {𝐸 𝑘 𝐴 } 𝑁 𝑎 𝐸𝑘=1 and the router 𝐺 𝐴 :EQUATIONEQUATION𝛼 𝑖,𝑘 = 𝑆𝑜 𝑓 𝑡𝑚𝑎𝑥 (EQUATIONwhere 𝐶 𝐻 is the head dimension. 𝑊 𝐾 ,𝑊 𝑉 ∈ R 𝐶 𝑀 ×𝐶 𝐻 are shared across attention experts to reduce computational complexity, while 𝑊 𝑘 𝑄 ∈ R 𝐶 𝑀 ×𝐶 𝐻 and 𝑊 𝑘 𝑂 ∈ R 𝐶 𝐻 ×𝐶 𝑀 are specific to each expert. As for a MoE MLP layer, we convert the GEGLU feed-forward layer [33] adopted in 𝜖 𝜃 into a MoE model: EQUATIONMoE Layer 𝑙 • • • Experts Meta-Router MoE Layer 𝑙 • • • Experts Meta-Router • • • • • • K-shot K-shot • • • • • •EQUATIONwhere Φ represents the GELU activation function [13] . Φ(𝑥 ŴG ) act as gating values. ⊙ represents element-wise product.The weights of attention experts {𝐸 𝑘 𝐴 } and MLP experts {𝐸 𝑘 𝐹 } are initialized using the weights from the multi-head attention layers and the feed-forward layers in 𝜖 𝜃 . For the specific details of this initialization process, please refer to the appendix. Prediction head. Recent research [1] has shown that different levels of feature maps from 𝜖 𝜃 contain varying levels of semantic information. To aggregate these levels of information, we introduce a feature aggregation module U, which consists of several convolutional layers and upsampling layers. From the lowest scale to the highest scale of {𝐹 𝑖 } 𝑖 ≤4 , the feature aggregation module upsamples the current level of feature map and concatenates it with the next level of feature map, followed with convolutions. This procedure generates a feature map with the shape R 𝐻 8 × 𝑊 8 ×𝐶 U . To generate the final predictions with a high resolution of 𝐻 ×𝑊 , an upsampling procedure is required. Instead of using the naive bilinear interpolation, we utilize the decoder D from the Stable Diffusion model as a parametric upsampling procedure. Since the output dimension 𝐶 T may vary for different tasks, we replace the last convolution layer in D, which originally produces an output with the shape R 𝐻 ×𝑊 ×3 , with task-specific convolution layers. Therefore, the prediction head can be represented as:EQUATIONwhere V denotes a 1 × 1 convolution layer that transforms the feature dimension fromR 𝐻 8 × 𝑊 8 ×𝐶 U to R 𝐻 8 × 𝑊 8 ×4. D T denotes the Stable Diffusion decoder with the task-specific convolution layer, which produces the final prediction Ŷ with the shape R 𝐻 ×𝑊 ×𝐶 T .The overall pipeline of our method is illustrated in Figure 2 , which can be divided into three stages: multi-task pre-training of the backbone, meta-training for the meta-routers, and fine-tuning of meta-routers for efficient adaptation to novel few-shot tasks. Multi-task pre-training. To capture the general knowledge for the domain of dense prediction, we perform multi-task pre-training for the backbone on a set of base dense prediction tasks {T 𝑖 𝐵 } 𝑖 ≤𝑁 𝐵 . There are two options to achieve multi-task pre-training for a MoE model. The first is to train a specific router for each task, while the other is to train a single router for all tasks conditioned on learnable task embeddings. Experiments have demonstrated that the former is more effective than the latter [8] . Therefore, we assign each dense prediction task a task-specific router for multi-task pre-training. Note that the Stable Diffusion encoder E and decoder D used in our model are fixed during the multi-task pre-training to preserve the pre-trained knowledge from the Stable Diffusion model.In the single-task setting, previous MoE methods [9, 28, 34] often use a load-balancing loss to promote similar utilization frequencies of experts across different batches. However, in the multi-task setting, this approach may force experts on conflicting tasks with learning gradients counteracting each other [3] . Therefore, we follow Mod-Squad [3] to utilize a mutual information loss to encourage a sparse but strong dependence between experts and tasks:EQUATIONwhere 𝑇 = {T 𝑖 𝐵 } 𝑖 ≤𝑁 𝐵 , 𝐸 𝑙 represents all experts in the MoE layer 𝑙, and 𝐼 (𝑇 ; 𝐸 𝑙 ) is the mutual information measurement between tasks 𝑇 and experts in 𝐸 𝑙 . An auto-balancing weight 𝑤 T 𝑖 𝐵 is learned for each task-specific loss L T 𝑖 𝐵 . We use cross-entropy loss for semantic segmentation task and L1 loss for the others in our experiments. The hyper-parameter 𝑤 𝑀𝐼 is set as 0.001. Here, we omit the computation details for clarity, which can be found in [3] . Meta-training for Meta-Routers. To enable efficient adaptation by selecting task-specific computation pathways, we propose to fine-tune the routers of the MoE backbone. However, it raises questions about which router to fine-tune and how to ensure its effectiveness. A naive approach is to randomly initialize a set of routers. However, due to the limited number of samples in the few-shot testing stage, such randomly initialized routers cannot ensure the adaptation efficiency for diverse novel tasks. To overcome this, we introduce Meta-Routers, which are optimized to efficiently adapt to any novel few-shot task. To obtain meta-routers, we freeze the well pre-trained backbone and employ episodes-based meta-learning to train an additional set of routers {𝐺 𝑙 𝑀 : ∀𝑀𝑜𝐸 𝑙𝑎𝑦𝑒𝑟 𝑙 }. In each episode, we simulate the few-shot setting during the testing phase. First, we randomly select a base task T and sample 𝑁 𝑆 examples as the support set S T . Then, we sample another 𝑁 𝑄 examples as the query set Q T . The objective of the episodes-based meta-learning is to make the model obtain optimal performance on the query set Q T after the few-shot fine-tuning based on the support set S T :EQUATIONwhere F (𝑋 Q ; S T ) represents the model's prediction for the query image 𝑋 Q after fine-tuning 𝑁 𝐹𝑇 steps on the support set S T , and L T is the task-specific loss. Note that the fine-tuning procedure of F (𝑋 Q ; S T ) only involves the meta-routers and the task-specific convolution layer in the prediction head. Calculating the gradient of L 𝑀 is a bi-level optimization problem, which involves computing the Jacobian matrix of the fine-tuning procedure F (𝑋 Q ; S T ), leading to extra computation cost [11] . In this regard, we opt for FOMAML (First-Order Model-Agnostic Meta-Learning) [11, 25] to strike a balance between speed and effectiveness, which can achieve nearly the same performance with MAML [11] while significantly reducing GPU memory requirements and training time, particularly when the number of fine-tuning steps is large [26] .Moreover, we enhance the meta-routers' capabilities by incorporating the task embedding extracted from the support samples. This transforms each meta-router 𝐺 𝑙 𝑀 into a task-conditioned router: Ĝ𝑙 𝑀 (𝑥, S T ) = 𝐺 𝑙 𝑀 (𝑥 + 𝜉 (S T )),where 𝜉 represents the task embedding extractor. For each sample (𝑋 𝑖 , 𝑌 𝑖 ) in the support set S T from any task T , 𝜉 first transforms 𝑌 𝑖 ∈ R 𝐻 ×𝑊 ×𝐶 T into 𝑌 𝑖 ∈ R 𝐻 ×𝑊 ×3 using linear projections. The specific details of this transformation are provided in the appendix. Then, 𝑋 𝑖 and 𝑌 𝑖 are concatenated, resulting in an input with shape R 𝐻 ×2𝑊 ×3 . Next, this concatenated input is fed into a vision transformer [5] consisting of three stacked transformer blocks. An extra token is appended to the sequence of embedded patches, whose state at the output of the transformer serves as the task token. For each sample in the support set S T , a task token is extracted. Finally, 𝜉 averages all these task tokens to output the final task embedding. The inclusion of task-conditioning allows the meta-routers to dynamically adapt themselves using task embeddings, thereby increasing their flexibility. This enhancement leads to improved adaptation performance compared to the naive meta-routers, which rely solely on fine-tuning a fixed set of initialization weights. Note that our method differs from existing task-conditioned routers [8, 50] in a fundamental way. These existing methods extract the task embedding in a dictionary-like manner: 𝑒𝑚𝑏𝑒𝑑 (𝑖𝑑 𝑡𝑎𝑠𝑘 (𝑥)), where 𝑖𝑑 𝑡𝑎𝑠𝑘 (𝑥) is the task index of current token 𝑥 and 𝑒𝑚𝑏𝑒𝑑 (•) represents the embedding layer. However, in the few-shot learning setting, the pre-trained embedding layer 𝑒𝑚𝑏𝑒𝑑 (•) is not adaptable to novel tasks. In contrast, our approach directly extract the task embedding from the support samples themselves, making it well-suited for few-shot scenarios.During meta-training, we only optimize the parameters of the meta-routers, the task embedding extractor and the task-specific convolution layers in the prediction head. Note that the task embedding extractor remains fixed during the fine-tuning procedure F (𝑋 Q ; S T ) and the testing phase. This ensures a consistent task embedding space, enabling meta-routers to achieve more efficient adaptation when prompted with a stable task embedding. Router fine-tuning. In the testing phase, given an arbitrary unseen dense prediction task T with 𝑁 𝑆 support examples, the model needs to rapidly adapt to this new task while avoid overfitting. Thanks to the versatile backbone acquired through multi-task pre-training and the meta-routers trained with episodes-based meta-learning, adapting the model to unseen few-shot tasks are super easy. You just need to fine-tune the meta-routers and the task-specific convolution layer on 𝑁 𝑆 support samples of the novel task T for 𝑁 𝐹𝑇 steps, which is consistent with the setup in the meta-training stage. 
experiments 5.1 experimental setup: 
 Datasets. We evaluate our approach on a variant of the Taskonomy dataset [45] , which only uses the Taskonomy-tiny partition. Taskonomy dataset contains indoor images with various annotations, where we choose ten dense prediction tasks of diverse output dimensions: semantic segmentation (SS), euclidean distance (ED), Z-buffer depth (ZD), texture edge (TE), occlusion edge (OE), 2DTable 1 : Quantitative comparison on a variant of Taskonomy dataset. 10-shot results are evaluated on each fold after being trained on the tasks from the other folds, while fully-supervised methods are trained on tasks from each fold independently (DPT) or trained on all folds (InvPT). The bolded results are the best results within their respective supervision settings. keypoints (K2), 3D keypoints (K3), reshading (RS), principal curvature (PC), surface normal (SN). Following VTM, we choose some single-channel tasks (ED, TE, OE) and transform them into multichannel tasks to increase task diversity, and standardize all labels to [0, 1]. All images and labels are resized to 256 × 256 resolution. Additional details are given in the appendix. Evaluation protocol. Following the standard evaluation setup of universal few-shot dense prediction [18] , we partition the ten tasks to construct a 5-fold split, in each of which two tasks are used for few-shot evaluation and the remaining eight are used for training. The split is shown in Table 1 . For evaluation on semantic segmentation (SS), we follow the standard binary segmentation protocol in few-shot semantic segmentation [32] and report the mean intersection over union (mIoU) for all classes. We use the mean angle error (mErr) for surface normal prediction (SN) [6] , and use root mean square error (RMSE) for other tasks. Note that all metrics are better when their values are smaller except for the metric of SS. To be consistent with VTM [18] , we test our method on the 'muleshoe' building of Taskonomy dataset. Implementation details. We use the released 1-5 version of Stable Diffusion in our experiments. We introduce MoE attention layers and MoE MLP layers into 𝜖 𝜃 every two layers. For MoE attention layers, we use 24 experts with top-k as 8. For MoE MLP layers, we use 12 experts with top-k as 4. For the multi-task pre-training, we set the learning rate to 1e-4 with a batch size of 16. For each training batch, we randomly sample 16 images with their labels for all training tasks. We use 5 warmup epochs with 50 total training epochs. For the meta-training of meta-routers, the learning rate is set as 5e-4, with 3000 episodes in total. Recent research [21] has found that utilizing a small learning rate exclusively for the nonfinetuned portion of the model can result in significant performance enhancement during the few-shot fine-tuning process. Therefore, we employ a 1/100 smaller learning rate to the rest part of the model for model fine-tuning in the testing phase. The fine-tuning steps 𝑁 𝐹𝑇 is set as 30. Further details can be found in the appendix.TasksIn Table 1 , we compare the 10-shot performance of our model with current State-Of-The-Art (SOTA) methods. Among them, 'HSNet', 'DGPNet' and 'VAT' are adapted from previous SOTA few-shot segmentation methods, following [18] . Our UniDense method consistently outperforms the previous SOTA few-shot method VTM across all tasks. Figure 3 illustrates the performance of UniDense and current SOTA methods as we increase the size of the support set from 10 to 275. Our method exhibits a significant advantage over VTM across all shot numbers. Remarkably, our model even surpasses the fully supervised method InvPT on 9 out of 10 tasks using a much smaller dataset (0.1%). This implies the potential value Table 2 : Comparing different few-shot adaptation methods while controlling all models' backbone as ViT [5] .Backbone SS (mIoU↑) SN (mErr↓) VTM w/o FT [18] ViT-B 0.0002 23.4212 Painter [40] ViT-L 0.1167 19.0734 Linear HeadViT-B 0.2681 13.0704 VTM [18] ViT-B 0.4097 11.4391 UniDense-ViTViT-B 0.4158 10.2803 of our method in specialized domains (e.g., medical images) where the number of available labels ranges from dozens to hundreds.Our method's superior performance is attributed to the efficient few-shot adaptation method of router fine-tuning with meta-routers and the reliable backbone based on the Stable Diffusion model. This section conducts ablation experiments to investigate the individual effectiveness of each design.Comparing adaptation methods with same backbone. Due to the architectural differences between our method and VTM, a direct comparison may be unfair. To demonstrate the effectiveness of our few-shot adaptation strategy, we perform an ablation experiment by replacing the stable diffusion backbone with a ViT backbone. Specifically, we customize MoE attention layers and MoE MLP layers within the transformer blocks of the DPT-Base model [27] to create a MoE ViT model. We then apply our three-stage UniDense pipeline to this MoE ViT model, resulting in the 'UniDense-ViT' model. Table 2 presents the results. The upper section of methods represents few-shot adaptation without fine-tuning, where Painter [40] adopts the in-context learning paradigm. The lower section of methods involves fine-tuning for few-shot adaptation, where 'Linear Head' denotes the method that fine-tunes only the linear mapping layer while freezing the backbone. Our UniDense-ViT method outperforms VTM with the same ViT-B backbone and achieves a significant improvement over the advanced in-context learning method Painter. These results validate the effectiveness of our few-shot adaptation strategy.Effect of Meta-Routers. To realize the idea of few-shot adaptation by selecting task-specific computation pathways, we propose to fine-tune the routers of the MoE backbone. However, this raises questions about which router to fine-tune and how to ensure its effectiveness. Here, we compare three different ways to construct routers for few-shot fine-tuning in Table 3 . 'Random Router' represents randomly initializing a set of routers for few-shot fine-tuning. 'Meta-Router w/o Cond.' represents using meta-routers without task conditions. As observed in the table, meta-routers significantly surpass randomly initialized routers, and the involvement of task conditions further improves the performance.Comparing different methods exploiting Stable Diffusion. While the pre-trained knowledge of Stable Diffusion (SD) greatly facilitates our method, it is worth exploring if there are more In this section, we perform visualization experiments to demonstrate that our method is capable of extracting task-specific knowledge by selecting relevant experts for different tasks, and that our method can effectively transfer knowledge between relevant tasks. We also visualize the predictions generated by our method to demonstrate its effectiveness. Activation frequencies of experts for different tasks. One advantage of our method is the ability to extract task-specific knowledge by selecting relevant experts for different tasks. To demonstrate this, we visualize the frequencies of activated experts for different few-shot tasks. In Figure 4 , the left part displays the expert frequencies obtained using the load-balancing loss, which promotes similar utilization frequencies of experts across different tasks. In contrast, the mutual information loss (right) adopted by our method encourages expert specialization, resulting in a sparse but strong dependence between experts and tasks.Transfer efficiency of expert knowledge. We are interested in understanding which task's knowledge is more useful than others for a specific task. To address this, we define the transfer efficiency of expert knowledge as the percentage of experts that the target few-shot task shares with the pre-training task. Figure 5 presents the results for all data folds. From the table, we observe that there is a higher transfer efficiency between 2D tasks (TE, K2) and between 3D tasks (ED, ZD, SN, RS, PC, K3, OE). This finding is reasonable and demonstrates that our framework can effectively transfer useful knowledge between relevant tasks.Qualitative comparison. In Figure 6 , we present visualizations of the predictions generated by our method and previous SOTA methods. Comparing our method to VTM, we observe that our method produces more accurate predictions across all tasks, which aligns with the results reported in Table 1 . 
conclusion and discussion: 
 In this paper, we tackle the universal few-shot learning problem from a novel perspective. Unlike traditional few-shot fine-tuning techniques that directly use all parameters of the pre-trained model and modify a specific set of weights, our approach focuses on selecting task-relevant computation pathways while keeping the model weights frozen. To achieve this, we introduce a novel framework called UniDense. Our method utilizes the Stable Diffusion model to construct a versatile MoE architecture and fine-tunes its routers to achieve few-shot adaptation. To ensure efficient router fine-tuning, we leverage episodes-based meta-learning to train a set of routers known as Meta-Routers. Additionally, we introduce a task embedding extractor that generates task conditions directly from support samples for the meta-routers. By fine-tuning these meta-routers for novel tasks, we enable highly efficient adaptation of the entire model. Experimental results on a modified version of the Taskonomy dataset demonstrate the superiority of our method.  
: 
 of computer vision. Their objective is to learn a mapping from input images to pixel-wise annotated labels. Given the high cost of pixelwise labeling for supervised methods, there is a strong demand for developing a few-shot learner which can flexibly and efficiently learn any dense task from a few labeled images. This specific task is referred to as universal few-shot dense prediction [18] , which has been recently proposed. The ultimate objective of few-shot learning is to emulate the human brain's capacity to learn arbitrary new tasks with minimal samples. The task of universal few-shot dense prediction serves as a testing ground for evaluating such learning approaches in the domain of dense prediction tasks.A crucial requirement for universal few-shot dense prediction is the ability of the model to adapt flexibly to diverse unseen tasks while being efficient enough to avoid over-fitting [18] . Several fewshot learning methods, including metric learning and in-context learning, have emerged to address this challenge by enabling rapid adaptation to novel few-shot tasks without the need for additional fine-tuning [10, 36, 37, 40] . However, numerous experiments consistently demonstrate that models with fine-tuning still achieve the highest performance [15, 20, 21, 24] . The naive approach involves fine-tuning the entire model, as depicted in Figure 1a . Nevertheless, the limited samples in the few-shot setting is insufficient for fullmodel fine-tuning, resulting in overfitting. To address this issue, current fine-tuning methods often resort to partial fine-tuning. A common approach is to fine-tune the last few layers [1, 38, 42] , as shown in Figure 1b . Some methods also fine-tune the biases [2, 18, 44] , as depicted in Figure 1c . However, all these methods achieve few-shot adaptation by modifying the model weights, which may disrupt the well-learned knowledge from the pre-training stage. Moreover, recent experiments have demonstrated that not all features of a pretrained model are beneficial for novel few-shot tasks and selecting task-specific features may lead to improved performance [4, 22] . However, prevalent fine-tuning methods simply use all parameters of the model and lack the ability to extract pre-existing task-specific knowledge embedded within the pre-trained model.To overcome these limitations, our paper approaches universal few-shot dense prediction from a novel perspective. Unlike conventional fine-tuning techniques that directly use all parameters of the model and modify a specific set of weights to achieve fewshot adaptation, our method focuses on selecting the task-specific computation pathways of the pre-trained model while keeping the model weights frozen, as shown in Figure 1d . Building upon this concept, we introduce a novel framework called UniDense for universal few-shot dense prediction. The basic idea of this framework is to construct a MoE (Mixture of Experts) backbone and achieve few-shot adaptation by fine-tuning its routers, which act as hyper-networks responsible for selecting task-specific computation blocks. Specifically, our UniDense framework follows a three-step pipeline: (1) We transform the UNet-like denoising autoencoder from Stable Diffusion (SD) [29] into a MoE model. This design allows us to leverage SD pre-trained knowledge, which has been empirically proven to be useful for a wide range of dense prediction tasks [1, 41, 48] . We then conduct multi-task training on this MoE model using datasets from various base training tasks (disjoint from testing tasks). After this, we assume that the model possesses the necessary general knowledge to tackle arbitrary dense prediction tasks. (2) We utilize episodes-based meta-learning to train a set of routers for this MoE model, called Meta-Routers. Each episode is designed to simulate the few-shot setting of testing phase, enabling the meta-routers to learn how to efficiently fine-tune on arbitrary unseen dense prediction tasks. (3) For novel few-shot tasks, we fine-tune these meta-routers to achieve efficient adaptation of the entire model by choosing task-specific computation pathways.Moreover, for each few-shot task, we leverage support samples to extract a task embedding, which serves as a conditioning factor for the meta-routers. This strategy allows meta-routers to dynamically adapt themselves for different few-shot task, instead of rely solely on fine-tuning a fixed set of initialization weights, thus increasing their flexibility and leading to improved adaptation performance.We follow the standard evaluation setup of universal few-shot dense prediction [18] to validate our method. A variant of the Taskonomy dataset [45] , consisting of ten tasks, is used for evaluation. The dataset is partitioned into a 5-fold split, where two tasks are selected for few-shot evaluation in each split, while the remaining eight tasks are used for training. This setup effectively simulates the few-shot learning scenario for unseen dense prediction tasks. Experimental results across all splits consistently demonstrate the superior performance of our proposed approach.Our paper makes the following key contributions: (1) We introduce UniDense, a novel framework for universal few-shot dense prediction, featuring a reliable and versatile MoE architecture based on the Stable Diffusion model, along with the technique of router fine-tuning to achieve efficient adaptation by selecting task-specific computation pathways instead of modifying model weights. (2) We propose meta-routers which are trained with episodes-based meta-learning to learn how to rapidly adapt. Moreover, a task embedding directly generated from the few-shot samples is used to dynamically adjust meta-routers, resulting in highly efficient model adaptation. (3) Experimental results on a variant of the Taskonomy dataset across all dataset splits and various shot numbers consistently demonstrate the superior performance of our approach. 
 
 Appendix:  


Decision: After rebuttal, reviewers agree on the acceptance of this paper, mostly positive comments about the introduction of meta routers. The major caveat is that it is a relatively complex method and potentially requires large computational resources. However, the merit of the meta-routing mechanism that dynamically adapts the diffusion process is reasonable to reviewers, moreover after rebuttal the outlook on the paper was more positive. Several ideas were highlighted as interesting hence the AC recommends acceptance.

Reviews: review: paper_summary: paper_summary: The paper introduces "UniDense," a novel framework designed to enhance the performance of few-shot dense prediction tasks across a variety of domains such as semantic segmentation, depth estimation, and surface normal prediction. The core innovation of UniDense lies in its use of diffusion models coupled with a meta-routing mechanism that dynamically adapts the diffusion process based on the task at hand. This approach allows UniDense to efficiently leverage a small number of examples to make accurate dense predictions, showcasing its versatility and effectiveness in a few-shot learning context.
strength_weakness: strength_weakness: 1. The introduction of meta-routers to guide diffusion models in the context of dense prediction tasks is an innovation. This approach addresses the challenge in few-shot learning by enabling more effective adaptation to new tasks with minimal examples.
2. The methodologies proposed in the paper are technically sound. The authors detail the implementation of UniDense, ensuring reproducibility and providing a clear understanding of the model's inner workings.
3. UniDense's ability to perform across different dense prediction tasks highlights its potential for wide-ranging applications in areas such as autonomous driving, robotics, and augmented reality, among others.
limitations: limitations: 0. **Not related to multimedia** This paper studies the dense prediction problem, which is unimodal. The authors claimed "This paper explores the application of pre-trained knowledge from the multi-model generative foundation model Stable Diffusion to a downstream task known as universal few-shot dense prediction". However, using stable diffusion for a computer vision problem can not be recognized as a contribution to MM processing.
1. **Can UniDense generalize to Unseen Domains**: While UniDense shows impressive performance on the evaluated tasks, the paper does not extensively discuss its ability to generalize to completely unseen domains or tasks that significantly deviate from the ones tested.
2. **Limited Discussion on Failure Cases**: The evaluation primarily focuses on scenarios where UniDense succeeds. A more comprehensive discussion on its limitations or failure cases, especially in extremely challenging few-shot scenarios, would provide a more balanced view of the model's capabilities.
3. **Lack of Discussion on Computational Resources**: The paper lacks a detailed analysis of the computational resources required for training and inference. It would be beneficial to understand the trade-offs between performance and computational efficiency and also compare with the SOTA in this aspect.

Since I am not quite familiar with the dense prediction task, I vote for borderline acceptance with a low confidence score.
paper_summary: paper_summary: Facing the risk of disrupting the pre-trained knowledge, this paper introduces UniDense for universal few-shot dense prediction. UniDense constructs a versatile MoE architecture and utilizes episodes-based meta-learning to train a set of meta-routers for the MoE model. The paper focuses on selecting the task-relevant computation pathways of the pre-trained model while keeping the model weights frozen. First, UniDense constructs a versatile MoE architecture for dense prediction based on the Stable Diffusion model. Then UniDense utilizes episodes-based meta-learning to train a set of routers for this MoE model, called Meta-Routers. Moreover, for each few-shot task, UniDense leverage support samples to extract a task embedding.
strength_weakness: strength_weakness: (1) The paper is logically.

(2) Sufficient experimental analysis.
limitations: limitations: (1) The method is relatively complex, and a set of Meta-Routersfor requires additional resource support.

(2) Except the introduction of diffusion model, other improvements are limited.
paper_summary: paper_summary: UniDense handles few-shot dense prediction tasks using three-stage training and meta-routers. In the first stage of training, balancing weights are updated, while in the second and third stages, meta-routers are refined. Meta-routers initialize routers through meta-learning. UniDense is the first to adopt Stable Diffusion for few-shot dense prediction tasks.
strength_weakness: strength_weakness: - Utilizing task embeddings from the support samples for meta-routers demonstrates a novel approach.
- The performance in the Taskonomy benchmark is impressive, achieving the highest scores across the tasks.
- The overall visualizations and figures are clear and helpful for understanding.
limitations: limitations: - The adoption of Stable Diffusion is mentioned as one of the main contributions. However, stable diffusion is adopted in rather a parallel manner with newly proposed methods, three-stage training and meta-routers. Can the authors explain more about the connection between the newly proposed methods and diffusion models? 
- Several architectures or equations have been selected from other works, such as SD, MoA, and mutual information loss from Mod-Squad, making it difficult to assess the performance boost attributed to the main methods. In Table 4, VTM with a ViT backbone and UniDense without MoE may seem to achieve comparable scores. Additionally, scores from other folds (Fold 2-5) would further clarify the ablation study.
paper_summary: paper_summary: This paper introduces a UniDense framework for the universal few-shot dense prediction problem.  The UniDense framework utilizes a versatile MoE architecture based on the Stable Diffusion model to select the task-specific computation pathway for few-shot adaptation. It employs episodes-based meta-learning to train Meta-Routers for this selection and uses task embeddings to dynamically adjust the routers. Experimental results across 10 different dense prediction tasks validate the effectiveness of UniDense.
strength_weakness: strength_weakness: Strengths         
1.	Introducing a pre-trained Stable Diffusion model to address the few-shot dense prediction problem is particularly compelling, especially with the incorporation of MoE attention mechanisms and MLP layer transformations within the Stable Diffusion architecture.  
2.	Extensive experiments show that the proposed method outperforms other baselines.
limitations: limitations: Weaknesses   
1.	In section 4.2, lines 314-317 mention that the model does not utilize the noise addition process of the Stable Diffusion model and also omits the language interface. This raises the question of whether it can still be considered a diffusion model. Without the diffusion process and relying solely on the UNet structure, the claim of using a diffusion model for dense prediction seems somewhat ambitious.     
2.	How does the task embedding extractor help the meta-routers? What is the difference between Stage B and Stage C in Figure 2 regarding the task embedding extractor's function?     
3.	[Minor] It would be better to explain that MoE stands for Mixture of Experts in the abstract. (L23)