RoME: Domain-Robust Mixture-of-Experts for MILP Solution Prediction across Domains

Dear Reviewer i5gC,

Thank you for your insightful and valuable comments. We sincerely hope our rebuttal adequately address your concerns. If so, we would deeply appreciate it if you could consider raising your score. If not, please let us know your further concerns, and we will continue actively responding to your comments.

Q1: The denominator in equation (2).

What is the denominator in equation (2) summation over (x')? My concern is the feasible region of a MIP can have uncountably infinite solution.

Thank you for your feedback! We would like to clarify Equation (2) as follows.

• First, it's important to note that the denominator in Equation (2) does not represent an infinite sum of solutions. In our implementation, we utilize Gurobi to run training instances for 3,600 seconds, resulting in the 50 solutions with the highest objective values. The denominator, therefore, is a summation over these top 50 solutions.
• More specifically, following the previous works [2-3], RoME also aims to approximate high-quality solution distributions. This means we focus solely on solutions with favorable objective values during training. Consequently, we use the top 50 solutions to estimate these high-quality distributions. As you noted, the total solution space can be infinite, making a summation over the entire space impractical.

Q2: The clarification of our methods on general MILPs.

The paper seems to claim its method works on general MIP at first, but starting from line 123, it started to assume that MIP only contains binary variables. Can you clarify this point?

Thank you for your insightful comments! To address your concerns, we have conducted experiments on general MILPs that include integer, binary, and continuous variables.

• Our method is applicable to general MILPs. In Section 6.1.3 of [1], the authors indicate that the model can be generalized to non-binary variables with certain modifications. Following the guidance in [1], we have adapted our model to accommodate general MILPs.
• We performed experiments on four challenging instances from MIPLIB, which are general MILP problems featuring integer variables. The results, presented in Table 1, demonstrate our method's superior performance compared to the baseline approaches. As in the main-text experiments, we first select a small set of similar instances (based on the instance-similarity metrics) and use them to train the PS model, which is then evaluated on the testing instances. By contrast, our method is applied to the same testing instances in zero-shot mode, requiring no additional training.

Table 1: Experiments on the general MILP instances.

Q3: Explanation on the "fully factorized solution" assumption.

What is your intuition for the "fully factorized solution" assumption in line 124? Besides the part that it makes it easy to design a loss function, is there any mathematical justification for that? Ignoring the part the definition of p(x|I) is not yet well-defined, it is doubtful that this assumption can be true, as given an instance I, the value of x_1 can change the feasible region of x_2. In addition, using this assumption, are we ignoring the continuous part of the solution? it seems like even if the continuous part give bad objective, we still assign high probability.

Thanks for your question! We would like to address your points regarding the definition of our probability distribution and the justification for the independence assumption.

• Clarification of the definition of the probability distribution $p(x|\mathcal{I})$. This is not a theoretical distribution over all possible solutions, but rather an empirical probability derived from a collected set of high-quality feasible solutions for a given instance $\mathcal{I}$. To generate the training data for an instance, we first solve it to collect a set of top-performing solutions (e.g., the best 50 solutions found within a time limit). The probability for a single binary variable, $p_\theta(x_i=1|\mathcal{I})$, is then computed as the frequency of that variable taking the value 1 across this set of elite solutions. In essence, it is the average value of that variable in the collected solution set, weighted by solution quality.
• The Assumption is widely adopted in the previous works. We acknowledge that variable independence is a strong assumption that may not hold true in practice. However, this is a widely adopted simplification in the field for reasons of computational tractability, as seen in related works [1-3]. The core intuition is that if a variable takes on a consistent value (e.g., 1) across the vast majority of top-tier solutions, then its marginal probability will be high.
• The variable dependencies are captured to a degree by the GNN encoder. The GNN processes the entire MILP instance as a single graph, and through its multi-layer message-passing mechanism, the final embedding for each variable is informed by the global problem structure, including its relationships with other variables and constraints. Therefore, the complex relationships between variables are considered during the feature extraction stage, even if they are not explicitly modeled in the final probabilistic output.
• Explanation on the continuous variables. The computational complexity of solving MILPs stems primarily from the combinatorial nature of the integer and binary variables. By predicting the values for a significant portion of the integer and binary variables, we effectively prune the search tree and guide the solver to a much smaller, more promising region of the solution space. Finding the optimal values for the remaining variables and continuous variables becomes significantly more tractable for the solver.

Q4 Typos

The definition of x^(i)_j in line 129 is a bit confusing. Why there is no j in the right-hand side?

Thank you for pointing this typo. We have corrected it to $x_j^{(i)} = p_{\theta}(x_j^{(i)}=1|\mathcal{I})$.

Q5 Sensitivity of $k_0$, $k_1$ and $\Delta$

In Append A.3 , It is said the the selection of k_0, k_1 and Delta have great influences. Given that in Table 3, different methods use different configuration of these values. Would you be able to provide tables/plots showing the impact of these values when varying?

Thank you for the insightful comments. To address them, we have conducted an additional sensitivity study on the CA benchmark. The results are summarized in Table 2. Please note that PS is trained directly on the CA dataset, whereas RoME has never seen CA during training. Despite this disadvantage, RoME still outperforms PS across all parameter settings.

Table 2: Sensitivity of hyperparameters.

Q6 Explanation on the Figure 3

I do not understand why, in Figure 3, the lines plot do not start second 0, but rather at the 10th second (4 out of 5)? Is it because the solver takes a while to find its first feasible solution or something else?

Thanks for your feedback! We observed that our approach usually finds the first feasible solution at around 10th second—except on the IS benchmark, which is relatively easy, so all methods locate a feasible point almost instantly. This observation matches prior works [2-4], which also plot results starting from the 10th second.

Q7 Quality of RoME

I am curious about the quality of the binary solution returned by RoME. I think it give us more insight if you can show the objective of the returned solution of RoME or how many of it was feasible (for all instances in your experiments). For the non-pure binary case, we can assign value for the continuous variables by the solving the LP when fixed all binary variables at the predicted solution.

Thank you for the comments! Methods in the Predict-and-Search family [2-3] learn a distribution of high-quality solutions and then obtain a partial assignment of variables. Following your suggestion, we fixed these predicted variables, solved the remaining problem, and reported the resulting objectives for PS and RoME on the SC and WA benchmarks in Table 3.

Table 3: Results of directly fixing variables.

[1] Solving Mixed Integer Programs Using Neural Networks.

[2] A GNN-Guided Predict-and-Search Framework for Mixed-Integer Linear Programming. ICLR 2023.

[3] Contrastive Predict-and-Search for Mixed Integer Linear Programs. ICML 2024.

[4] Differentiable Integer Linear Programming. ICLR 2025.

Dear Reviewer i5gc,

We sincerely thank you for your time and efforts during the rebuttal process. We are writing to gently remind you that the author-reviewer discussion period will end in less than 10 hours. We have responded to your further comments and eagerly await your feedback, and we sincerely hope that our response has properly addressed your concerns. We would deeply appreciate it if you could kindly point out your further concerns so that we could keep improving our work. We sincerely thank you once more for your insightful comments and kind support.

Best,

Authors

Dear Reviewer 4Hsn,

Thank you for your insightful and valuable comments. We sincerely hope our rebuttal adequately address your concerns. If so, we would deeply appreciate it if you could consider raising your score. If not, please let us know your further concerns, and we will continue actively responding to your comments.

W1 & Q1 Need to quantify performance drop vs. Gurobi

Although the model performs well on average, for some problem, a performance drop can be observed compared to Gurobi. This performance drop should be quantified and reported.

Thank you for highlighting the importance of quantifying the performance drop against Gurobi. Following standard practice in multi-task learning, we introduce a “win” metric—i.e., the number of test instances (out of 100) on which each method attains the best objective. First, we compare Gurobi , PS, and RoME together on the CA and SC benchmarks, as shown in Table 1.

Table 1: The total win counts of each method.

RoME dominates on both benchmarks. We then separate the instances on which each learning-enhanced solver wins or loses against Gurobi and report the mean absolute objective gap ($\text{Gap}_{\text{abs}}$) in Table 2.

Table 2: The win/lose gap against Gurobi of each method.

RoME not only wins far more often but also gains more when it wins and loses less when it falls short. For a concrete view, the Table 3 lists the objective values on the first 20 CA instances:

Table 3: The objective of partial instances obtained by each method.

W2 & Q2 Effect of expert count on scalability unexplored

The effect of the number of experts in the MoE architecture is not explored, which is essential to understand the scalability and architectural sensitivity of the proposed approach.

Thank you for stressing this point. As Fig. 1(b) in the main text shows, scaling the number of experts together with the number of training tasks consistently improves performance. We further evaluate scalability by varying the expert count on four benchmarks (For each benchmark, we select 20 instances) in Table 4:

Table 4: Effect of expert count on RoME performance.

W3 Typos

The paper includes minor inconsistencies and typos (e.g., “T-SNE” should be “t-SNE” with lowercase “t”).

Thank you. We have corrected these typos in the manuscript.

Dear Reviewer SKRg,

Thank you for your insightful and valuable comments. We sincerely hope our rebuttal adequately address your concerns. If so, we would deeply appreciate it if you could consider raising your score. If not, please let us know your further concerns, and we will continue actively responding to your comments.

W1 & Q1 Rationale of fuzzy consistency loss remains unclear

The reason why the fuzzy consistency loss works seems unclear. In the high-dimensional task embedding space, why can forcing the original and perturbed outputs to be the same help generalize? Do you have more insights into this part?

Thank you for raising this point! To clarify the effect of the fuzzy-consistency loss, we have conducted an ablation study in which we removed this loss—that is, we no longer require the outputs of the original and perturbed task embeddings to be the same. We evaluate this change on the Set Cover benchmark and on the MIPLIB instance bab2 (a real-world instance which also contains the set_covering constraint type). The results are shown Table 1.

Table 1: Ablation study of fuzzy consistency loss.

The improvement confirms the effectiveness of the fuzzy-consistency loss. We attribute the gains to three factors:

• More accurate routing. A task embedding represents the structural pattern of a MILP instance and is directly consumed by the MoE gate. The loss encourages similar embeddings to select the same expert subset. In Section 4.4 we observed that both SC and bab2 activate Expert #3. When the loss is removed, SC still selects Expert #3, but bab2 mistakenly switches to Expert #2, degrading performance.
• Alignment with cross-domain reducibility. From a computational-complexity viewpoint, most MILP classes are NP-complete, implying that different MILPs can, in principle, be reduced to one another—a fact confirmed by recent advances in graph-based combinatorial optimization [1]. Building on this insight, we exploit structural similarities (e.g., the IS and CA domains both have the set_packing constraints) so that related domains route to the same expert subset, enhancing generalization. Perturbing task embeddings partially simulates cross-domain shifts; forcing the perturbed and original embeddings to yield identical outputs tightens this alignment and further boosts generalization.
• Enhanced robustness. From an adversarial-training viewpoint, embedding perturbations stabilize the training of DRO, helping RoME capture high-level patterns that persist across domains.

[1] UniCO: On Unified Combinatorial Optimization via Problem Reduction to Matrix-Encoded General TSP. ICLR 2025.

Q2 Gap between benchmark and real-world problem sizes

In terms of the testing problem sizes, what's the gap between them and the real-world problem size?

Thank you for your feedback! Appendix C lists the sizes of all test instances, covering both our synthetic benchmarks and the real-world MIPLIB benchmarks. MIPLIB instances originate from industrial production, transportation, and other operational settings. And their variable and constraint counts can reach hundreds of thousands. Here we select some of them which are reported in Table 2, and the complete information is detailed in Appendix C.

Table 2: Statistical information of some partial benchmarks.

Dear Reviewer SAaz,

Thank you for your kind support and valuable comments. We sincerely hope our rebuttal adequately address your concerns. If so, we would deeply appreciate it if you could consider raising your score. If not, please let us know your further concerns, and we will continue actively responding to your comments.

W1 Model interpretability & performance

The complexity of the model pose challenges for interpretability.

We appreciate this concern. Model complexity is indeed a common challenge for most learning algorithms. Below we address your concerns in detail.

1.RoME is designed as a pretrained model for MILP solution prediction, not merely a task-specific predictor.

Unlike prior works such as PS [1] and ConPS [2], which are tailored to a single domain, RoME is a pretrained model and targets multiple domains simultaneously. Although this slightly increases model complexity, it also delivers substantial performance gains across domains. We analyze the reasons for the success of RoME below.

• Scalability is essential for a cross-domain MILP model. Inspired by progress in large language models, computer vision, and drug discovery, we adopt a MoE architecture. The GNN encoder is identical to earlier work [1-4] and remains the primary computational cost. We add multiple lightweight MLP-based experts and decoder heads, so these introduce minimal inference overhead. Here we provide an intuitive comparison between prior work and our approach. RoME has 697.89K parameters (storage 2.66MB), compared to 449.05K parameters (storage 1.71MB) in PS [1], showing only a slight increase in model size. Meanwhile, our approach achieves substantial gains in cross-domain generalization ability. In addition, as Fig. 1(b) in the main text shows, scaling the number of experts together with the number of training tasks consistently improves performance.
• RoME meets the three pillars of a foundation model. In recent years, foundation models have emerged across a wide range of fields. Empirically, their success relies on three core components: a robust training approach, a scalable architecture, and access to large-scale data. (i) Training approach: we use a two-level DRO objective to balance inter- and intra-domain robustness. (ii) Scalable architecture: the MoE design enables the model to perceive and differentiate diverse domains. (iii) Large training data: RoME holds the ability to scale along with the large data. Our work represents the first step toward a foundation model in the MILP domain. While still in its early stages, we hope it lays the groundwork for future development in this direction.
• Lower overall cost compared with single-domain baselines. Once pre-trained, RoME can be reused for any target domain without retraining, whereas previous methods must be re-trained from scratch for each new domain and degrade sharply outside it. In practice, this makes RoME more economical despite its larger architecture.

2.RoME learns a deeper understanding in the routing mechanism.

Current works typically adopt GNN-based models, which are inherently treated as black-boxes. However, with the recent progress in large language models, an increasing number of interpretability techniques have emerged, particularly for MoE architectures, offering new opportunities to better understand model behavior. Following this, we have presented an interpretability study in Section 4.4 of the main text. Below we give a further analysis.

• Expert-activation patterns across tasks. We examine which experts are activated for different tasks and find that tasks with similar structure consistently select the same subset of experts. This indicates that each expert has learned intrinsic cross-domain patterns rather than memorizing single tasks.
• Latent-space visualization of task embeddings. We project the learned task embeddings into 2-D space (T-SNE). Structurally similar tasks cluster tightly together, confirming that RoME captures correct task representations. This, in turn, shows that the routing mechanism maps tasks to experts in a coherent, task-aware manner.

The model performance heavily depends on routing & task embedding quality.

We acknowledge that routing precision and task embedding quality are the core targets of our intra-domain DRO strategy, which acts as the bridge between training and architecture. By perturbing task embeddings during training, intra-DRO properly activates the routing gate, ensuring that only the appropriate experts are activated—even under small representation shifts. This is pivotal for RoME’s cross-domain generalization; its effectiveness stems from the joint design of our DRO training strategy and MoE architecture, as detailed in Appendix D.4.

W2 & W3.1 Data efficiency

RoME requires more data for training.

In fact, in data-limited scenarios, our approach demonstrates stronger performance. This is because RoME is a pretrained model that can be trained on diverse synthetic datasets to learn transferable, cross-domain patterns. In contrast, prior methods are trained on task-specific data and can only be applied to the corresponding target domain. Unfortunately, collecting sufficient training data for each new domain is often impractical. Additionally, to further demonstrate the scalability of our approach, and in line with your suggestion, we have conducted the following analysis. We trained RoME and the PS baseline using 30 %, 50 %, and 80 % of the original training data, then evaluated both models on the widely adopted Set Cover benchmark. The results (see the Table 1) show that RoME consistently surpasses PS at every data budget, and the performance gap widens as the data volume increases.

Table 1: Performance of each method with the limited data.

In Section 4.3 and Appendix D.1, RoME takes a large step forward by delivering zero-shot performance on the challenging industrial benchmark MIPLIB, whereas previous methods first have to build and train on a dedicated MIPLIB dataset—often infeasible in practice. To mirror real-world constraints more closely, we follow [3]: (i) we test RoME in a data-scarce setting where only 80% of the original IIS training data are available, and (ii) we fine-tune RoME on the training data constructed in [3] while freezing the GNN encoder and finetuning only the expert and head networks. With fewer than 10 steps, RoME already surpasses training-specific baseline such as PS. The complete results are reported in Table 2.

Table 2: Performance of finetuning on the IIS datasets.

W3.2 & W4 More complicated problem structure

RoME should explore more complicated problem structure during training.

We sincerely thank you for your insightful suggestions! As illustrated in Figure 1(a) of the main text, the key difference between prior works and RoME lies in the training domains. Compared to previous methods [1–4], our approach takes a significant step forward by being trained on diverse domains, each with distinct problem structures, whereas prior works are trained on a single fixed domain. For example, in our main experiments, RoME is trained on three different MILP families, each representing a unique problem structure. During training, we randomly sample instances from these families, meaning that the problem structure dynamically varies across training batches. In contrast, prior methods are trained on data from a single domain, where the structural pattern remains constant. Moreover, to evaluate cross-domain generalization, we not only test RoME on unseen synthetic domains, but also evaluate it on the real-world MIPLIB benchmark, which includes a wide variety of industrial instances spanning domains such as manufacturing, energy, and transportation. To give a clearer picture of the structural diversity, we list the constraint types of a subset of MIPLIB instances used in our experiments, as shown in Table 3.

Table 3: Constraints type of different benchmarks.

[1] A gnn-guided predict-and-search framework for mixed-integer linear programming. ICLR 2023.

[2] Contrastive Predict-and-Search for Mixed Integer Linear Programs. ICML 2024.

[3] Apollo-MILP: An Alternating Prediction-Correction Neural Solving Framework for Mixed-Integer Linear Programming. ICLR 2025.

[4] Differentiable Integer Linear Programming. ICLR 2025.