🤖 AI

Forge: Quality-Aware Reinforcement Learning for NP-Hard Optimization in LLMs

The paper introduces OPT-BENCH, a comprehensive framework utilizing quality-aware Reinforcement Learning with Verifiable Rewards (RLVR) to train Large Language Models on NP-hard optimization problems, demonstrating significant improvements in solution quality and generalization across diverse reasoning tasks compared to existing models and binary reward approaches.

Original authors: Xiaozhe Li, Xinyu Fang, Shengyuan Ding, Yang Li, Linyang Li, Haodong Duan, Qingwen Liu, Kai Chen

Published 2026-05-12

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Xiaozhe Li, Xinyu Fang, Shengyuan Ding, Yang Li, Linyang Li, Haodong Duan, Qingwen Liu, Kai Chen

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 have a very smart robot assistant (a Large Language Model, or LLM) that is great at answering questions like "Is this math problem solved correctly?" or "Did I spell this word right?" For a long time, we've trained these robots by saying "Good job!" if the answer is right and "Try again" if it's wrong. This is like grading a multiple-choice test where there's only one right answer.

But what if the task isn't just about getting the answer right, but getting the answer best?

This is the problem the paper FORGE tackles. It's like asking the robot not just to find a route to the grocery store, but to find the shortest, fastest route possible, even if there are millions of other valid routes that just take longer.

Here is a simple breakdown of how they did it:

1. The Problem: "Good Enough" vs. "The Best"

Imagine you are packing a suitcase.

Old Way (Binary Rewards): You ask the robot to pack. If it fits everything in the bag, you say "Success!" If it spills over, you say "Fail." The robot learns to just fit the items, even if it leaves half the suitcase empty or packs heavy items on top of fragile ones.
The FORGE Way (Quality-Aware Rewards): You tell the robot, "Success is good, but if you can fit more stuff in the same space, or arrange it so it's lighter, you get a bigger reward." The robot learns to keep trying to improve the packing until it's perfect.

The paper argues that current AI models are great at finding "valid" solutions (like a route that works) but terrible at finding "optimal" solutions (the absolute best route). This is a big deal for real-world problems like logistics, scheduling, and network design, which are known as NP-hard problems (math problems that are incredibly hard to solve perfectly).

2. The Solution: The "Forge" Factory

The authors built a factory called FORGE-ENGINE to train these robots to be better optimizers. Think of it as a gym for the AI's brain, but instead of lifting weights, it's solving complex puzzles.

The factory has three main machines:

The Generator: This machine creates millions of practice puzzles. It can make them easy (like a 5-piece puzzle), medium, or hard (like a 1,000-piece puzzle).
The Validator: This is the strict referee. It checks if the robot's solution actually follows the rules (e.g., "Did you visit every city exactly once?").
The Heuristic Solver (The Secret Sauce): This is the most important part. It's a super-fast, traditional computer program that solves the puzzle almost perfectly. It acts as a "gold standard" or a coach.
- The Analogy: Imagine the robot is a student taking a test. The Validator checks if the answer is written correctly. The Heuristic Solver is the teacher who has the answer key. If the robot gets 80% of the points, the teacher doesn't just say "Wrong." The teacher says, "You got 80%. Try to get to 90%." This gives the robot a continuous score instead of a simple "Pass/Fail."

3. The Training Method: "Climbing the Mountain"

You can't throw a robot into a 1,000-piece puzzle immediately; it would get confused and give up. So, the paper uses a Curriculum Learning strategy:

Easy Stage: The robot solves small, simple puzzles to learn the rules.
Medium Stage: The puzzles get bigger. The robot learns to plan ahead.
Hard Stage: The robot tackles massive, complex puzzles.
The Replay Trick: The authors noticed that if you only move forward (Easy → Hard), the robot forgets how to do the easy stuff. So, they made the robot replay the easy and medium levels periodically. This keeps its skills sharp while it learns the hard stuff.

4. The Results: A Smarter Brain

They tested their new robot (named FORGE) on 10 different types of difficult puzzles (like planning the shortest path for a delivery truck or scheduling meetings without conflicts).

The Score: The robot didn't just find a solution; it found great solutions. It beat the famous GPT-4o model by a huge margin. While GPT-4o found valid solutions about 62% of the time, FORGE found them 93% of the time. More importantly, FORGE's solutions were much closer to the "perfect" answer.
The Bonus Effect: Here is the coolest part. When they trained the robot on these hard optimization puzzles, it didn't just get better at puzzles. It got better at everything else, too.
- It got better at math.
- It got better at logic.
- It got better at following instructions.
- The Analogy: It's like training a chess player to be a grandmaster. In the process, they don't just get better at chess; they get better at strategy, patience, and planning in their daily life. The paper suggests that learning to "optimize" (find the best solution) teaches the AI a general skill of thinking deeply and refining its answers, which helps it in all kinds of tasks.

Summary

The paper introduces FORGE, a new way to train AI. Instead of just teaching AI to get the "right" answer, they teach it to find the best possible answer by giving it a constant score on how good its solution is. This turns the AI into a master optimizer that not only solves hard math puzzles better than current top models but also becomes smarter at general reasoning, logic, and following instructions.

Technical Summary: FORGE-ENGINE

Problem Statement

