Tensor Seeks Layout: Formalizing Layout Selection for ML Compilers
本文通过将布局选择（layout selection）建模为一个组合优化问题，证明了其计算难度，并提出了针对有界树宽图的最优算法以及针对一般实例的加权 MaxSAT 编码，从而首次对机器学习编译器中的布局选择进行了正式研究，并证明了简单的启发式算法与最优解相比，性能可能会下降高达 5 倍。
原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。 阅读完整免责声明
现代人工智能依赖于庞大的数学模型，这些模型通过处理海量数据来识别语音、翻译语言或生成图像。为了让这些模型快速运行，工程师们使用专门为这种重型任务设计的专用计算机芯片。然而，这些芯片不仅仅是执行指令，它们还必须高效地移动数据。模型的运行速度往往并不取决于芯片的原始算力，而更多地取决于数据在其内存中的排列方式。想象一个图书馆，书籍存放在书架上。如果读者需要寻找一组特定的书，所花费的时间完全取决于这些书是散落在不同的走廊，还是整齐地排列在同一个书架上。在计算机芯片的世界里，这种排列被称为“布局”（layout）。当一个计算机程序执行计算时，它期望数据的排列方式符合某种特定格式，但程序的前一步骤可能会将数据以另一种格式留下。如果两者不匹配，计算机就必须停止并重新排列数据才能继续进行，这个过程既浪费时间又浪费能量。
多年来，用于为这些芯片准备模型的软件一直依赖于一系列粗略的猜测和经验法则来决定如何排列数据。这些规则在处理简单任务时效果尚可，但随着模型变得更加复杂，这些猜测开始失效，导致了显著的减速。来自维也纳理工大学和亚马逊的研究团队致力于改变这种方法。他们不再依赖直觉，而是将数据排列问题视为一个正式的数学谜题。他们构建了一个精确的模型，可以计算每种可能排列方式的精确成本，包括在不同格式之间移动数据所需的时间。通过这样做，他们可以为任何给定的模型确定单一的最佳数据组织方式，而不是仅仅希望一套规则能尽可能接近目标。
研究人员发现，寻找这种完美排列是一个极其困难的任务。用计算机科学的语言来说，这个问题非常复杂，以至于没有计算机能够为每一种可能的情况都快速求解，尤其是在模型规模不断扩大的情况下。他们证明了，即使是针对仅涉及基础矩阵运算的简化版本问题，其可能性的数量也如此之巨，以至于标准计算机无法在合理的时间内找到答案。这一发现排除了通过单一、快速且通用的算法来解决未来所有模型问题的可能性。然而，团队也找到了前进的方向。他们表明，虽然该问题在一般情况下很难，但当模型的结构类似于具有有限分支的树状结构时，它就变得可以处理了。对于这些在许多实际应用中常见的特定结构，他们设计了一种能够快速找到完美解的方法。对于不符合这种模式的更复杂的结构，他们开发了一种将问题转化为现有强大求解器可以处理的格式的方法，从而使他们即使在不存在完美的数学捷径时也能找到最佳排列。
为了测试他们的想法，研究人员将这种新方法实现到了一个用于亚马逊 Trainium 芯片的真实编译器中，这些芯片旨在运行人工智能模型。他们将这种新方法与目前行业内使用的标准方法（即依赖旧有经验法则的方法）进行了对比。结果令人震惊。在某些复杂的模型（特别是用于图像识别的模型）上，旧有的经验法则导致模型的运行速度比必要速度慢了多达五倍。这是因为简单的规则无法看到全局；它们会为某一步骤安排完美的数据，却为下一步创造了一团乱麻，迫使计算机不断浪费时间进行重新排列。而新方法通过同时观察整个步骤序列，避免了这些代价高昂的重新排列，并保持了数据的流畅流动。
然而，这项研究也揭示了一个关键的局性。虽然新方法总能根据其自身的计算找到数学上的最佳排列，但这并不总是能转化为实际硬件上的最快速度。在某些情况下，新方法产生的结果在理论上是完美的，但在实际表现上却不如旧有的简单规则。研究人员将这种差异归因于成本模型本身。用于预测任务耗时的软件并不完全准确；它低估了某些类型数据移动所需的时间。由于新方法非常擅长寻找其自身存在缺陷的预测下的最低成本，它有时会选择一个在纸面上看起来很便宜、但在现实中却很昂贵的排列方式。这一发现表明，未来改进的最大障碍不是更好的搜索算法，而是更好的预测任务实际耗时的方法。
这项工作为该领域提供了一条清晰的路径。它证明了将布局选择视为一个正式的优化问题是一种可行且强大的策略，能够在简单规则失效的地方实现巨大的加速。它还阐明了性能的最终极限不在于寻找最佳解的能力，而在于引导搜索的预测准确性。对于具有规则、可预测结构的模型，这种基于求解器的算法已经是更优的选择。对于更加混乱和复杂的模型，重点必须转向完善成本模型，使数学上的最优解与芯片的物理现实相一致。通过将“寻找最佳解”与“预测成本”这两个问题分离，研究人员为编译器开发者提供了一个衡量进步的新工具，以及一个明确的下一个努力目标。
您所在领域的论文太多了？
获取与您研究关键词匹配的最新论文每日摘要——附技术摘要，使用您的语言。