SHIELD: Multi-task Multi-distribution Vehicle Routing Solver with Sparsity & Hierarchy in Efficiently Layered Decoder

We thank the reviewer for the positive recognition on the effectiveness of our method SHIELD as well as the contributions we have written. We would like to address your concerns as follows:

[Q1: Model runtime]: Thanks for your comment. Actually, we reported the inference time in the tables 13-21 of Appendix A.8. Following the reviewer suggestions, we further averaged the runtime of the models reported in Appendix A.8 and updated Table 1 to reflect them. In specific, it refers to the average solving time on 1000 test instances over 9 countries. For convenience, we tabulate the number of parameters and runtimes here. From the table, we can see that utilizing SHIELD improves the inference time of a denser MVMoE-Deeper, and also renders MTMDVRP100 trainable.

[Q2: Other VRP Approaches]: We acknowledge the reviewer's feedback regarding other VRP methods. However, we note that many of recent works are focused on the single-task VRP setting, and that the multi-task and single-task research directions are orthogonal to each other. Generalizing single-task VRP solvers to the Multi-task Multi-distribution setting is non-trivial, as shown in the case for POMO-MTVRP. Following the ICLR guidelines, for Multi-task VRP solvers, we have included all recent published works as baselines in the form of POMO-MTVRP and MVMoE. We have also provided multiple benchmark models that expand on MVMoE so that the overall model capacity as equivalent as possible. Please let us know if you have any other recommendations. We refer the reader to Appendix A.1 for a comprehensive literature review and discussion on neural single-task VRP solvers.

[Q3: Explanation on $\mathcal{Q}$]: We have updated Figure 1 in the manuscript to reflect the overall process from sampling to training and finally inference. As shown, we designate USA13509, JA9847, and BM33708 as known maps (i.e., in-distribution). These form the set of distributions $\mathcal{Q}'$, which is a subset of $\mathcal{Q}$. Also, we have a set of known tasks (i.e., in-task) $\mathcal{K}'$, containing CVRP, VRPB, OVRP, VRPTW, VRPL, and OVRPTW, which is a subset of $\mathcal{K}$. At each step, we sample a map $q' \in \mathcal{Q}'$ uniformly, followed by sampling a set of coordinates (cities) from the map. Next, we sample a task $k'$ from the set of $\mathcal{K}'$. This forms a batch of training data and is passed through the model. Note that for the proposed MTMDVRP setting, we consider 9 countires map and 16 VRP tasks as $\mathcal{Q}$ and $\mathcal{K}$ in this paper, while $\mathcal{Q}$ and $\mathcal{K}$ could be any world maps or related-tasks in theory. During inference, we apply the model to the entire set of tasks $\mathcal{K}$ and distributions $\mathcal{Q}$, whereby the remaining 10 tasks and 6 distributions are considered to be out-task and out-distribution. For a formal definition, please see Section 4.1 in our revised paper.

[Q4: Training on Uniform only]: We do appreciate the reviewer pointing out this key point and have further trained MTMDVRP50 models on uniform distribution, followed by evaluating them on the 9 distributions discussed in this paper. We retain the same table format as Table 1, but please note that in this case, all distributions are technically out-of-distribution, as the models are trained purely on uniform data. We exclude runtimes in the following table as we have updated Table 1 with overall runtimes for the various models. Also, due to time constraints, we were only available to train the models on MTMDVRP50 settings. As shown, when the models are not trained on any structured distributions and only on the uniform one, there is a degradation in performance all around. This necessitates the training of a sufficiently flexible model on varied and structured distributions, so as to train a strong foundation model capable of realistic applications.

We thank the reviewer for acknowledging our good writing, and positive performance and impact of SHIELD. We understand that the main concern is regarding the motivation of MoD and MTMDVRP. We hope the following rebuttal could address your concerns.

[W1: Motivation of MoE vs MoD]: While the reviewer highlights the differences in sparsity between MoE and MoD, we would like to clarify that the MoE mechanism used in MVMoE achieves only partial sparsity by selecting specific experts, that is, determining which sub-network to activate. Consequently, the forward computation in MVMoE requires a similar amount of compute as the POMO-MTVRP framework. This similarity arises because MVMoE replaces the dense multi-layer perceptron (MLP) blocks in POMO-MTVRP with an MoE layer, which retains the same input and output dimensions. Although only a portion of the MoE layer is activated during inference or training, the overall parameter usage remains comparable for each input token.

In contrast, the MoD employed in SHIELD presents 3 layers in the decoder, each of which independently decides a fraction of the tokens to be processed. This encourages the network to prioritize tokens that are more important. Figure 3 (right panel) showcases this effect, whereby starting nodes and ending nodes in sub-routes utilize more layers, whereas those in the middle do not. As such, the computation can vary for each input token, enhancing the model's flexibility in representation learning and decision-making.