While Large Language Models (LLMs) have achieved significant success in reasoning benchmarks (math, coding, logic) via Reinforcement Learning with Verifiable Rewards (RLVR), existing evaluation paradigms suffer from a critical limitation: they focus exclusively on correctness (feasibility) rather than optimality. In many real-world scenarios, such as logistics or resource allocation, a merely "valid" solution is insufficient; the goal is to find the solution that maximizes or minimizes an objective function. Current benchmarks often rely on binary rewards (pass/fail), which fail to provide the dense, continuous feedback necessary to guide models toward optimal solutions. Furthermore, existing NP-hard benchmarks lack mechanisms to generate ground-truth optimal values, preventing fine-grained reward computation and reverting training to inefficient binary feedback loops.

Methodology: FORGE-ENGINE

The authors propose FORGE-ENGINE, a comprehensive framework designed to train and evaluate LLMs on NP-hard optimization problems using Quality-Aware RLVR. The framework consists of three core components:

1. Data Construction Pipeline

The framework employs a scalable, four-stage pipeline to generate training data across 10 distinct NP-hard tasks spanning five categories: Graph Clustering, Resource Scheduling, Graph Partitioning, Subset Selection, and Path Planning.

Controllable Generators: Produce infinite instances with graded difficulty (Easy, Medium, Hard) using a "construction-by-design" approach to guarantee the existence of feasible solutions.
Rule-based Verifiers: Automatically validate constraint satisfaction and compute objective values, ensuring absolute correctness without human labeling.
Heuristic Solvers: A critical innovation where task-specific heuristic algorithms (e.g., multi-start nearest-neighbor with 2-opt for TSP) compute near-optimal baselines for every instance. These baselines serve as the ground truth for reward calculation.
Curriculum Calibration: Difficulty levels are empirically calibrated to achieve target success rates (e.g., Easy: 70-90%, Hard: 10-40%) to facilitate progressive learning.

2. FORGE-BENCH

A rigorous evaluation benchmark comprising 1,000 high-complexity instances (100 per task) designed to be unseen during training. Unlike prior works, it evaluates models using two complementary metrics:

Success Rate (SR): The percentage of instances where the model generates a valid solution satisfying all constraints.
Quality Ratio (QR): A continuous metric measuring solution quality relative to the heuristic baseline (e.g., $QR = \text{Heuristic} / \text{Model}$ for minimization). This distinguishes between merely feasible and near-optimal solutions.

3. Quality-Aware RLVR Framework

The training strategy replaces sparse binary signals with fine-grained, continuous rewards derived from the heuristic baselines.

Reward Function: A hierarchical reward system evaluates outputs on three dimensions:
1. Format Reward ( $R_{format}$ ): Ensures structured reasoning (e.g., <thought> tags).
2. Feasibility Reward ( $R_{feasibility}$ ): Penalizes constraint violations (e.g., -1.5) and passes the optimality score for valid solutions.
3. Optimality Reward ( $R_{optimal}$ ): A continuous score based on the ratio of the model's objective value to the heuristic baseline, providing dense feedback for optimization.
Multi-Stage Curriculum Replay: To mitigate the "cold start" problem and catastrophic forgetting, the training employs a cyclic strategy (Easy $\to$ Medium $\to$ Hard) that iterates through difficulty levels multiple times, reinforcing foundational skills while progressively tackling complex optimization.

Key Results

Experiments were conducted on Qwen2.5-7B-Instruct-1M, trained on 15,000 examples (15K total across three stages).

In-Domain Performance: The resulting model, FORGE, achieved a 93.1% Success Rate and 46.6% Quality Ratio on FORGE-BENCH. This significantly outperforms GPT-4o (62.1% SR, 36.2% QR) and substantially exceeds the base model's performance (29.6% SR, 14.6% QR).
Out-of-Domain (OOD) Generalization: Training on FORGE-ENGINE transferred to diverse reasoning tasks, yielding gains in:
- Instruction Following: +6.1%
- Knowledge Reasoning: +4.1%
- Mathematics: +2.2%
- Logic: +1.2%
Ablation Studies:
- Reward Design: Quality-aware rewards improved solution quality by 28.8% over binary rewards, preventing reward hacking and guiding the model toward true optimality.
- Curriculum Strategy: Multi-stage curriculum replay outperformed linear single-pass strategies, breaking performance ceilings and preventing stagnation.
- Task Diversity vs. Quantity: The study found that training on a diverse set of 10 tasks yielded better generalization than training on fewer tasks with more data, suggesting that exposure to varied problem structures is a key driver of reasoning capabilities.
- RLVR vs. SFT: While Supervised Fine-Tuning (SFT) on optimal solutions improved in-domain performance, it suffered from catastrophic OOD degradation (-12.0 points). In contrast, RLVR fostered robust, transferable constraint satisfaction strategies.

Significance and Claims

The paper claims that FORGE-ENGINE represents the first comprehensive framework to successfully apply RLVR to NP-hard optimization problems by integrating heuristic solvers to provide fine-grained optimality rewards.

The authors assert that this work bridges the gap between feasibility and optimality in LLM reasoning. By moving beyond binary correctness, the framework enables models to engage in iterative optimization search. Crucially, the paper provides empirical evidence that optimization training fosters a transferable "optimization mindset"—the ability to generate candidates and self-refine—which enhances broader reasoning capabilities beyond the specific optimization tasks. The findings suggest that for scaling RLVR, task diversity and quality-aware rewards are more critical factors than sheer data quantity.

The work is presented as a foundational step for integrating LLMs with optimization-based reasoning, offering a scalable infrastructure and benchmark for future research in complex reasoning.