UniteFormer: Unifying Node and Edge Modalities in Transformers for Vehicle Routing Problems

Thank you for your constructive and valuable comments and suggestions.

Q0:Related Work

A0: (1) Scope of prior work: Our study builds on an established body of research in neural VRP solvers. Prior works such as POMO, LEHD, RELD, GCN-BS, and Pointerformer (most of them are published in NeurIPS/ICLR) predominantly adopt GNNs or Transformer-based architectures within the VRP domain. To remain focused and highlight our specific contribution to modality-flexible neural solvers, we limited our related work to this area. Our motivation further stems from the lack of models that effectively integrate GNNs and Transformers for VRPs under varying or incomplete input modalities, which guided the design of our unified architecture.

(2) Broader works outside VRPs: We appreciate the reviewer highlighting six relevant works[*1]–[*6] from other domains. While they share certain architectural elements, our method differs fundamentally in both objective and design, e.g.,: HGT[*1] models type-specific message passing in heterogeneous graphs, whereas UniteFormer addresses modality heterogeneity (node/edge/hybrid) for route optimization. GraphGPT[*2] focuses on Euclideanizing graphs for representation learning; we focus on path construction by jointly encoding structure and distances. AHEAD[*3] detects graph anomalies via triple-attention, while UniteFormer uses REINFORCE-based training to generate feasible tours in VRPs. WeaveNet[*4] solves bipartite matching over complete graphs; UniteFormer supports incomplete or partial input modalities, which is critical for real-world VRPs. [*5] addresses 3D point cloud registration using geometric alignment; our work is focused on graph-based sequential decision-making. MGT[*6] models long-range molecular interactions for property prediction, while UniteFormer focuses on combinatorial path generation and input-modality adaptation, not regression or graph-level inference.

Although there may be surface-level architectural similarities (e.g., use of attention or GNNs), UniteFormer is fundamentally designed for path construction under modality-variant inputs, which is non-trivival and crucial for real-world VRPs. To our knowledge, UniteFormer is the first unified model to support all three modalities within a single framework. Each component is purposefully designed to handle node, edge, and hybrid in a scalable and flexible manner, combining the strengths of GNNs (for edge reasoning) and attention mechanisms (for node embedding). However, we agree that referencing broader literature like [*1]–[*6] will strengthen the generality and novelty of our method.

Q1:Existing methods inadequate?

A1: (1) The core limitation of existing methods lies not in the use of masked inputs, but in their reliance on a single input modality, i.e., node-only, edge-only, or hybrid, making them less adaptable to real-world scenarios with missing or varying input types. For example: (i)Edge-only methods like MatNet may underutilize node coordinates even when available; (ii)Node-only methods such as POMO and LEHD cannot directly process edge input without redesign and retraining; (iii)Models like GCN-BS or DAR, built for hybrid inputs, perform poorly when only one modality is present due to their architectural dependence on both. To our knowledge, no existing work offers a unified framework that flexibly handles node-only, edge-only, and hybrid inputs within one model, despite its clear practical value. This motivates our design of a unified architecture that adapts to all modalities without the need for retraining.

(2) Regarding your point on GNNs learning to internally complement node and edge features: we agree in principle. In fact, our model includes GCNs as a core component, precisely to capture structural dependencies from edge and node. However, through extensive experiments and ablation studies, we find that GCNs alone struggle to perform consistently across all input modalities, especially when only node or edge data is provided. To address this, we design a parallel mixed encoder combining GCNs and attention, enabling effective capture of both node- and edge-level semantics. This design ensures robust performance across all input modalities, as evidenced by: (i)Quantitative comparisons in Table 1, (ii)Component-level ablations in Section 4.2, and (iii)Additional targeted ablations (see A3).

Q2:Other Tasks and Three Inputs

A2: (1) VRPs are a core class of combinatorial optimization problems (COPs) with widespread applications in logistics and transportation. Given their practical relevance and structural complexity, developing specialized and efficient models for VRPs is both common and well-justified in existing literature, such as LEHD, GCN-BS, DAR, ICAM, focusing solely on VRPs. Our work follows this tradition, ensuring fair comparisons with established baselines. Moreover, our model further extends this research by supporting all three input modalities, offering greater flexibility than those prior single-modality methods. Moreover, UniteFormer can also handle asymmetric problems like ATSP, which node-only models such as POMO and LEHD cannot directly address.

