DAMEX: Dataset-aware Mixture-of-Experts for visual understanding of mixture-of-datasets

Summary

This manuscript is motivated from training a universal object detector on the mixture of unlimited counts of datasets. Concretely, the previous works cannot scale the MoE due to budget, and the conventional MoE using balanced route suffers from knowledge sharing issues. To tackle the aforementioned issues, the authors devised a Dataset Aware Mixture of Experts (DAMEX), in which a novel MoE loss is proposed to learn the dataset index distribution. Then, extensive experiments are conducted on various settings to show the effectiveness.

Strengths

1. The paper is well-organized and well-motivated. The authors provide detailed explaination to the conventional MoE and provides certain discussion on the motivation. 2. The idea of this manuscript is self-consistency, where the motivation is to tackle the knowledge distribution shift among datasets for existing MoE, and the $\mathcal{L}_{DAMEX}$ pushes router learning the dataset index distribution seems to be work in intuition. 3. The conducted experiments are wide and adequate to demonstrate the effectiveness of the proposed method.

Weaknesses

1. The contribution of this manuscript is limited. All of the technical contributions are summarized as loss function of DAMEX, while it seems not settle the dataset distribution shift well. Concretely, as shown in Tab. 1, DAMEX fails to balance the results among datasets. As far as I am concerned, the dataset shift can be treated as hardness. To balance it, a larger improvement should be gained at harder datasets. Comparing with baseline DINO-MoE, the improvement of DAMEX concentrates on KITTI (3.9), COCO (1.1), DOTA (1.3), Clipart (2.4), Watercolor (3.6). There goes without strong correlation with dataset hardness. Welcome further discussion. 2. The manuscript has unfair expression. The authors claimed they outperform the state-of-the-art with almost 10% AP. I cannot agree with a paper published in 2019 could be called as sota, and the DINO looks to be sota.

Questions

Why the authors repeatedly mention GPU? Is there any efficient design for DAMEX? How to implement MoE on multi-GPU has been mentioned in Sec. 3, is this necessary to keep mentioning GPU setting?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

2 fair

Presentation

3 good

Contribution

2 fair

Limitations

As discussed in Weakness, I don' t agree with the authors tackle the dataset shift. But this can be further discussed.

Summary

The paper aims to develop a universal object detector that is applicable to a wide range of object detection datasets. For this aim, the authors proposes Dataset-aware Mixture of Experts (DAMEX). DAMEX is an extension of the vanilla Mixture of Experts layer to the multi-dataset scenario, in which each dataset is aimed to be routed to a specific expert. The authors also suggest replacing the fully connected layer in the transformer blocks by DAMEX layers. They incorporate such transformer blocks into DINO object detector. The experiments show that the approach improves upon previous SOTA as well as training a detector specifically for each dataset.

Strengths

- The paper adapts MoEs to obtain universal object detection, which is intuitive. - The approach does not increase the number of parameters significantly while obtaining a MoE for several datasets. - The gains in terms of AP are significant over previous baseline on UODB benchmark. Also, the improvement over training a different DINO for each dataset is notable with ~4 AP. The method also outperforms using a Vanilla MoE by 1.3 AP. - The analysis part is comprehensive and it shows that the method does work as expected. For example, the model learns to assign the experts to the datasets properly. - Generally speaking, the paper is clearly written.

Weaknesses

1. I think this task is closely related to open vocabulary object detection, on which there are already several different baseline methods. Such models are also evaluated on a set of multiple object detection datasets, e.g. using Object Detection in the Wild benchmark [A]. As an example method, GLIP [A] can easily be prompted given the union of the label spaces of each dataset in UODB benchmark. By considering recent advances, I'd expect at least an example of these models to be considered as a baseline. 2. The authors use UODB benchmark (compared to Wang et. al, 2019) to evaluate their method. This dataset normally uses Pascal VOC style AP. Here, the authors use COCO-style AP. That's why, the AP values are significantly low in this paper (compared to Wang et. al, 2019) owing to the choice of this evaluation measure. It took me a while to dig into this and see the difference in evaluation. I'd expect this to be very clear in the metrics section. 3. Please note that k index to define summation is not used in p_i or e_i. I'd recommend checking that equation. [A] Grounded Language-Image Pre-training, CVPR 2022

Questions

