用蒙特卡洛树搜索优化张量计算图，提升模型推理速度
Optimizing Tensor Computation Graphs with Equality Saturation and Monte Carlo Tree Search
- 结合等价饱和与蒙特卡洛树搜索，智能选择最优重写规则
- 在内存受限下仍实现最高11%的推理加速
- 适合需要高效推理的深度学习部署场景
深度神经网络的实际性能常受延迟影响，需优化推理时间而不损失精度。现有方法通过顺序替换子图来重构计算图，但易陷入相位排序问题——早期重写可能阻断后续更优操作。近年研究提出等价饱和技术，先构建中间表示（IR）存储多个优化版本，再从中选出最佳方案，缓解该问题。然而实际中内存限制导致无法捕获所有优化版本，重建了构造阶段的相位排序难题。本文提出一种张量图重写方法，利用蒙特卡洛树搜索识别最有潜力的重写规则，构建更优的中间表示；并设计新型提取算法，可快速准确估算IR中张量程序的运行时性能。相比现有方法，本方案使神经网络推理速度提升最多达11%。
原文摘要 · Abstract (English)
The real-world effectiveness of deep neural networks often depends on their latency, thereby necessitating optimization techniques that can reduce a model's inference time while preserving its performance. One popular approach is to sequentially rewrite the input computation graph into an equivalent but faster one by replacing individual subgraphs. This approach gives rise to the so-called phase-ordering problem in which the application of one rewrite rule can eliminate the possibility to apply an even better one later on. Recent work has shown that equality saturation, a technique from compiler optimization, can mitigate this issue by first building an intermediate representation (IR) that efficiently stores multiple optimized versions of the input program before extracting the best solution in a second step. In practice, however, memory constraints prevent the IR from capturing all optimized versions and thus reintroduce the phase-ordering problem in the construction phase. In this paper, we present a tensor graph rewriting approach that uses Monte Carlo tree search to build superior IRs by identifying the most promising rewrite rules. We also introduce a novel extraction algorithm that can provide fast and accurate runtime estimates of tensor programs represented in an IR. Our approach improves the inference speedup of neural networks by up to 11% compared to existing methods.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。