However, UniteFormer's architecture is inherently general and we extend it to the 0-1 Knapsack Problem(KP). The results below for KP50 with 200 training epochs on both POMO and UniteFormer show that UniteFormer maintains impressive performance beyond VRPs. This further confirms its potential of UniteFormer as a general-purpose neural solver for other tasks beyond VRPs.

(2) While prior methods like POMO is restricted to a single input modality and require architectural changes to accommodate others, real-world VRPs demand greater flexibility. Node-based inputs are common in Euclidean settings, edge-based inputs are essential for non-Euclidean or asymmetric tasks like ATSP, and hybrid inputs combining coordinates and distances are increasingly important in practical logistics scenarios. UniteFormer is the first framework, to our knowledge, that supports all three modalities, i.e., node, edge, and hybrid, within a unified architecture, reducing the need for retraining. As shown in Table 1 and Section 4.2, it delivers strong and consistent performance across all input types. This modality flexibility is crucial in real-world deployments, e.g., node coordinates may be unavailable or noisy, while edge-level data is often more reliable.

Q3:Structural Rationality

A3: First, to clarify, our architecture includes both GCNs and attention, and there is no absence of GCNs. So we guess what you mean by this sentence is: "If attention alone... then the inclusion of a GCN seems odd."

(1) Complementarity: As described in Section 3.2, we adopt a parallel GCN-attention encoder to jointly process node and edge features and capture their cross-modal interactions. While attention mechanisms (e.g., POMO, LEHD) are effective for node-centric representations, GCNs are well-suited for learning from structured graph data, especially when edge input is available or critical. This hybridization is particularly beneficial in tasks like CVRP and ATSP, where edge weights (e.g., asymmetries, capacities) encode essential semantics that attention alone may overlook. In fact, Section 3.3(lines 157–161) further explains that attention modules are introduced to complement the limited node-level expressiveness of GCNs, ensuring a balanced and expressive node embedding.

(2) Input Modality: UniteFormer is designed to handle three input modalities. Given this, attention and GCNs serve distinct yet complementary roles: GCNs are naturally suited for Edge-only and hybrid inputs, enabling relational reasoning over graphs. Attention excels in processing Node-only data, especially when no explicit graph structure is available. Empirically, attention-only methods like POMO work well for Node-only inputs but perform poorly when edge features dominate. As shown in Table 1, our variant "POMO-edge"(a modified POMO model applied to edge-based input) performs notably worse than UniteFormer. This confirms that attention alone is insufficient for edge-dominant tasks. Moreover, in problems like ATSP, where edge asymmetry is key and node coordinates cannot infer distance, GCN-based reasoning is indispensable. As demonstrated in Appendix F, our hybrid model surpasses attention-only methods like MatNet on ATSP.

(3) Empirical Validation: To validate our architectural design, we conducted four additional ablation studies (see response A4 to Reviewer HhBh), examining (i) the complete removal of the GCNs module and (ii) reordering by placing self-attention before GCNs. As shown in the table below, both changes led to noticeable performance drops, underscoring the necessity of GCN and the soundness of our original design.

In addition, we further conducted ablation studies on the mixed-score attention and parallel encoding structures (see the table below), confirming their effectiveness, particularly for edge-only inputs. Our initial submission already included ablation results for three architectural components (see Appendix E), further demonstrating the value of self-attention and the overall structure.

Thank you for your constructive and valuable comments and suggestions. We address your concerns point-by-point as follows:

Q1: Given that REINFORCE can suffer from high variance, what specific techniques (e.g., baseline choices, larger batch sizes, other variance reduction methods) were employed to ensure stable and efficient training of UniteFormer?

A1: Thank you for your question regarding the stability of REINFORCE training in UniteFormer. To mitigate the known high variance issue of REINFORCE, we adopt a shared baseline strategy, as detailed in Appendix B, which is also widely used and validated in prior works such as POMO and MatNet.

(1) A shared baseline strategy: Specifically, we sample multiple trajectories per instance (e.g., via augmentation or diverse decoding), and use the average reward across these trajectories as the baseline. This shared baseline significantly reduces gradient variance and has been shown to improve training stability in VRP tasks. Empirical evidence from POMO shows that this shared baseline is more robust to local minima than the greedy rollout baseline in AM. This technique ensures that the advantage term $R(\tau_i)-b(s)$ has zero mean, and detailed formulas for the gradient variance can be found in Appendix B.

In addition, we have adopted Batched Trajectory Sampling and Large Mini-Batch Training techniques to ensure more stable and efficient training of UniteFormer.

