想象一下，你正试图寻找一条从家到朋友家最快的驾驶路线。你有一张地图，但交通状况时刻在变，而且你可能可以走的路线有数百万条。
在计算机科学领域，特别是针对机器学习时，这正是“编译器”所做的事情。它试图为计算机执行复杂的数学任务（称为“张量程序”）找到最高效的方法。问题在于，编写代码的可能方式太多了，如果通过在计算机上实际运行每一种方式来检查，速度会极其缓慢且成本高昂。这就像是为了看哪条路最快而尝试开遍每一条可能的路线一样；你会还没找到最好的路就耗尽汽油了。
旧方法：拍摄“快照”
以前，尝试解决此类问题的计算机程序（被称为“自动调度器”）就像是一个拍摄最终目的地快照的摄影师。它们会观察生成的代码，猜测其运行速度，然后决定它是否足够好。
论文指出这并非良策，因为：
- 它忽略了过程： 它不理解代码是如何生成的。两条不同的路线可能会到达同一个终点，但其中一条可能是平坦的高速公路，而另一条可能是颠簸的土路。快照看起来是一样的，但体验（和速度）却截然不同。
- 它会被微小的细节迷惑： 如果你修改了代码中一些实际上并不改变其功能的词汇，旧系统可能会认为这是一条完全不同且更差的路线。
新想法：“世界模型”（GPS 模拟器）
作者提出了一种受世界模型（World Models）启发的全新方法。请记住，这不仅仅是一个摄影师，而是一个高科技 GPS 模拟器。
它不只是观察最终目的地，而是会在它的“大脑”里（一个被称为“潜空间”的数学空间）逐步模拟整个旅程。
以下是使用烹饪类比的工作原理：
- 食材（初始状态）： 你从一个原始食谱开始（未优化的代码）。
- 厨师的动作（行动）： 编译器做出决策，比如“切洋葱”、“翻炒 5 分钟”或“加盐”。
- 旧方法： 厨师观察最终的菜肴并猜测：“嗯，味道还可以。”
- 新方法（世界模型）： 厨师拥有心理模拟能力。他会想象：“如果我先切洋葱再翻炒，口感会是 X；如果我先翻炒再切，口感会是 Y。”他在脑海中模拟烹饪的过程，从而在实际动手做饭之前，就能预测出最终的味道。
他们是如何构建它的
研究人员构建了一个包含三个部分的系统：
- 翻译器（编码器）： 它将杂乱的计算机代码转化为一种干净的、计算机可以轻松理解的数学“思想”（向量）。
- 模拟器（转换模型）： 这是核心创新。它获取当前代码的“思想”，并逐一应用“厨师的动作”（调度动作）。它会在计算机内存中预测代码在每一步之后的样子，而无需实际运行代码。
- 裁判（排序模型）： 模拟完成后，裁判会观察预测的最终结果，并判断：“这条路线可能最快，”或者“那条看起来很慢。”
实验结果
他们针对两种类型的计算机进行了测试：一款强大的 CPU（Intel Xeon）和一款高端显卡（NVIDIA RTX 4090）。
- 更快的产出： 他们发现，相比于之前的最佳方法（称为 Ansor），他们能更快地找到更好的代码调度方案。
- 更少的工作量： 他们实现了与旧方法同样出色的效果，但实际进行的“试驾”（测量）次数减少了 10 倍。
- 现实世界速度： 当他们使用此技术运行实际的 AI 模型（如图像识别或语言模型）时，程序的运行速度比标准版本快了 4 到 5 倍，在某些情况下甚至快了高达 58 倍。
核心结论
该论文声称，通过教会计算机理解优化过程（旅程）而非仅仅是结果（快照），我们可以更高效地找到最快的代码。这就像拥有一个能在脑海中模拟交通状况的 GPS，从而找到最佳路线，而不是仅仅根据一张照片来猜测。
论文中提到的局限性：
- 该系统是一个帮助挑选最佳路线的“裁判”；它本身并不发明路线。如果搜索引擎最初没有提出任何好的路线，裁判也无法修复它。
- 如果“旅程”极其漫长且复杂，计算机大脑中的模拟可能会产生累积的小误差，从而降低预测的准确性。
- 它的设计目的是为了对选项进行排序（哪个更快？），而不是精确预测到毫秒级的运行时间。
技术摘要：迈向编译器世界模型
问题陈述
张量程序优化对于现代机器学习系统至关重要，然而，由于搜索空间巨大，寻找最有效率的程序变得异常困难。为了避免在设备上进行测量所带来的高昂成本，现有的自动调度器（如 Ansor、AutoTVM）依赖于学习到的代价模型来对候选程序进行排序。然而，这些模型通常将候选程序视为静态代码快照（例如，最终的 TensorIR 文本）。
这种静态视角引入了两个关键局限性：
- 丢失变换上下文： 它忽略了生成程序的逐步变换过程。因此，评估器无法理解后续步骤中的调度决策是如何依赖于早期决策所建立的上下文的。
- 对语法噪声敏感： 静态编码器容易被表面的代码变化所误导。例如，两个结构等价的程序（例如，循环大小分别为 N=1024 与 N=1023 的矩阵乘法）可能会因为边界处理（循环剥离/loop peeling）产生截然不同的 TensorIR 文本，导致模型尽管它们的硬件行为完全一致，却会错误地分类其性能潜力。
方法论
作者提出了一个编译器世界模型（Compiler World Model）框架，将调度评估重新定义为程序状态上的动作条件潜空间动力学（action-conditioned latent dynamics）。该框架不再评估最终的静态代码产物，而是通过连续的潜空间模拟程序状态的演变。
该框架由三个在 TVM AutoScheduler 中实例化的学习组件组成：
1. 张量程序状态表示（编码器）
- 输入： TensorIR 文本（包括调度前和调度后的状态）。
- 架构： 基于可训练的 CodeBERT 架构，并针对长输入进行了适配。它使用滑动窗口对重叠块进行分词和编码，并通过掩码平均池化（masked mean pooling）将其聚合为稠密向量表示 z=E(s)。
- 训练： 使用对比学习，将来自相同工作负载源的相似 TensorIR 状态映射到潜空间中相近的位置，同时将不相似的状态推开。这创建了一个即使在存在语法变化时也能保持语义等价性的状态空间。
2. 动作条件的多步状态转移
- 目标： 在不显式实例化中间抽象语法树（AST）的情况下，预测由一系列调度动作诱导的终端状态表示。
- 机制： 从初始状态嵌入 z0 开始，模型逐步模拟调度动作序列 a1:T 产生的效果。
- 动作被序列化并编码为稠密嵌入。
- 转移模型 F(⋅) 采用基于 TransH 的架构（最初用于知识图谱表示学习）来实现，用于预测下一个潜状态：zt=F(zt−1,at)。
- 这将调度动作建模为特定关系超平面中的几何变换，从而有效地追踪决策溯源并解决引用歧义（例如，在融合后循环变量代表什么）。
- 训练： 使用结合了状态匹配项（最小化与真实未来状态的距离）和基于边际的排序项（确保预测状态比负样本更接近真实情况）的多步展开损失（multi-step rollout loss）。
3. 基于排序的候选程序评估
- 输入： 预测的终端状态表示 z^T、结构化动作特征（如动作计数、序列长度）以及可选的硬件特征。
- 输出： 同一工作负载内候选程序的相对优先级标量分数。
- 架构： 一个使用 LambdaRank 目标函数训练的 XGBoost 排序模型。
- 训练策略： 该模型并不预测绝对运行时间。相反，它学习在同一个工作负载组内根据实测执行时间对候选程序进行排序，将问题视为一个学习排序（learning-to-rank）任务。
核心贡献
- 世界模型视角下的编译器： 本文将世界模型范式引入张量优化，建立了一个通过动作条件下的潜状态演变而非静态快照来对候选调度进行建模的框架。
- 动作-状态轨迹数据集： 构建了一个基于 TVM/TenSet 的数据集，将调优日志组织成对齐的动作-状态轨迹（包括：调度前状态、动作序列、中间状态、调度后状态），用于多步预测。
- 高效搜索集成： 在 TVM AutoScheduler 中实现了该框架，证明了通过建模潜动力学可以显著降低自动调优成本，同时保持极高的搜索质量。
实验结果
该方法在 Intel Xeon Gold 6430 CPU 和 NVIDIA RTX 4090 GPU 上进行了评估，涵盖了七种神经网络模型（包括 ResNet、BERT、GPT 和 OPT）以及 103 个提取的子图。
- 性能对比 Ansor： 在 64 次试验的预算下，该方法在 GPU 上的代表性子图延迟降低了 1.37×，在 CPU 上降低了 1.54×。
- 样本效率： 该方法仅使用 1,000 次试验（减少了 10 倍的测量次数）即可达到与 Ansor-10K（10,000 次试验）相当的搜索质量。具体而言，在代表性子图上，其算术平均值略高于 Ansor-10K，且几何平均值差距在 2.2% 以内。
- 端到端加速： 与 PyTorch/PyTorch-opt (cuDNN) 相比，该方法实现的几何平均加速比分别为 4.61× (CPU) 和 3.67× (GPU)，峰值加速可达 58.61×。
- 消融实验： 移除学习到的转移模块（退回到静态初始或最终状态编码）会导致延迟显著增加（中位数约 56–69%），这证实了建模中间状态演变是最关键的组成部分。
重要性与局限性
重要性：
本文认为，将编译器优化视为动态的状态转移过程，而非静态的代码评估问题，能够捕捉调度决策之间的语义依赖关系。这种方法有效地抑制了语法噪声并解决了引用歧义，从而实现更高效的搜索，且所需的硬件测量次数更少。
局限性（如作者所述）：
- 对搜索空间的依赖： 该方法改进的是候选程序的评估，而非生成。其有效性受限于底层调度器提出强力调度的能力；一个更好的评估器无法找回从未被提出的调度方案。
- 误差累积： 在潜空间中进行长序列或复杂调度轨迹的预测时，可能会出现预测误差累积，从而降低可靠性。
- 排序 vs. 预测： 模型是针对工作负载内的排序进行训练的，而非绝对延迟预测。这使其非常适合在线优先级排序，但不适用于需要跨不同工作负载进行校准性能估计的场景。
每周获取最佳 machine learning 论文。
受到斯坦福、剑桥和法国科学院研究人员的信赖。
请查收邮箱确认订阅。
出了点问题，再试一次？
无垃圾邮件，随时退订。