🤖 machine learning

Toward Compiler World Models: Learning Latent Dynamics for Efficient Tensor Program Search

This paper proposes a world-model-inspired evaluator that learns latent dynamics of scheduling actions to efficiently rank tensor program candidates, achieving significant latency improvements over existing auto-schedulers like Ansor while drastically reducing the number of required measurements.

Original authors: Haolin Pan, Lianghong Huang, Xvlin Zhou, Mingjie Xing, Yanjun Wu

Published 2026-06-09

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Haolin Pan, Lianghong Huang, Xvlin Zhou, Mingjie Xing, Yanjun Wu

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 trying to find the absolute fastest route to drive from your home to a friend's house. You have a map, but the traffic conditions change constantly, and there are millions of possible routes you could take.

In the world of computer science, specifically for Machine Learning, this is exactly what a "compiler" does. It tries to find the most efficient way for a computer to perform complex math tasks (called "tensor programs"). The problem is that there are so many possible ways to write the code that checking every single one by actually running it on the computer is incredibly slow and expensive. It's like trying to drive every single possible route just to see which one is fastest; you'd run out of gas before you found the best one.

The Old Way: Taking a Snapshot

Previously, computer programs that tried to solve this (called "auto-schedulers") acted like a photographer taking a snapshot of the final destination. They would look at the finished code, guess how fast it would be, and decide if it was good.

The paper argues this is a bad idea because:

It ignores the journey: It doesn't understand how the code got there. Two different routes might end up at the same spot, but one might have been a smooth highway while the other was a bumpy dirt road. The snapshot looks the same, but the experience (and speed) is different.
It gets confused by small details: If you change a few words in the code that don't actually change how it works, the old system might think it's a completely different (and worse) route.

The New Idea: A "World Model" (The GPS Simulator)

The authors propose a new approach inspired by World Models. Think of this not as a photographer, but as a high-tech GPS simulator.

Instead of just looking at the final destination, this new system simulates the entire journey step-by-step in its "mind" (a mathematical space called a "latent space").

Here is how it works, using a cooking analogy:

The Ingredients (Initial State): You start with a raw recipe (the unoptimized code).
The Chef's Moves (Actions): The compiler makes decisions, like "chop the onions," "sauté for 5 minutes," or "add salt."
The Old Way: The chef looks at the final dish and guesses, "Hmm, this tastes okay."
The New Way (World Model): The chef has a mental simulation. They imagine: "If I chop the onions then sauté them, the texture will be X. If I sauté then chop, the texture will be Y." They simulate the process of cooking in their head to predict the final taste without actually cooking the whole meal first.

How They Built It

The researchers built a system with three main parts:

The Translator (Encoder): It turns the messy computer code into a clean, mathematical "thought" (a vector) that the computer can understand easily.
The Simulator (Transition Model): This is the core innovation. It takes the "thought" of the current code and applies the "chef's moves" (scheduling actions) one by one. It predicts what the code will look like after each step, all inside the computer's memory, without actually running the code.
The Judge (Ranking Model): Once the simulation is done, the judge looks at the predicted final result and says, "This route is likely the fastest," or "That one looks slow."

The Results

They tested this on two types of computers: a powerful CPU (Intel Xeon) and a high-end Graphics Card (NVIDIA RTX 4090).

Faster Results: They found better code schedules much faster than the previous best method (called Ansor).
Less Work: They achieved results that were just as good as the old method, but they only had to run "test drives" (measurements) 10 times fewer.
Real-World Speed: When they used this to run actual AI models (like image recognition or language models), the programs ran 4 to 5 times faster than standard versions, and in some cases, up to 58 times faster.

The Bottom Line

The paper claims that by teaching the computer to understand the process of optimization (the journey) rather than just the result (the snapshot), we can find the fastest code much more efficiently. It's like having a GPS that simulates traffic in its head to find the best route, rather than just guessing based on a photo of the destination.

Limitations mentioned in the paper:

This system is a "judge" that helps pick the best route; it doesn't invent the routes itself. If the search engine doesn't propose any good routes to begin with, the judge can't fix it.
If the "journey" is extremely long and complex, the simulation in the computer's head might make small mistakes that add up, making the prediction less accurate.
It is designed to rank options against each other (which is faster?), not to predict the exact time down to the millisecond.

Technical Summary: Toward Compiler World Models

Problem Statement

Tensor program optimization is essential for modern machine learning systems, yet identifying the most efficient program is hindered by a massive combinatorial search space. To avoid the prohibitive cost of on-device measurements, existing auto-schedulers (e.g., Ansor, AutoTVM) rely on learned cost models to rank candidate programs. However, these models typically treat candidates as static code snapshots (e.g., final TensorIR text).

This static perspective introduces two critical limitations:

Loss of Transformation Context: It ignores the step-by-step transformation process that generated the program. Consequently, evaluators cannot understand how a scheduling decision in a later step depends on the context established by earlier decisions.
Sensitivity to Syntactic Noise: Static encoders are easily misled by superficial code variations. For instance, two structurally equivalent programs (e.g., matrix multiplication with loop sizes $N=1024$ vs. $N=1023$ ) may produce significantly different TensorIR texts due to boundary handling (loop peeling), causing the model to misclassify their performance potential despite identical hardware behavior.