(2) Batched Trajectory Sampling: Similar to POMO, we employ parallel exploration from diverse starting states, which leads to multiple diverse yet optimal solution representations. This practice effectively acts as a kind of entropy-regularized exploration, and avoids overfitting to narrow trajectory spaces, i.e., a common cause of instability in REINFORCE. Moreover, by leveraging the symmetries in routing problems (e.g., permutation invariance in TSP tours), we encourage the policy to generalize more robustly, thereby improving sample efficiency.

(3) Large Mini-Batch Training: As suggested in prior literature (e.g., AM and POMO), we train the policy network with larger mini-batches (e.g., 64–256 problem instances per batch), which smooths out noisy gradients and contributes to more stable convergence.

In summary, we combine a low-variance shared baseline, multiple trajectory rollouts, and large batch training to ensure that the REINFORCE training of UniteFormer remains stable and sample-efficient. These strategies have been shown to be effective in prior work (especially POMO) and are further validated in our empirical results across TSP, CVRP, and ATSP tasks.

We appreciate your feedback again and will include a more explicit discussion in Appendix B.

Q2: Were other policy gradient algorithms (like Actor-Critic variants or PPO) considered, and if so, what were the reasons for choosing REINFORCE over them?

A2: Thank you for your thoughtful question. In our work, we adopt the REINFORCE algorithm as the underlying policy gradient method, following a widely recognized and empirically validated precedent in neural combinatorial optimization, especially for routing problems. This choice is grounded in both empirical evidence and practical considerations, as discussed and validated in prior influential works such as AM and POMO. Particularly,

• AM explicitly states in the abstract:

  "We show how to train this model using REINFORCE with a simple baseline based on a deterministic greedy rollout, which we find is more efficient than using a value function."

• Similarly, POMO adopts REINFORCE and addresses variance through a shared baseline:

  "We optimize the model using the REINFORCE algorithm, with a shared baseline across the multiple trajectories per instance. This baseline significantly reduces the variance of the gradient and improves stability."

These works demonstrate that, when paired with appropriate variance reduction techniques (e.g., rollout or shared baselines), REINFORCE can train powerful and generalizable neural solvers. As a result, REINFORCE has become a de facto standard in the VRP literature.

In addition, we chose REINFORCE over more complex policy gradient algorithms such as Actor-Critic variants or PPO for several reasons:

(1) Simplicity and robustness in discrete, structured action spaces: Routing problems involve large, dynamic, and discrete action spaces (e.g., permutations). Designing a stable critic network that generalizes over such combinatorial structures is highly non-trivial and often leads to convergence issues.

(2) Compatibility with autoregressive decoding: Our decoder constructs solutions sequentially over permutations. REINFORCE allows us to compute unbiased policy gradients directly from sampled trajectories, without the need for approximating value functions over structured outputs.

(3) Lack of evidence for superior performance of Actor-Critic methods: To our knowledge, actor-critic or PPO-style algorithms have not shown consistent advantages over REINFORCE in VRP tasks. In fact, empirical attempts to integrate these methods (e.g., applying PPO to POMO according to our preliminary experiments) have often resulted in non-convergent training. This further explains why recent state-of-the-art models continue to rely on REINFORCE-based frameworks.

As detailed in Appendix B, our training setup with REINFORCE, which is combined with shared baselines, large batches, and architecture design, results in stable convergence and high-quality solutions across TSP, CVRP, and ATSP benchmarks.

We acknowledge that more advanced policy gradient algorithms could be interesting for future work, especially if tailored to handle structured action spaces with learned critics. However, for the scope of this paper, REINFORCE strikes the best balance between effectiveness, reproducibility, and compatibility with our architecture. So we find REINFORCE to be the most suitable and efficient choice for our setting.

Thank you for your constructive and valuable comments and suggestions. We address your concerns point-by-point as follows:

Q1: Notation Explanation.

A1: We appreciate your concern regarding the clarity of the mixed-score MHA layer and the associated notation in Appendix A. In the main text, the mixed-score MHA layer is briefly introduced by listing the required input arguments, specifically, $\text{mixed-scoreMHA}(h_0,h_v,D_{ij})$, to emphasize the functional interface rather than the internal details. A more detailed explanation of this mechanism has already been provided in Appendix A.1, particularly in Equation 1, where the general input embedding is denoted as $h_i$. In our implementation, $h_i$ corresponds to: a zero vector $h_0$, and a randomly selected one-hot vector $h_v$. These two are used to construct the inputs to the MHA layer as follows: $q = W^Q h_0,\quad k = W^K h_v,\quad v = W^V h_v.$ We opted for generic notation in the appendix to reflect the flexibility of $h_i$ as a general embedding, without tying it to specific choices like $h_0$ and $h_v$, which are task- and design-specific. That said, we agree that clarifying this correspondence more explicitly would improve readability, and we will incorporate this clarification in the final version.

