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

• URL: http://arxiv.org/abs/2606.09312v1
• Date: Mon, 08 Jun 2026 10:17:27 GMT
• Title: Toward Compiler World Models: Learning Latent Dynamics for Efficient Tensor Program Search
• Abstract summary: We propose an evaluator that schedules evaluation as action-conditioned latent dynamics over program states.<n>Our method improves representative-subgraph latency over Ansor by 1.37$times$ on GPU and 1.54$times$ on CPU under the same 64-trial budget.<n>It also accelerates full-model inference over PyTorch/PyTorch-opt(cuDNN) by 4.61$times$/3.67$times$ geometric mean.
• Score: 9.85275171877854
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Tensor program optimization is essential for modern machine learning systems, but its search space is enormous. Existing auto-schedulers reduce measurement cost with learned cost models, yet they usually evaluate each candidate as a static code snapshot, ignoring the schedule trajectory that produced it. This makes them insensitive to action dependencies and vulnerable to superficial code variations. We propose a \emph{world-model-inspired} evaluator that models schedule evaluation as action-conditioned latent dynamics over program states. Starting from the initial program, it rolls out scheduling actions in a continuous latent space with a lightweight transition model, avoiding expensive AST mutation and repeated code encoding. The final dynamic representation is combined with action and hardware features to rank candidates. Implemented in TVM AutoScheduler, our method improves representative-subgraph latency over Ansor by 1.37$\times$ on GPU and 1.54$\times$ on CPU under the same 64-trial budget. It also matches Ansor-10K within 2.2% geometric mean using 10$\times$ fewer measurements, and accelerates full-model inference over PyTorch/PyTorch-opt(cuDNN) by 4.61$\times$/3.67$\times$ geometric mean.

Related papers

• REFACTOR-VLA: Unsupervised Library Learning of Typed Motor Programs [0.0]
  Most vision-language-action (VLA) models are monolithic.<n>We introduce REFACTOR-VLA, a wake/sleepmodel system for learning reusable skills.
  arXiv  Detail & Related papers  (2026-09-01T13:19:05Z)
• AI4AI-Bench: Benchmarking LLM Agents in Algorithmic Design for Recursive Self-Improvement [4.962568775070268]
  Recursive self-improvement (RSI) asks whether an AI system can improve the process that produces AI systems.<n>We present AI4AImbox-Bench, 10 frozen research repositories spanning 10 training algorithm families.
  arXiv  Detail & Related papers  (2026-08-20T17:56:59Z)
• Learning-Augmented Scalable Linear Assignment Problem Optimization via Neural Dual Warm-Starts [19.540758462427878]
  We propose a learning-augmented framework that accelerates exact assignment solvers while maintaining optimality and worst-case guarantees.<n>We introduce RowDualNet, a lightweight row-independent architecture that avoids the $mathcalO(N2)$ memory bottleneck of graph-based models.
  arXiv  Detail & Related papers  (2026-05-10T07:15:49Z)
• $\ abla$-Reasoner: LLM Reasoning via Test-Time Gradient Descent in Latent Space [71.23672814629448]
  $nabla$-Reasoner is an iterative generation framework that integrates differentiable optimization over token logits into the decoding loop.<n>$nabla$-Reasoner achieves over 20% accuracy improvement on a challenging mathematical reasoning benchmark.
  arXiv  Detail & Related papers  (2026-03-05T08:42:54Z)
• Form Follows Function: Recursive Stem Model [0.0]
  We introduce Recursive Stem Model (RSM) to solve compute-heavy and NP puzzles.<n>RSM treats early iterations as detached "warm-up" steps, and applies loss only at the final step.<n>On Sudoku-Extreme, RSM reaches 97.5% exact accuracy with test-time compute.
  arXiv  Detail & Related papers  (2026-03-03T00:55:00Z)