1. Do you think open vocabulary object detection methods are also valid baselines for this task? Specifically, why don't you use a model such as GLIP to compare your method against? 2. You mention that each GPU stores the weights of one expert. Following from this, do you pay specific attention to sampling the batch? What happens if all images (16 in total across all GPUs) come from the same dataset, in which case they should be sent to the same GPU, potentially causing memory issues? 3. The learning rate of the proposed method is tuned to be 1.414e-4. Did you tune the learning rate for each dataset while obtaining the single domain models (the first two lines in Table 1)? If so, how? 4. Why did you prefer using COCO-style AP instead of the performance measure of the benchmark that is Pascal VOC style AP?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

3 good

Contribution

3 good

Limitations

Yes

Summary

This work, motivated by the goal of developing a “universal detector,” seeks to understand how best to train a model across a large set of existing curated, labeled datasets that might differ in collection strategy, labeling standards, categories of interest, etc. They posit that the best strategy is to train dataset-specific features within a model that can be used to jointly predict on new images via mixture-of-experts, where dataset tokens are used to route images through the network based on the “expertise” of the learned features effectively and with many fewer parameters than prior multi-headed approaches. They show nice improvement using their method on the Universal Object-Detection Benchmark, and demonstrate that their method reduces representation collapse, a common issue with MoE training. After reading the authors rebuttal and discussing with them, I will maintain my score of 6.

Strengths

This work, which extends MoE-based dataset mixing to object detection, is differentiated from prior work using MoE for dataset mixing in image classification via the inclusion of specific per-dataset-expertise-routing layers within a larger shared-across-datasets architecture. They distribute expertise via a load balancing loss (similar to prior work), which they adapt to the detection setting by applying it only to foreground tokens (reducing the influence of background on expertise routing). They use explicit dataset tokens to route information to specific experts during training and inference, which requires train- and test-time dataset labels. DAMEX particularly seems to do well over vanilla MoE in a few-shot setting, as seen in the figure in Table 2.

Weaknesses

The claims of performance improvements of 10% on average are slightly misleading, as there are probably lots of additional contributing factors that are not attributable to their method that led to prior reported numbers being lower. They see only a 2.3% improvement on average over vanilla DINO, and their contribution of a dataset-specific routing layer only improves performance over DINO-MoE by 1.3%. On many test sets their method underperforms by a significant margin. I would appreciate a more nuanced and representative claim, and some additional investigation and analysis as to why their dataset performs better on, e.g., DOTA, and much worse on Watercolor or LISA. Since the dataset tokens are explicit, I believe the model would need to be completely retrained to add a single additional source of data, which is (perhaps prohibitively) inefficient. Maybe an extension of this work would be to enable add-on routing/expertise generation for new mixing datasets without end-to-end retraining across all datasets. Nit: I’m not sure the citation style matches the NeurIPS standard? Grammar is inconsistent and frequently in error throughout, with too many small mistakes to reasonably capture in this review. I would recommend significant text editing to increase the readability of a camera-ready manuscript for this work.

Questions

How would you add an additional dataset into the mix with this method? Would it require training from scratch? It seems in figure 3 that despite the load balancing, expertise is heavily relying on “expert 1” for many of the classes. How is this related to the class distribution across datasets?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

2 fair

Contribution

2 fair

Limitations

More nuanced evaluation and discussion of failure modes of their method on some datasets but high performance on others would help the reader gain intuition as to when this method would be worth implementing and applying. What about the individual dataset structure contributes to when this dataset-specific routing helps vs hurts? I do appreciate the discussion of label unification as a current limitation. The comment “Going forward it is important that we ensure that the universal object detector are unbiased and maintain fairness in their predictions” feels like a throwaway. How would this work possibly contribute to that or how could it be adapted to contribute to that goal? How is the current paradigm possibly entrenching bias?

Summary

This paper introduces a DAMEX layer based on the idea of assigning samples conforming to the charactersitiscs of a dataset to the corresponding expert. The underlying thought is to build dataset-relevant modules and ensemble them all together. Previous approaches leverage Mixture of Experts to scale their model while maintaining approximately the same inference speed. This work deals with the scenario of training multiple datasets together for object detection.

Strengths

The writing is clear. Results seem to prove the effectiveness of the proposed approach. The added number of parameters is small.

Weaknesses

