Yuandong Tian on ML for Systems: From Game Search to System Optimization

At the NeurIPS 2021 MLFS workshop, Yuandong Tian, a research scientist and manager at Facebook AI Research, shared a core insight: the key to reinforcement learning's enormous success in complex games like Go, chess, Dota, and StarCraft is that it learned to provide efficient heuristic guidance for search. So can the same idea be transferred to the field of system optimization?

System optimization typically involves a massive number of design choices—there may be tens of thousands or even hundreds of thousands of parameters to tune. Traditional approaches rely on human experts to understand the system and spend a lot of time tuning parameters, which is essentially finding a path through a huge search tree using human knowledge. Yuandong Tian believes that machine learning is fully capable of finding a path to system optimization goals without relying so heavily on human knowledge.

Essential Differences Between System Optimization and Game Search

When applying reinforcement learning to system optimization, one interesting difference is that we don't care about "how" the solution is obtained—only what the final solution is. For example, with 100 parameter knobs, which one to adjust first doesn't matter; what matters is whether the final configuration satisfies all constraints. This provides options for parameterizing optimization problems that differ from traditional reinforcement learning settings.

Yuandong Tian listed several possible strategies: one-step prediction (directly predicting the solution from the problem description), progressive prediction (predicting partial solutions step by step), and iterative refinement (starting from an initial solution and continuously improving it until reaching a local optimum). One could even learn the action space itself—automatically identifying the best parameterization for optimization from the system description. Due to time constraints, he focused on three works under the "iterative refinement" direction.

Case 1: Two-Stage Solving in Neural Network Planning

Network planning is a typical challenge in data center interconnection: given several data center nodes, one must decide on fiber capacity purchase plans while satisfying all traffic demands, controlling costs, and ensuring robustness against potential failures. Such problems can be modeled as linear programming (LP) and in principle solved with solvers like Gurobi or CPLEX, but scalability is a concern—solving can take days or even weeks.

In practice, humans design heuristic methods to handle topological constraints, decomposing the problem by geographic structure and solving each part separately. But the core pain point of such methods is the difficulty of finding feasible solutions. Yuandong Tian's team proposed a two-stage approach: in the first stage, a reinforcement learning agent iteratively refines the solution from scratch, adding a small amount of capacity to some edge each time until the entire network configuration satisfies all traffic constraints; in the second stage, this initial feasible solution is handed to an LP solver as an excellent starting point to accelerate the final solve.

The key to why this approach works is that checking the feasibility of a solution is far easier than finding the minimum-cost solution. The team also introduced two acceleration techniques—source aggregation (merging multiple nodes to reduce graph size) and stateful failure checking (since capacity is only added each time, already-passed failure scenarios need not be rechecked). Experimental results show that the method, called Neural Plan, achieves costs comparable to LP solvers on small topologies, while on larger topologies where LP solvers time out and cannot produce meaningful results, Neural Plan still outputs good costs. The code has been open-sourced.

Linear programming (LP) is a classic modeling framework in operations research: maximizing or minimizing a linear objective function under several linear inequality constraints. Its standard form can be written as min cᵀx s.t. Ax ≤ b, x ≥ 0, where x is the decision variable vector, c is the objective coefficients, and A, b describe the constraints. An important property of LP is that the feasible region is a convex polytope, and the optimal solution must occur at a vertex; the simplex method and interior-point methods are designed based on this. Gurobi and CPLEX are industrial-grade LP/MILP solvers that can solve small-scale problems in seconds, but their algorithms such as branch-and-bound have exponential complexity in the worst case, and solving time explodes when facing large-scale instances in network planning with tens of thousands of edges and massive failure scenarios. This is precisely the practical motivation for the article's proposal to use learned heuristics to provide high-quality initial feasible solutions to solvers, thereby pruning and accelerating the process.

Case 2: Iterative Rewriting Framework for Code Decompilation

The second work focuses on code decompilation—recovering high-level programming language source code from binary executable code, which is the inverse process of compilation. The challenge is that there are many hardware architectures and programming languages, and traditional methods require extensive manual rewriting of the entire decompilation framework for each new language supported. The team's goal is to build an end-to-end general framework that simultaneously accounts for the functionality and semantics of binary programs.

Their approach is also divided into two stages. The first stage generates a code sketch from assembly code: the assembly code is converted into a graph structure, an N-ary tree encoder captures the inter-instruction and intra-instruction dependencies of the low-level code (such as shared memory, branch jumps, etc.), and then an attention mechanism decodes it into an abstract syntax tree (AST) rather than directly generating token-level source code, thereby exploiting the syntactic structure of the source code. This stage can produce syntactically correct programs that may contain errors, such as misidentified tokens, missing lines, or redundant lines.

