🔢 mathematics

Learning to Optimize at Scale: A Benders Decomposition-TransfORmers Framework for Stochastic Combinatorial Optimization

This paper proposes a learning-augmented Benders decomposition framework that leverages a pre-trained TransfORmer model to rapidly generate high-quality approximate solutions for scenario subproblems, enabling the efficient solution of large-scale two-stage stochastic capacitated lot-sizing problems with arbitrary time horizons while maintaining zero infeasibility.

Original authors: Seung Jin Choi, Kimiya Jozani, Josh Cooper, Esra Buyuktahtakin Toy

Published 2026-07-28

📖 4 min read🧠 Deep dive

CC BY 4.0

Original authors: Seung Jin Choi, Kimiya Jozani, Josh Cooper, Esra Buyuktahtakin Toy

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are the captain of a massive cargo fleet, trying to decide exactly when and where to load ships to meet customer orders. The catch? You don't know exactly how many customers will show up, or how much cargo they'll need, until the ships are already sailing. This is the heart of a field called stochastic optimization: the science of making the best possible plans when the future is foggy and full of surprises. In the real world, this isn't just about ships; it's about factories deciding how much to produce, power grids balancing energy, and hospitals managing supplies. The problem is that as the number of possibilities grows, the math required to find the perfect plan becomes so huge that even the world's fastest supercomputers can get stuck, like a car trying to drive through a traffic jam that never ends.

To solve these massive puzzles, mathematicians have long used a clever trick called Benders decomposition. Think of this as a team of detectives working on a giant mystery. Instead of one detective trying to solve the whole case at once, they split the work. One detective (the "Master") makes the big, long-term decisions, like "Should we open a factory?" Then, a team of specialists (the "Subproblems") checks if those decisions actually work for every possible future scenario, like "What if it rains?" or "What if demand spikes?" They send feedback notes back to the Master to refine the plan. This works great for small mysteries, but when the case gets huge, the specialists spend so much time checking every tiny detail that the Master never gets a chance to make a final decision.

This is where a new paper by Seung Jin Choi and colleagues from Virginia Tech steps in with a fresh idea. They asked: What if we could give those specialists a superpower? Instead of spending hours calculating every single possibility, what if we could train a smart computer brain—a Transformer (the same type of AI that powers many modern chatbots and translation tools)—to instantly guess the best moves? The authors propose a hybrid framework they call ML-Benders. In their system, the AI acts as a rapid-fire surrogate, quickly predicting high-quality solutions for the complex "what-if" scenarios. It doesn't replace the math entirely; rather, it acts as a turbocharger, generating strong hints (called "cuts") that guide the Master detective much faster toward the right answer.

The team tested this on a classic production planning problem called the Two-Stage Stochastic Capacitated Lot-Sizing Problem (TSSCLSP). They trained their AI model on relatively short planning horizons, specifically looking at 90 time periods (like 90 days). The real magic, however, happened when they asked the model to solve problems three times larger, stretching out to 270 time periods, without ever seeing a problem that big during its training. This is like teaching a student to solve a 10-page math test and then handing them a 30-page test, expecting them to figure it out using the same logic.

The results were impressive. When the AI was tested on its home turf (the 90-period problems), it slashed the time needed to find a solution by nearly 20% and reduced the error gap by a massive 91.5% compared to the old, slow method. But the most exciting finding was its ability to scale. Even when faced with the giant 270-period problems, the system successfully generated valid, workable plans for every scenario without getting stuck or producing impossible results. While the final plans for these giant problems weren't perfect (leaving a gap of about 19.60% compared to a theoretical perfect solution), the fact that the system could solve them at all is a big deal. In the past, problems of this size were considered too difficult to tackle with this specific approach.

The paper highlights a specific technique called "expandable generation," which works like a sliding window. Imagine the AI is reading a long story; it reads the first chapter, then uses the ending of that chapter as context to predict the next chapter, and so on, sliding forward until the whole story is written. This allowed a model trained on short stories to write long novels. The authors emphasize that this doesn't mean the AI is perfect; in the giant 270-period tests, the solutions were good enough to be feasible but still had room for improvement. However, the study proves that combining the rigorous logic of classical math with the speed of modern AI can unlock solutions for problems that were previously too big to handle, offering a promising new path for solving complex, real-world planning challenges.

Technical Summary: Learning to Optimize at Scale via Benders-Decomposition-TransfORmers

Problem Definition

The paper addresses the Two-Stage Stochastic Capacitated Lot-Sizing Problem (TSSCLSP), a classic NP-hard optimization problem involving production planning under uncertain demand and capacity constraints. The objective is to minimize total costs, comprising first-stage fixed setup costs and the expected second-stage production and inventory holding costs across multiple demand scenarios.