Regarding your other points:

(1) The notation $\Theta_{ij}^{knn}$ refers to the edge adjacency matrix derived from the KNN graph structure. Although this has already been mentioned in line 139 of the main text, we acknowledge that it still lacks a detailed explanation. We will provide a more detailed description in the final version to improve clarity and completeness.

(2) The notation $\text{MHA}(h,H)$ is a shorthand representation of a generic attention layer, where the output from the previous layer (e.g., hidden states) serves as both the query and context input for the next layer. In contrast, $\text{MHA}(q,k,v)$ explicitly denotes the multi-head attention computation used in the decoder stage, where the query, key, and value inputs may come from different sources. We will clearly distinguish and define these usages in the final version to avoid ambiguity.

Q2: Model Size of UniteFormer.

A2: Thank you for your question. We provide a comparison of the model size (in terms of parameter count and storage footprint) between UniteFormer and other baseline neural solvers. Specifically, UniteFormer has approximately 4.48 million parameters, resulting in a model size of 47,287 KB, which is larger than POMO (as UniteFormer needs to deal with different types of input simultaneously), but smaller than MatNet.

This increase in size is due to UniteFormer’s unified architecture, which supports all three input modalities within a single model, without requiring separate training runs for each modality. In contrast, baseline methods typically require separate models to be trained independently for each input type. Moreover, our inference time remains competitive and is on par with most baseline methods, as shown in Table 1 of the main paper.

We will include a detailed table comparing parameter counts for all methods in the final version to improve transparency and reproducibility.

Q3: Training Settings of UniteFormer.

A3: Following the standard practice in prior works like POMO and MatNet, our UniteFormer trained on TSP-n is mainly evaluated on the same instance size TSP-n. However, we have also included the generalization test on TSPLib and CVRPLib (i.e., trained on 100-node instances but tested on various sizes) in Table 3 of the paper. To further demonstrate strong generalization to larger problem sizes, we apply models trained on TSP-100/CVRP-100 to solve TSP-200/CVRP-200 instances (see the table below, which provides the gap values under three input modalities). Here, UniteFormer still achieves competitive performance, highlighting its capacity to generalize beyond the training size.

Q4: Ablation Study.

A4: Thank you for your valuable suggestion. As part of our original submission, we have already conducted ablation studies on three input modalities and three major architectural components, as detailed in Section 4.2 and Appendix E. These experiments were selected to highlight the most representative design choices under the space constraints of the main paper.

Following your suggestion, we have further performed four additional ablation studies to assess the impact of the following components:

• Removing the mixed-score attention layer (denoted by w.o. mixeMHA),
• Changing the parallel-attention decoding to a single-layer decoding architecture (denoted by w.o. decoding),
• Removing the GCN layer (denoted by w.o. GCN), and
• Placing the self-attention layer before the GCN module (denoted by w.o. Order).

The experimental results are summarized in the table below and further confirm that each of these components contributes meaningfully to the model’s performance.

We appreciate your feedback and will include these new ablation studies in the final version to provide a more comprehensive analysis of the architecture.

Q5: Minor remarks [line 20]: Improvement solvers.

A5: Thank you for the valuable comment. You're right that there exists a broad class of machine learning–based methods that function as improvement solvers (e.g., learning local search heuristics or integrating with classical optimization frameworks), which are not necessarily autoregressive in nature.

Our original sentence was indeed narrow in scope. We have revised it as follows to better reflect the diversity of neural methods for VRPs:

These methods include both autoregressive models that learn construction policies from data[21, 9, 25], as well as learning-based improvement solvers that enhance classical optimization procedures. However, many of these models, particularly construction-based ones, make an overly simplifying assumption: they rely solely on either node coordinates or edge distances as input.

Q6: Minor remarks [line 41]: Complementary.

A6: Thank you for your insightful question. While node positions in Euclidean TSP do determine pairwise distances, we use the term "complementary structural information" to highlight that our model learns distinct yet mutually reinforcing aspects of the problem via separate encoders.

Specifically, the edge-aware sub-encoder directly models pairwise distances, which are critical for evaluating tour costs and capturing local geometric structure. In contrast, the node-focused sub-encoder uses attention to extract global, node-centric patterns, such as a node’s relative position in the graph (e.g., centrality, clustering). Although both sub-encoders ultimately rely on the same raw coordinates, they capture complementary views: one encodes pairwise relational structure, while the other emphasizes global positional context.