The second stage is iterative rewriting: the compiler is used to recompile the generated source code back into assembly, compare it with the input assembly, and after errors are found, correct them and repeat the process. Experiments were conducted on multiple synthetic program benchmarks under Clang O0 configuration. Results show that the approach combining AST and attention mechanisms outperforms naive sequence-to-sequence models in token accuracy, and performance further improves with error correction. On the stricter program accuracy metric, the gain from error correction is particularly significant—for example, on a certain math library program it rose from 64% all the way to 100%.

The subsequent version Code Breif (Code First 2.0) further unified the pipeline: starting from an empty AST tree and iteratively refining, modeling the assembly graph as a multi-type graph containing data flow edges, instruction structure, and control flow edges, and during training using known assembly-source pairs to make each step's output match the source AST. On harder datasets containing library function calls, performance further improved compared to the previous generation. Ablation studies confirmed the necessity of components such as graph neural network encoding of program structure, positional encoding, and memory augmentation. In addition, data type recovery, as a parallel process, also adopts an iterative refinement structure, further improving decompilation performance.

An abstract syntax tree (AST) is a tree-like representation of the syntactic structure of source code by the compiler front end: internal nodes correspond to syntactic constructs such as operators, statements, or declarations, and leaf nodes correspond to lexical units such as identifiers and literals. Compared with the raw token sequence, the AST explicitly encodes nesting and precedence relationships—for example, in the expression a+bc, bc is located in a deeper subtree in the tree, so associativity can be expressed without parentheses. Precisely because of this, letting the model directly decode into an AST rather than generating source code token by token can build syntactic validity as a structural constraint into the output, significantly reducing the probability of generating illegal programs. At the same time, ASTs and source code can be converted bidirectionally, making it convenient in the iterative rewriting stage to restore the generated result to source code and then compile it back to assembly for comparison, which constitutes the key interface of the article's two-stage framework.

Challenges and Outlook

Yuandong Tian concluded that machine learning for learning heuristics for system optimization has been proven feasible, but many challenges remain ahead. First, the enormous state-action space is difficult to handle, requiring careful design of state and action representations; second, there are many correlated actions in real systems; and third, the evaluation speed of system optimization is slow, severely constraining the application of reinforcement learning.

In the network planning case, the team accelerated feasibility evaluation through stateful failure checking and source aggregation. But for some system problems, running the system itself may take days or even weeks, and at that point how to learn a good value function to accelerate evaluation becomes a key issue.

At the end of the talk, Yuandong Tian introduced the team's new work CompilerG—a robust, high-performance reinforcement learning environment for compiler optimization. There are many tunable knobs in compilers, used to make programs run faster or binaries smaller, but tuning is extremely time-consuming and relies on domain expertise. CompilerG aims to let machine learning or reinforcement learning play a role without requiring users to have compiler domain knowledge.

The value function in reinforcement learning is used to estimate the long-term return starting from a certain state and is a core component of policy evaluation and improvement. In system optimization scenarios, if a real evaluation requires running the system for days, the agent cannot learn through massive trial and error as in games. A feasible approach at this point is to use a learned value function or surrogate model to replace real evaluation, quickly predicting the payoff of candidate actions and calling expensive simulation only at key nodes. This is in the same vein as the logic in AlphaGo of using a value network to evaluate the winning probability of a board position and reduce the number of Monte Carlo rollouts, and it is also the deeper technical correspondence in the article for transferring "providing heuristics for search" from games to system optimization.

One-line Summary

Reinforcement learning can provide efficient heuristics for system optimization, outperforming traditional methods in multiple cases through iterative refinement.

Paragraph Summary

Based on Yuandong Tian's sharing at the NeurIPS 2021 MLFS workshop, the article proposes that the key to reinforcement learning's success in complex games is providing efficient heuristics for search, and this idea can be transferred to system optimization. System optimization does not care about the solution path but only the final solution, so parameterization methods such as one-step prediction, progressive prediction, and iterative refinement can be adopted. The article focuses on three iterative refinement works: network planning uses a reinforcement learning agent to iteratively generate feasible solutions, which then serve as starting points for LP solvers, outperforming traditional solvers on large-scale topologies; code decompilation uses a two-stage framework of converting assembly to a graph, decoding to an AST, and iteratively correcting through compilation comparison, significantly outperforming sequence-to-sequence baselines in token and program accuracy; the subsequent Code Breif further unifies the pipeline and improves performance. Finally, it points out challenges such as the huge state-action space, correlated actions, and slow evaluation, and introduces the CompilerG environment to lower the domain knowledge barrier for compiler optimization.