Moreover, we argue that since all parameters are shared, the sparsity introduced by MoD forces the network to learn the best set of parameters that can be generalizable across all tasks and distributions. This behavior aligns with that of Information Bottleneck as mentioned in Section 4.2, where studies found that highly predictive representations with minimal complexity improved generalization as a whole, suggesting the need to balance the shared and task-specific information. Therefore, by introducing sparsity on the network through MoD, we force the network to learn a set of minimal representations capable of solving various tasks and distributions. Our empirical results in Table 2 justify the importance of enforcing sparsity: when we increase the number of tokens allowed, thereby reducing sparsity, the network's in-task in-distribution performance improves, but its OOD generalization starts to worsen. On the extreme end, in the case of MVMoE-Deeper where all tokens are available, the network has the worst performance compared to the sparser variants.

[W2: Contribution of Soft-clustering]: While soft-clustering has been previously explored, its capabilities in learning foundation models for generalization in combinatorial optimization problems have not yet been studied, to the best of our knowledge. Note that the soft clustering approach in (Goh et el., 2024) was only centered on learning neural TSP solvers and the neural solver was trained specifically for one distribution. Differently, in our paper, we explore the generalization properties of foundation VRP models. With the introduction of the MTMDVRP scenario, we test the robustness of the soft clustering across 16 different problems and 9 distributions. Technically, we introduce the context-based prompt so as to make the soft clustering mechanism more suitable to the MTMDVRP. We believe this work is the first to identify the key benefits of having such a clustering mechanism in the space of foundational neural combinatorial optimization solvers.

Thank you to the authors for the responses. I do appreciate the authors’ effort during the discussion phase, but I still hold the opinion in the initial comments. Here are the detailed reasons that I decide to give my rating as 3.

1. As the generalization performance of neural solvers has always been an important issue in the community, both multi-task generalization and multi-distribution generalization have been studied before. I do not think that highlighting the scenario of multi-task multi-distribution generalization could become one of the strengths for a high-quality paper. Meanwhile, I do not think that there are some specific designs and advantages of the proposed method on the multi-task multi-distribution. In other words, it seems that employing MoD to potentially emphasizing important nodes, and employing soft-clustering can also benefit multi-scale or multi-distribution scenarios as well.

2. The main ideas in the proposed methods provide rare new insights. I agree that MoE and MoD are achieving sparsity from different perspectives, but I do not think that this can be regarded as a main reason for that MoD should be better. Even that there are some intuitive explanations with Figure 3 and Figure 4, the results are not significant and convincing enough. For example, in figure 3, the explanations about right panel are very weak. In figure 4, it is hard for me to recognize the relationship among 3 maps on emphasized nodes. Therefore, I tend to regard this part as an application of existing method MoD with weak motivation. Besides, with the consensus that the idea of soft-clustering has been studied before, I do not think that extending it to CVRP problems and multi-task multi-distribution scenarios with rare new findings is very interesting.

3. The supplemented results on set-X are very limited, which is only a simple test of the pretrained models on MTMDVRP. As there are various powerful solvers that can achieve better performance on set-X, more comparison with them under the same training settings would be interesting.

3. While there are powerful solvers that can achieve better performance on Set-X, the models here are trained on a multitude of different tasks and distributions. This set of parameters do well on our scenario, in addition to performing well on a form of size generalization. Additionally, the feedback proposed to us in general was regarding size generalization, which we interpret to be an upward generalization in overall size. As such, we believe that the performance on Set-X, which contain most of the larger problems in CVRPLib, would satisfy such a scenario. In fact, powerful solvers such as LEHD, fail to train with reinforcement in a reasonable amount of time, as described in their paper [4]. Couple their high training cost with the addition of a multi-task multi-distribution scenario would result in a large increase in training time as convergence gets significantly harder due to the complexity of the scenario.

[1] Lin, Zhuoyi, et al. "Cross-problem learning for solving vehicle routing problems." arXiv preprint arXiv:2404.11677 (2024).

[2] Bi, Jieyi, et al. "Learning generalizable models for vehicle routing problems via knowledge distillation." Advances in Neural Information Processing Systems 35 (2022): 31226-31238.

[3] Zhou, Jianan, et al. "Towards omni-generalizable neural methods for vehicle routing problems." International Conference on Machine Learning. PMLR, 2023.

[4] Luo, Fu, et al. "Neural combinatorial optimization with heavy decoder: Toward large scale generalization." Advances in Neural Information Processing Systems 36 (2023): 8845-8864.

[5] Liu, Fei, et al. "Multi-task learning for routing problem with cross-problem zero-shot generalization." Proceedings of the 30th ACM SIGKDD Conference on Knowledge Discovery and Data Mining. 2024.