Traditional solution methods, including stochastic dynamic programming, Lagrangian relaxation, and classical Benders decomposition, struggle with scalability as the problem size increases (e.g., large time horizons and numerous scenarios). The primary bottleneck in classical Benders decomposition is the iterative, computationally expensive solving of numerous scenario-specific subproblems to generate optimality and feasibility cuts.

Methodology

The authors propose ML-Benders, a learning-augmented framework that integrates a pre-trained TransfORmer model into the Benders decomposition loop to accelerate convergence.

1. Hybrid Decomposition Framework

The problem is decomposed into:

Master Problem (MP): Determines first-stage binary setup decisions ( $Y_t$ ) and approximates the expected recourse cost.
Subproblems (SP): For a fixed master solution, independent Linear Programs (LPs) are solved for each scenario to evaluate recourse costs and generate dual variables.

2. TransfORmer as a Surrogate Solver

Instead of solving every SP exactly during the Benders iterations, the framework employs a TransfORmer model to rapidly predict high-quality approximate solutions for the binary setup variables ( $Y_t$ ) within each scenario.

Workflow: The MP solution is encoded as a masked input. The TransfORmer performs a single forward pass to predict the full horizon of $Y_t$ variables.
Cut Generation: Once the TransfORmer predicts the $Y_t$ values, the subproblem reduces to a fast LP (with fixed binaries). This LP is solved exactly to obtain the precise costs and dual variables ( $\pi, \theta$ ) required to construct valid Benders cuts.
Architecture: The model utilizes an encoder-decoder structure with self-attention and cross-attention mechanisms. It includes specific "heads" for setup, production, and inventory, with piecewise transforms for continuous variables.

3. Expandable Generation Mechanism

A core technical innovation is a sliding-window generation technique that allows a model trained on a fixed horizon (e.g., $T=90$ ) to solve instances of arbitrary, longer lengths (e.g., $T=270$ ).

Process: The model generates a solution for the first time window. For subsequent windows, the input concatenates the last portion of the previous output (as historical context) with known future master variables and masked tokens for the unknown future.
Robustness: To ensure feasibility, the system employs 5-shot sampling per subproblem, selecting the best feasible prediction to construct cuts.

Key Contributions

Accelerated Convergence: The framework replaces the exact solving of subproblems with a TransfORmer-based surrogate, significantly reducing the time required to generate strong optimality and feasibility cuts.
Scalability via Expandable Generation: The paper demonstrates a novel mechanism where a model trained on short horizons ( $T=90$ ) can generalize to solve long-horizon instances ( $T=270$ ) without retraining, effectively addressing the "zero-shot" generalization challenge in stochastic optimization.
Zero Infeasibility: Despite using a learning-based surrogate for the initial solution, the framework maintains zero infeasibility in the generated subproblem solutions because the final cuts are derived from exact LP solves based on the TransfORmer's predictions.
High-Quality Cuts: The TransfORmer-generated cuts are shown to be dense and highly informative, guiding the master problem more effectively than classical approaches.

Experimental Results

The framework was evaluated on TSSCLSP instances using benchmarks from Atamtürk and Munoz [2004] and Yilmaz and Büyüktahtakın [2023].

Performance at Training Horizon ( $T=90$ )

Compared to classical Benders decomposition under a 300-second time limit:

Solution Quality: ML-Benders achieved an average optimality gap of 3.55%, a 91.5% reduction compared to the 41.81% gap of classical Benders.
Efficiency: The method required 29.4% fewer iterations and 19.9% less time to reach convergence.
Cut Budget Sensitivity: Under a strict limit of 50 cuts, classical Benders failed to converge (77.03% gap), whereas ML-Benders achieved a 2.80% gap, demonstrating the superior quality of the generated constraints.

Scalability to Long Horizons ( $T=270$ )

The model, trained exclusively on $T=90$ instances, was tested on $T=270$ instances:

Generalization: The framework successfully generated complete, feasible solutions for all scenarios in the long-horizon setting.
Trade-offs: The average optimality gap increased to 19.60% (compared to 3.55% at $T=90$ ), and iterations increased from 51.6 to 66.9. However, the ability to solve instances of this scale without subproblem infeasibility validates the robustness of the expandable mechanism.

Significance and Claims

The paper positions this work as a significant step toward bridging the gap between deep learning and classical mathematical optimization.

Trust in ML-Augmented Solvers: By ensuring that the final cuts are mathematically valid (derived from exact LPs) and that no infeasible subproblems are generated, the framework addresses reliability concerns in using ML for optimization.
Generalizability: The authors claim that the expandable generation mechanism proves that policies learned on small, manageable problem sizes can be generalized to solve much larger, industrially relevant instances that were previously intractable for this specific decomposition approach.
Framework Versatility: While focused on TSSCLSP, the authors assert the approach is general and adaptable to a wide range of dynamic two-stage stochastic optimization problems.

The work is presented as a non-archival workshop paper at NeurIPS 2025, highlighting the potential of TransfORmers as powerful surrogate solvers embedded within classical decomposition algorithms.