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

TL;DR

RoME tackles cross-domain generalization in MILP solution prediction by learning a single model that adapts to diverse problem distributions through a domain-robust Mixture-of-Experts (MoE) and a distributionally robust training objective. The MoE comprises a shared graph encoder, several expert networks, and a multi-head task decoder that routes each instance via learned task embeddings, producing $p$-dimensional marginals for the binary variables. The training objective combines inter-domain group-DRO and intra-domain embedding perturbations, plus regularizers for expert diversity and robust routing, yielding strong cross-domain and zero-shot performance. Empirically, a RoME model trained on three domains achieves a $67.7\%$ average improvement across five diverse domains and shows measurable gains on MIPLIB in a zero-shot setting.

Abstract

Mixed-Integer Linear Programming (MILP) is a fundamental and powerful framework for modeling complex optimization problems across diverse domains. Recently, learning-based methods have shown great promise in accelerating MILP solvers by predicting high-quality solutions. However, most existing approaches are developed and evaluated in single-domain settings, limiting their ability to generalize to unseen problem distributions. This limitation poses a major obstacle to building scalable and general-purpose learning-based solvers. To address this challenge, we introduce RoME, a domain-Robust Mixture-of-Experts framework for predicting MILP solutions across domains. RoME dynamically routes problem instances to specialized experts based on learned task embeddings. The model is trained using a two-level distributionally robust optimization strategy: inter-domain to mitigate global shifts across domains, and intra-domain to enhance local robustness by introducing perturbations on task embeddings. We reveal that cross-domain training not only enhances the model's generalization capability to unseen domains but also improves performance within each individual domain by encouraging the model to capture more general intrinsic combinatorial patterns. Specifically, a single RoME model trained on three domains achieves an average improvement of 67.7% then evaluated on five diverse domains. We further test the pretrained model on MIPLIB in a zero-shot setting, demonstrating its ability to deliver measurable performance gains on challenging real-world instances where existing learning-based approaches often struggle to generalize.

Figures (6)

• Figure 1: Illustration of cross-domain training. (a) A single model is trained on a collection of MILP domains with varying constraint types miplib2021, such as PAC (set packing), KPS (knapsack), COV (set covering), and PAR (set partitioning), and evaluated on unseen domains such as MIPLIB. (b) We conduct experiments across five distinct domains, reporting the average objective gap in percentage relative to the best-known solution while varying the number of experts and training domains. Each value in the heatmap denotes the average percentage gap to the best-known solution. The results show that training on more diverse domains and selecting an appropriate number of experts can effectively reduce this objective gap.
• Figure 2: The overview of RoME. RoME employ a MoE architecture with structure-aware representation learning. For cross-domain training, RoME proposes a robust training object containing binary cross-entropy loss $\mathcal{L}_{\text{BCE}}$, expert diversity regularization $\mathcal{L}_{\text{Div}}$ and fuzzy membership consistency loss $\mathcal{L}_{\text{Robust}}$. To further enhance generalization across structurally diverse MILP families, RoME uses a group-level domain robust scheme to construct DRO loss.
• Figure 3: The average primal gap of different methods over 100 instances as the solving process proceeds. We use Gurobi for implementation and set the time limit to be 1,000 seconds.
• Figure 4: The primal gap of different methods on five easier instances from IIS dataset as the solving process proceeds. We use Gurobi for implementation and set the time limit to be 1,000 seconds.
• Figure 5: The primal gap of different methods on five harder instances from MIPLIB as the solving process proceeds. We use Gurobi for implementation and set the time limit to 1,000s.