• INC: An Indirect Neural Corrector for Auto-Regressive Hybrid PDE Solvers [61.84396402100827]
  We propose the Indirect Neural Corrector ($mathrmINC$), which integrates learned corrections into the governing equations.<n>$mathrmINC$ reduces the error amplification on the order of $t-1 + L$, where $t$ is the timestep and $L$ the Lipschitz constant.<n>We test $mathrmINC$ in extensive benchmarks, covering numerous differentiable solvers, neural backbones, and test cases ranging from a 1D chaotic system to 3D turbulence.
  arXiv  Detail & Related papers  (2025-11-16T20:14:28Z)
• $\texttt{SPECS}$: Faster Test-Time Scaling through Speculative Drafts [55.231201692232894]
  $textttSPECS$ is a latency-aware test-time scaling method inspired by speculative decoding.<n>Our results show that $textttSPECS$matches or surpasses beam search accuracy while reducing latency by up to $sim$19.1%.
  arXiv  Detail & Related papers  (2025-06-15T05:50:05Z)
• Position-Aware Depth Decay Decoding ($D^3$): Boosting Large Language Model Inference Efficiency [26.173523821684306]
  A token-position aware layer skipping framework is proposed to save 1.5x times operations efficiently while maintaining performance.<n>Experiments on large language models with $7 sim 70$ billion parameters show that $D3$ can achieve an average 1.5x speedup compared with the full-inference pipeline.
  arXiv  Detail & Related papers  (2025-03-11T15:15:54Z)
• Scalable 3D Registration via Truncated Entry-wise Absolute Residuals [65.04922801371363]
  A $3$D registration approach can process more than ten million ($107$) point pairs with over $99%$ random outliers. We call our method TEAR, as it involves minimizing an outlier-robust loss that computes Truncated Entry-wise Absolute Residuals.
  arXiv  Detail & Related papers  (2024-04-01T04:43:39Z)
• Pruner: A Draft-then-Verify Exploration Mechanism to Accelerate Tensor Program Tuning [11.139231636612479]
  We propose Pruner and MoA-Pruner to speed up program tuning.<n>Pruner is a "Draft-then-Verify" exploration mechanism that accelerates the schedule search process.<n>MoA-Pruner introduces a Momentum online Adaptation strategy to address the cross-platform online unawareness.
  arXiv  Detail & Related papers  (2024-02-04T06:11:12Z)
• Improving Dual-Encoder Training through Dynamic Indexes for Negative Mining [61.09807522366773]
  We introduce an algorithm that approximates the softmax with provable bounds and that dynamically maintains the tree. In our study on datasets with over twenty million targets, our approach cuts error by half in relation to oracle brute-force negative mining.
  arXiv  Detail & Related papers  (2023-03-27T15:18:32Z)
• Hidet: Task Mapping Programming Paradigm for Deep Learning Tensor Programs [11.338285393619042]
  We propose to embed the scheduling process into tensor programs and use dedicated mappings, called task mappings, to define the computation assignment and ordering. With the proposed paradigm, we implement a deep learning compiler - Hidet.
  arXiv  Detail & Related papers  (2022-10-18T05:32:13Z)
• Matching Pursuit Based Scheduling for Over-the-Air Federated Learning [67.59503935237676]
  This paper develops a class of low-complexity device scheduling algorithms for over-the-air learning via the method of federated learning. Compared to the state-of-the-art proposed scheme, the proposed scheme poses a drastically lower efficiency system. The efficiency of the proposed scheme is confirmed via experiments on the CIFAR dataset.
  arXiv  Detail & Related papers  (2022-06-14T08:14:14Z)
• Revisiting Smoothed Online Learning [70.09792747315323]
  We investigate the problem of smoothed online learning, in which the online learner suffers both a hitting cost and a switching cost. To bound the competitive ratio, we assume the hitting cost is known to the learner in each round, and investigate the greedy algorithm which simply minimizes the weighted sum of the hitting cost and the switching cost.
  arXiv  Detail & Related papers  (2021-02-13T14:15:55Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.