[6] Zhou, Jianan, et al. "MVMoE: Multi-Task Vehicle Routing Solver with Mixture-of-Experts." arXiv preprint arXiv:2405.01029 (2024).

We thank the reviewer for the valuable feedback which helps a lot to further strengthen our work. We understand the main concern about the importance of the MTMDVRP setting and we hope our response below with new empirical results would clear your concerns.

[W1: More details about MTMDVRP]: Great feedback! We have revised the manuscript to address this in Section 4.1. The MTMDVRP scenario is particularly significant since it 1) reflects a more practical application encountered by the industry and 2) represents a critical step toward developing foundation models for academic research.

From the application perspective, suppose a well-known logistic company X has established its presence in a handful of countries/cities. It is able to train a model based on collected data across these countries/cities. Now, if company X wishes to expand its market to newer ones, it definitely will face data from new distributions, and potentially faces new tasks. As such, it is highly beneficial that its trained model is able to be quickly applied to new incoming data - suggesting the need for the model to be robust to new unseen tasks and distributions. We note that while the original MTVRP setting has taken steps toward learning a unified model for various VRP variants, all customers within the instances are uniformly generated. This does not reflect practical distributions, where customer locations are following underlying patterns within a country/city.

From the machine learning perspective, foundation models typically exhibit similar behavior, whereby they are singular models that have learnt general representations capable of diverse applications across tasks and distributions. The MTMDVRP scenario is a crucial scenario that exhibits such challenges, the models trained have to be sufficiently flexible to generalize across not only tasks but also distributions. As explained before, facing new tasks and distributions is a very practical scenario that industrial companies encounter.

This is the MTMDVRP scenario we are introducing - from 9 countries, we sub-select 3 countries and denote them as the ones that company X has established presence (in our work, we use USA13509, JA9847, and BM33708 during training). Then, we challenge the models to train only on a subset of tasks (i.e., CVRP, OVRP, VRPB, VRPL, VRPTW, OVRPTW) drawn from these 3 distributions, and observe their OOD performances on data drawn from the remaining 6 distributions and 10 tasks. As such, we can observe how the models generalize across either task, distribution, or both.

[W2: Benefits of MoD]: Thanks for your comment. We understand the reviewer's concerns regarding the overall benefits of MoD. We have opted to remove the lines 285-288, and added a more detailed analysis of our results in lines 420-428. Specifically, in our experiments in Table 1 and ablation study in Table 2 of the manuscript, we have presented the results of varying the processing level. We note that the combined MVMoE-Deeper and MVMoE in Table 1 actually represent two extreme cases in model design, illustrating the trade-off between over-processing and under-processing tokens. In specific, MVMoE-Deeper has a decoder with the same number of layers as SHIELD and SHIELD-MoD (i.e., without clustering), but all tokens are needed to be processed at every layer (the case of over-processing). In contrast, MVMoE has only a single decoder layer that processes tokens (under-processing). These models exemplify the scenarios mentioned by the reviewer: over-processing (MVMoE-Deeper) and under-processing (MVMoE). Our method, positioned between these extremes, processes only a subset of tokens in each layer, with this operation performed independently for every layer for flexibility.

From Table 1, we observe that in general, MVMoE-Deeper is a superior model to MVMoE, suggesting that increasing the capacity of the decoder is beneficial to the model. Unfortuately, a consequence of doing so renders MVMoE-Deeper to be untrainable on larger instances. SHIELD presents itself as compromise between the two ends, we opt for a larger capacity decoder but force the network to learn which tokens should be processed more frequently. Consequently, we see that both SHIELD and SHIELD-MoD are capable of overall stronger in-task in-distribution performance as shown in Table 1. Apart from that, both models also shows stronger generalization properties across both tasks. Furthermore, the results in Table 2 of the manuscript further demonstrate how the number of tokens processed (i.e., processing level) at each layer affects model performance. Increasing processing power improves in-task, and in-distribution performance but comes at the cost of reduced generalization. There is also a turning point at SHIELD(30%) whereby the model starts to deterioriate as a whole both in and out of task and distributions. Please see our analyses in Section 5.2.

Thank you for your rebuttal. However, my decision remains the same. I believe this paper could be improved and resubmitted.

Dear Reviewer #n7rT,

Thank you for your previous constructive comments, which has indeed helped us to improve our paper. Following the comments from you and other reviewers, we have carefully revised the paper and uploaded the updated version in accordance with the ICLR policy. As outlined in the ICLR reviewer guidelines:

The discussion phase at ICLR is different from most conferences in the AI/ML community. During this phase, reviewers, authors and area chairs engage in asynchronous discussion and authors are allowed to revise their submissions to address concerns that arise. It is crucial that you are actively engaged during this phase. Maintain a spirit of openness to changing your initial recommendation (either to a more positive or more negative) rating.