1. There is no introduced novelty in the proposed approach. It is good to know dataset specific experts are good for improving performances on object detection but the proposed DAMEX is in principle a linear routing layer that maps the input to different experts. The training procedure is also very straightforward: training different datasets on different GPUs. Essentially, this is equivalent to model ensemble a per-block router. Most of credits should go to GShard and Tutel. 2. More clarification is needed for the improvement of performances. Throughout the results part, clarifications of why the proposed method is better than the others are missing. For example, I miss why DAMEX could deal with imbalanced datasets? Why it has better domain adaptation than DINO? 3. Scalability is limited. From the experimental setting and the results, the maximum number of experts is 8, which is pretty small. It seems that it is highly correlated with the number of GPUs as well. From Tab. 5, we can see that when multiple experts are distributed to a single dataset, the performance drop is significant, which demonstrates that the approach is not scalable. 4. Important experimental details are missing. For example, in Tab. 5, which two datasets are distributed together when # Experts = 2? The number of iterations needed for training is not given. 5. From the visualization in Fig. 3, DAMEX has more condensed cluster such as Expert 1. This may indicate that other experts are not exploited well and it might work to delete some experts with no obvious decrease on performance. This potentially would have robustness issues. Again, it is more like doing majority voting from the ensemble perspective.

Questions

See weaknesses please.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

2 fair

Presentation

3 good

Contribution

2 fair

Limitations

No. The authors mentioned the negative societal impact but does not have a solution yet.

Summary

This paper tackles the problem of mixture-of-datasets training for object detection. The authors propose a mixture-of-expert-based (MoE) model that utilizes dataset-specific features to tackle mixing of heterogeneous datasets. The backbone is based on DINO transformer and one expert is assigned to one dataset. Results on the UODB dataset shows the efficacy of the proposed method.

Strengths

1. The idea is straightforward and the intuition of using one expert per dataset makes sense. The method seems easy to implement and the authors promise to release code so the reproducibility is good. 2. The proposed method is better than the previous SOTA by a large margin.

Weaknesses

1. The design of assigning one expert for each dataset seems limited. The MoE model learns to route the tokens to a specific dataset expert, but in the real world there are many more than 11 datasets/styles. For images that are watercolor-comic etc., would the model get confused? Also, please discuss the need for a universal object detector trained on limited supervised datasets when large vision-language foundation models (VLMs) like [4*] can already get 60+ mAP (vs. 41 in this paper) on COCO. It would help to see the performance of the VLMs on these lesser-tested datasets like DOTA and Watercolor. 2. The performance on common object detection benchmarks (MSCOCO, VOC, etc.) is lower than baselines. This may suggest that some features are difficult to learn under the MoE setting.

Questions

1. The line of work on mixture-of-datasets for video understanding should also be discussed [1*, 2*, 3*] in related work. Also, open-vocabulary object detection based on large vision-language models (VLMs) [4*] should also be considered. 2. A inference speed comparison is needed. What is the FPS, for example, of the proposed method vs. others? 3. Other minor comments: The citation format makes some sentences strange: “Firstly, these datasets have been collected over time Everingham et al. [2015] and …” (Line 29) The text in Figure 2 is too small to read. [1*] Akbari, H., Yuan, L., Qian, R., Chuang, W. H., Chang, S. F., Cui, Y., & Gong, B. (2021). Vatt: Transformers for multimodal self-supervised learning from raw video, audio and text. Advances in Neural Information Processing Systems, 34, 24206-24221. [2*] Duan, H., Zhao, Y., Xiong, Y., Liu, W., & Lin, D. (2020, August). Omni-sourced webly-supervised learning for video recognition. In European Conference on Computer Vision (pp. 670-688). Cham: Springer International Publishing. [3*] Liang, J., Zhang, E., Zhang, J., & Shen, C. (2022). Multi-dataset Training of Transformers for Robust Action Recognition. Advances in Neural Information Processing Systems, 35, 14475-14488. [4*] Yuan, L., Chen, D., Chen, Y. L., Codella, N., Dai, X., Gao, J., ... & Zhang, P. (2021). Florence: A new foundation model for computer vision. arXiv preprint arXiv:2111.11432. --------------------------Post rebuttal I have read the author rebuttal and other reviewers. The authors have addressed my concerns and questions. I'm increasing my score.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The authors have discussed the limitation and potential issues of the research.

What about the potential implications of the work on bias and fairness, and mentioned in the text by the authors?

As I mentioned in my review, "The comment “Going forward it is important that we ensure that the universal object detector are unbiased and maintain fairness in their predictions” feels like a throwaway. How would this work possibly contribute to that or how could it be adapted to contribute to that goal? How is the current paradigm possibly entrenching bias?" I would appreciate if the authors respond to this comment, as I believe engaging more deeply with broader questions about the implications of our work is an important and necessary opportunity for growth in our field. I thank the authors for their thoughtful responses to my other questions and concerns.

Post rebuttal

I've read the authors' response. The author response addresses most of my concerns. I'm happy to increase my score.