[Paper Review] SurCo: Learning Linear Surrogates For Combinatorial Nonlinear Optimization Problems

Aaron Ferber, Taoan Huang+5 more|arXiv (Cornell University)|Oct 22, 2022

Metaheuristic Optimization Algorithms Research4 citations

TL;DR

SurCo proposes learning linear surrogate costs for combinatorial nonlinear optimization problems by differentiating through a linear solver to enable efficient, gradient-based optimization. It achieves better solutions faster than state-of-the-art methods in real-world applications like embedding sharding, photonic design, and route planning.

ABSTRACT

Optimization problems with nonlinear cost functions and combinatorial constraints appear in many real-world applications but remain challenging to solve efficiently compared to their linear counterparts. To bridge this gap, we propose $ extbf{SurCo}$ that learns linear $\underline{ ext{Sur}}$rogate costs which can be used in existing $\underline{ ext{Co}}$mbinatorial solvers to output good solutions to the original nonlinear combinatorial optimization problem. The surrogate costs are learned end-to-end with nonlinear loss by differentiating through the linear surrogate solver, combining the flexibility of gradient-based methods with the structure of linear combinatorial optimization. We propose three $ exttt{SurCo}$ variants: $ exttt{SurCo}- exttt{zero}$ for individual nonlinear problems, $ exttt{SurCo}- exttt{prior}$ for problem distributions, and $ exttt{SurCo}- exttt{hybrid}$ to combine both distribution and problem-specific information. We give theoretical intuition motivating $ exttt{SurCo}$, and evaluate it empirically. Experiments show that $ exttt{SurCo}$ finds better solutions faster than state-of-the-art and domain expert approaches in real-world optimization problems such as embedding table sharding, inverse photonic design, and nonlinear route planning.

Motivation & Objective

• Address the challenge of solving combinatorial optimization problems with nonlinear, non-convex cost functions using efficient, off-the-shelf solvers.
• Overcome limitations of black-box optimization and heuristic methods that are slow or fail to ensure combinatorial feasibility.
• Bridge the gap between nonlinear objectives and linear solvers by learning surrogate linear costs that preserve problem structure.
• Enable generalization across problem instances by training models to predict surrogate coefficients from instance descriptions.
• Improve solution quality and runtime efficiency in real-world domains such as GPU sharding, photonic design, and stochastic routing.

Proposed method

• Propose SurCo-zero: end-to-end training of linear surrogate costs via backpropagation through a linear combinatorial solver to minimize nonlinear loss.
• Introduce SurCo-prior: train a model to predict surrogate coefficients from instance-specific features, enabling fast inference on unseen problems.
• Develop SurCo-hybrid: fine-tune SurCo-prior predictions using SurCo-zero on individual test instances to combine distributional and instance-specific knowledge.
• Leverage differentiable solvers to compute gradients of the nonlinear objective through the surrogate solution, enabling gradient-based optimization of surrogate coefficients.
• Use a surrogate optimization problem: min_x∈Ω ĉᵀx, where ĉ is predicted by a neural network and x is combinatorially feasible.
• Train using nonlinear, differentiable loss functions that reflect the true objective, allowing the model to learn effective linear approximations.

Experimental results

Research questions

• RQ1 Can learning linear surrogate costs enable efficient and high-quality solutions to nonlinear combinatorial optimization problems?
• RQ2 Does differentiating through a linear solver improve sample efficiency and solution quality compared to black-box or heuristic methods?
• RQ3 Can a model trained on a distribution of problems generalize to unseen instances while maintaining high solution quality?
• RQ4 How does combining distribution-level priors with instance-specific fine-tuning improve performance in nonlinear optimization?
• RQ5 Can SurCo outperform state-of-the-art methods in real-world applications with complex, nonlinear objectives and combinatorial constraints?

Key findings

• SurCo-zero achieves better solution quality than state-of-the-art methods in embedding table sharding with comparable or lower runtime.
• SurCo-prior outperforms reinforcement learning and other learning-based methods in solving held-out problems from the same distribution.
• SurCo-hybrid improves upon SurCo-prior by fine-tuning on individual instances, yielding higher-quality solutions in nonlinear route planning and photonic design.
• The method finds viable photonic devices more quickly than baseline approaches, demonstrating practical utility in inverse design.
• SurCo generalizes well across diverse real-world problems, including industrial-scale applications like GPU sharding and stochastic path planning.
• The framework demonstrates that learning linear surrogates with differentiable solvers enables scalable, high-performance optimization of nonlinear combinatorial problems.