We kindly request that you review our revised paper and let us know if there are any remaining concerns. As the rebuttal phase is still active with a week to go, we would be grateful for the opportunity to address any further points you might raise. We deeply appreciate the time and effort you have dedicated to reviewing our work and look forward to your feedback.

Dear Reviewer #n7rT,

We have made significant efforts to address all your concerns. For your convenience, we have tabulated the key revisions and TL;DR points below based on your previous constructive comments.

We hope the table above provides clarity and assists in your review process.

We thank the reviewer for their positive feedback on the practicality of MTMDVRP, as well as the recognition of its outstanding performance and the detailed ablation studies. We would like to address your concerns as follows:

[W1: Meaning of symbols]: Sorry for the confusion. We have revised the manuscript to better reflect the notations used and have added a notation table in Appendix A.6 for all symbols. For convenience, we provide the description of $\mathcal{D}_t = \{z_t, l_t, t_t, o_t \}$ here as well, copied verbatim from the manuscript (lines 241-247).

For our setup, we adopt the following feature set. At each epoch, we are faced with a problem instance $i$ such that $\mathcal{S}_i=\{x_i, y_i, \delta_i, w^o_i, w^c_i \}$, where $x_i$ and $y_i$ are the respective coordinates, $\delta_i$ the demand, $w^o_i$ and $w^c_i$ the respective opening and closing times of the time window. This is passed through the encoder resulting in a set $\mathbf{H}$ of $d$-dimensional embeddings. At the $t$-th decoding step, the decoder receives this set of embeddings $\mathbf{H}$, the clustering embeddings $\mathbf{C}$, and a set of dynamic features $\mathcal{D}_t = \{z_t, l_t, t_t, o_t \}$, where $z_t$ denotes the remaining capacity of the vehicle, $l_t$ the length of the current partial route, $t_t$ the current time step, and $o_t$ indicates if the route is an open route or not.

[W2: Clarity of Section 4]: We value your feedback! We have carefully revised the manuscript to enhance clarity. The changes can be found in the highlighted section of Section 4. As a whole, we provided more detailed description of MTMDVRP in lines 198-210, changed Figure 1 to reflect the overall process of our approach, and added clear details regarding the variables used in lines 241-247. Additionally, we have improved the explanation and description of the MoD layer in lines 279-285, as well as explained in greater detail the contextual clustering approach in lines 318-325, and lines 334-337. We have also added a full algorithm description of the soft clustering in Appendix A.4, with a full set of mathematical notations in Appendix A.6.

With regards to section 4.3, we have opted to remove the confusing notation of $k$ and $N-k$, and only maintain $\beta$. $\beta$ refers to the percentage of tokens allowed through a MoD layer. By producing a set of scores via the router, we take the top $\beta$-th percentile of tokens.

With regards to section 4.4, we have revised the section accordingly and changed some notations to synchronize the overall use of them. $\gamma_k$ is a one-hot encoded vector for task $k$. For a task $k$, $\gamma_k = [\gamma^1_k,\gamma^2_k,\gamma^3_k,\gamma^4_k]$, where $\gamma^1_k$ denotes open, $\gamma^2_k$ denotes time-window, $\gamma^3_k$ denotes route length, and $\gamma^4_k$ denotes backhaul constraints. $W_\theta^\top$ refers to a set of learnable parameters used to convert the one-hot encoded vector $\gamma_k$ into a latent representation $\alpha_k$.

We hope this provides greater clarity and welcome any additional suggestions you may have!

[W3: Large-scale Generalization]: Thanks for your suggestion. The focus of this paper is regarding the generalization properties across tasks and distributions. Nevertheless, we recognize the importance of size generalization and have further conducted multiple experiments on larger-scale CVRP. We apply our MTMDVRP100 models to two sets of data from Set-X in the CVRPLib, used in the DIMACS competition. Set-X-1 contains problems of various distributions and sizes, ranging from 101 nodes to 251 nodes (28 instances). Set-X-2 contains problems of various distributions and sizes, ranging from 502 nodes to 1001 nodes (32 instances). Below are the average performance. Note that since we are using MTMDVRP100 models, MVMoE-Deeper could not be run as it is too large to be trained on a A100-80GB GPU.

As shown by the above 2 tables, our model provides significant benefits in terms of size generalization as well. In addition to this, we have added another model, SHIELD-Ep400. This model is a checkpoint of SHIELD at the 400th epoch, where we find that its in-task in-distribution performance is similar to that of MVMoE.

Thank you for your response. The revised version has indeed improved the paper's readability, enhancing its overall clarity. The score I provided in the initial review already reflects my assessment of the methodology and experimental sections. I will temporarily retain this score, though I believe the rating could reasonably be raised to 7 if such an option were available.