In addition, although distances can be derived from coordinates in Euclidean space, we find that neural networks benefit from directly learning the true distance values rather than relying on implicitly inferred ones. This direct modeling of edge-level information leads to more effective learning and improved solution quality.

Moreover, the notion of complementary becomes even more pronounced in non-Euclidean settings, such as the ATSP (the results are presented in Appendix F), where distances cannot be solely inferred from coordinates. This demonstrates that our model provides a flexible interface for directly handling real-world problems where true pairwise distances are explicitly defined and cannot be inferred from positions alone.

Q7: Minor remarks [line 115]: Figure.

A7: Thank you for your careful observation. The sentence in line 115-"Without edge weights, node features are passed to the GCNs directly" is specifically referring to the Node-only input modality illustrated in Figure 2(b). In this subfigure, the arrows clearly indicate the direct flow of node embeddings into the GCNs under the Node-only setting, and are intended to highlight the different input modes handled by our framework. In contrast, Figure 1 is designed to provide a high-level overview of the model architecture, emphasizing the modular composition and interactions between the main components. It does not explicitly depict the detailed flow of inputs for the three input modalities, which are instead elaborated in Figure 2. As image updates are not permitted during the rebuttal phase, we will revise Figure 1 in the final version to better reflect these input pathways and avoid confusion.

Thank you for your constructive and valuable comments and suggestions. We address your concerns point-by-point as follows:

Q1: Computational complexity analysis.

A1: Thank you for raising this important point. We acknowledge that our encoder architecture in UniteFormer might be heavier and increase computational demands, particularly for large-scale instances. This design was motivated by the need to learn rich structural representations across multiple input modalities, which has shown strong empirical benefits in terms of solution quality.

However, we would like to emphasize that UniteFormer is more practical in real-world scenarios due to its unified architecture: it can handle three different input modalities within a single model. In contrast, existing methods typically require separate training for each input setting, leading to higher cumulative training cost and reduced flexibility. UniteFormer's ability to generalize across modalities without retraining helps reduce overall training overhead and model maintenance cost, making it more efficient in deployment settings.

In terms of inference efficiency, our model achieves comparable inference time to baseline methods, as shown in Table 1 of the main paper. In addition, we also explicitly acknowledge the computational trade-off in the original paper. In Appendix G, we have highlighted that reducing the computational footprint of UniteFormer while preserving its expressive power is an important direction for future work. We are dedicated to the design of more lightweight yet expressive encoder modules to improve scalability for large instances.

We appreciate your suggestion and will incorporate a formal complexity analysis and efficiency discussion in the final version to make these trade-offs more transparent.

Q2: Problem Description and Formalization.

A2: Thank you for your valuable feedback. We have provided the problem descriptions and formalizations in the paper (see Appendix C). Due to space constraints in the main text, we have already included the detailed problem descriptions and formalizations for both the TSP and CVRP in Appendix C, and referenced them in the Problem Setting subsection of the experimental section (see Section 4).

We acknowledge that a more explicit formulation in the main paper would improve clarity, especially for readers less familiar with these problems. If space permits in the final version, we will include a concise but formal problem definition in the main text to ensure that the model's alignment with the problem requirements is clearly understood.

Q3: Hyperparameter Sensitivity.

A3: Thank you for your thoughtful suggestion. In this work, our primary focus is on analyzing the impact of the three input modalities and the architectural components of the proposed model. As such, we have conducted extensive ablation studies on both input types and architecture modules, which are presented in Section 4.2 and Appendix E.

For hyperparameter settings, we follow the configurations used in strong baselines such as POMO and MatNet to ensure fair comparison. Therefore, we did not include a detailed hyperparameter sensitivity analysis in the initial submission. However, we appreciate the importance of evaluating model robustness across different settings. In response to your comment, we conducted a sensitivity analysis on hyperparameter $K$, which controls the number of nearest neighbors used in the KNN graph for the GCN modules. We tested values of $K = 10, 20, 30, 40, 50$ on TSP50, and found that $K = 20$ yields the best performance, with relatively stable results across the range. The results are summarized in the table below.

We will include this analysis in the final version and additionally plan to expand the hyperparameter sensitivity study to better support practical deployment and reproducibility.

Q4: Reproducibility.

A4: Following common practice, we will release the full source code and all trained models upon acceptance of the paper. As external links are not permitted during the rebuttal phase, we promise to include links to the public code in the final version.