Methodology

The authors propose a Compiler World Model framework that reframes schedule evaluation as action-conditioned latent dynamics over program states. Instead of evaluating a final static code artifact, the framework simulates the evolution of the program state through a continuous latent space.

The framework consists of three learned components instantiated within the TVM AutoScheduler:

1. Tensor Program State Representation (Encoder)

Input: TensorIR text (both pre-schedule and post-schedule states).
Architecture: A trainable CodeBERT-based architecture adapted for long inputs. It uses a sliding window to tokenize and encode overlapping chunks, aggregating them via masked mean pooling into a dense vector representation $z = E(s)$ .
Training: Uses contrastive learning to map similar TensorIR states (from the same workload source) close together in the latent space while pushing dissimilar states apart. This creates a state space where semantic equivalence is preserved despite syntactic variations.

2. Action-Conditioned Multi-Step State Transition

Goal: Predict the terminal-state representation induced by a sequence of scheduling actions without explicitly materializing intermediate ASTs.
Mechanism: Starting from an initial state embedding $z_0$ $z_{0}$ , the model rolls out the effect of a scheduling action sequence $a_{1:T}$ $a_{1 : T}$ step-by-step.
- Actions are serialized and encoded into dense embeddings.
- A transition model $F(\cdot)$ , implemented using a TransH-based architecture (originally from knowledge graph representation learning), predicts the next latent state: $z_t = F(z_{t-1}, a_t)$ .
- This models scheduling actions as geometric transformations in a relation-specific hyperplane, effectively tracking decision provenance and resolving referential ambiguities (e.g., what a loop variable represents after fusion).
Training: Uses multi-step rollout loss combining a state-matching term (minimizing distance to ground-truth future states) and a margin-based ranking term (ensuring the predicted state is closer to the ground truth than to negative samples).

3. Ranking-Based Candidate Evaluation

Input: The predicted terminal-state representation $\hat{z}_T$ , structured action features (e.g., action counts, sequence length), and optional hardware features.
Output: A scalar score for relative candidate prioritization within the same workload.
Architecture: An XGBoost ranking model trained with a LambdaRank objective.
Training Strategy: The model does not predict absolute runtime. Instead, it learns to rank candidates within a workload group based on measured execution times, treating the problem as a learning-to-rank task.

Key Contributions

World-Model Perspective for Compilers: The paper adapts the world-model paradigm to tensor optimization, establishing a framework where candidate schedules are modeled through action-conditioned latent state evolution rather than static snapshots.
Action-State Trajectory Dataset: Construction of a TVM/TenSet-based dataset that organizes tuning logs into aligned action-state trajectories (pre-schedule state, action sequence, intermediate states, post-schedule state) for multi-step prediction.
Efficient Search Integration: Implementation of the framework in TVM AutoScheduler, demonstrating that modeling latent dynamics significantly reduces autotuning costs while maintaining high search quality.

Experimental Results

The method was evaluated on an Intel Xeon Gold 6430 CPU and an NVIDIA RTX 4090 GPU across seven neural network models (including ResNet, BERT, GPT, and OPT) and 103 extracted subgraphs.

Performance vs. Ansor: Under a 64-trial budget, the method improves representative-subgraph latency over Ansor by 1.37× on GPU and 1.54× on CPU.
Sample Efficiency: The method achieves comparable search quality to Ansor-10K (10,000 trials) using only 1,000 trials (10× fewer measurements). Specifically, it slightly outperforms Ansor-10K in arithmetic mean and matches it within 2.2% in geometric mean on representative subgraphs.
End-to-End Speedup: Compared to PyTorch/PyTorch-opt (cuDNN), the method delivers geometric mean speedups of 4.61× (CPU) and 3.67× (GPU), with peak speedups reaching up to 58.61×.
Ablation Studies: Removing the learned transition module (reverting to static initial or final state encoding) resulted in significant latency increases (median ~56–69%), confirming that modeling the intermediate state evolution is the most critical component.

Significance and Limitations

Significance:
The paper argues that treating compiler optimization as a dynamic process of state transitions, rather than a static code evaluation problem, allows evaluators to capture the semantic dependencies between scheduling decisions. This approach effectively suppresses syntactic noise and resolves referential ambiguities, leading to more efficient search with fewer hardware measurements.

Limitations (as stated by the authors):

Dependency on Search Space: The method improves candidate evaluation, not generation. Its effectiveness is bounded by the underlying scheduler's ability to propose strong schedules; a better evaluator cannot recover schedules that are never proposed.
Error Accumulation: Prediction errors may accumulate over long or complex scheduling trajectories in latent space, potentially reducing reliability for difficult cases.
Ranking vs. Prediction: The model is trained for within-workload ranking, not absolute latency prediction. This makes it ideal for online prioritization but less suitable for settings requiring calibrated performance estimates across different workloads.