TCL: Enabling Fast and Efficient Cross-Hardware Tensor Program Optimization via Continual Learning

📄 arXiv: 2604.12891v1 📥 PDF

作者: Chaoyao Shen, Linfeng Jiang, Yixian Shen, Tao Xu, Guoqing Li, Anuj Pathania, Andy D. Pimentel, Meng Zhang

分类: cs.LG, cs.AR

发布日期: 2026-04-14

备注: introduces TCL framework for cross-hardware tensor program optimization with active learning, Mamba-based cost model, and continual knowledge distillation; includes extensive experiments on CPU and GPU platforms

💡 一句话要点

TCL：通过持续学习实现快速高效的跨硬件张量程序优化

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 张量程序优化 深度学习编译器 主动学习 Mamba模型 知识蒸馏

📋 核心要点

  1. 现有深度学习编译器依赖大型离线数据集进行张量程序优化，导致数据收集成本高昂，且跨平台迁移性不足。
  2. TCL框架通过RDU采样器降低数据收集成本，使用Mamba模型捕获长程依赖，并利用持续知识蒸馏实现跨平台知识迁移。
  3. 实验表明，TCL在CPU和GPU上优化主流深度学习模型时，调优时间显著减少，推理延迟也得到降低。

📝 摘要（中文）

本文提出了一种名为TCL的高效且可迁移的编译器框架，用于在不同硬件平台上快速优化张量程序，旨在解决现有方法依赖大型离线数据集导致的高收集成本和跨平台迁移性差的问题。TCL基于三个核心组成部分：(1) RDU Sampler，一种数据高效的主动学习策略，通过联合优化代表性、多样性和不确定性，仅选择10%的张量程序，显著降低数据收集成本，同时保持接近原始模型的准确性；(2) 一种新的基于Mamba的成本模型，通过减少参数化和轻量级序列建模，有效地捕获长程调度依赖性，并在预测精度和计算成本之间实现良好的权衡；(3) 一个连续知识蒸馏框架，有效地逐步地在多个硬件平台之间迁移知识，同时避免了传统多任务学习通常引起的参数爆炸和数据依赖问题。大量实验验证了每个独立组件和整体TCL框架的有效性。在CPU和GPU平台上优化一系列主流深度学习模型时，与Tenset-MLP相比，TCL平均实现了16.8倍和12.48倍的更快调优时间，以及1.20倍和1.13倍的更低推理延迟。

🔬 方法详解

问题定义：现有深度学习编译器在优化张量程序时，依赖于大量的离线数据进行训练，这导致了高昂的数据收集成本。此外，针对特定硬件平台训练的模型，在迁移到其他平台时，性能往往会显著下降，即缺乏良好的跨平台迁移性。因此，如何降低数据收集成本，并提升跨平台优化能力，是本文要解决的核心问题。

核心思路：TCL的核心思路是利用数据高效的主动学习策略、轻量级的序列建模成本模型和持续知识蒸馏框架，来实现快速且高效的跨硬件张量程序优化。通过RDU采样器，减少了数据收集量；通过Mamba模型，提升了成本模型的预测精度和效率；通过持续知识蒸馏，实现了跨平台的知识迁移，避免了参数爆炸和数据依赖问题。

技术框架：TCL框架主要包含三个核心模块：RDU Sampler、Mamba-based Cost Model和Continuous Knowledge Distillation Framework。首先，RDU Sampler负责从大量的张量程序中选择最具代表性、多样性和不确定性的样本，用于训练成本模型。然后，Mamba-based Cost Model利用选定的样本进行训练，预测不同张量程序在目标硬件上的性能。最后，Continuous Knowledge Distillation Framework利用知识蒸馏技术，将已训练好的模型知识迁移到新的硬件平台上，实现跨平台优化。

关键创新：TCL的关键创新在于三个方面：(1) 提出了RDU Sampler，一种数据高效的主动学习策略，能够显著减少数据收集成本；(2) 提出了基于Mamba的成本模型，能够有效地捕获长程调度依赖性，并在预测精度和计算成本之间实现良好的权衡；(3) 提出了连续知识蒸馏框架，能够有效地逐步地在多个硬件平台之间迁移知识，避免了传统多任务学习的参数爆炸和数据依赖问题。与现有方法相比，TCL在数据效率、预测精度和跨平台迁移性方面都具有显著优势。

关键设计：RDU Sampler通过联合优化Representativeness、Diversity和Uncertainty三个指标来选择样本。Representativeness衡量样本的代表性，Diversity衡量样本的多样性，Uncertainty衡量模型对样本预测的不确定性。Mamba-based Cost Model采用Mamba架构，通过选择性状态空间模型来捕获长程依赖关系。Continuous Knowledge Distillation Framework采用逐步蒸馏的方式，将知识从一个平台迁移到另一个平台，避免了一次性蒸馏带来的信息损失。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，TCL在CPU和GPU平台上优化主流深度学习模型时，与Tenset-MLP相比，平均实现了16.8倍和12.48倍的更快调优时间，以及1.20倍和1.13倍的更低推理延迟。这些结果验证了TCL框架在提升调优效率和降低推理延迟方面的有效性。

🎯 应用场景

TCL框架可应用于各种深度学习编译器的后端优化，提升模型在不同硬件平台上的性能。该研究成果对于加速深度学习模型的部署和推理，降低计算成本，以及推动深度学习在边缘设备上的应用具有重要意义。未来，TCL可以扩展到更多类型的硬件平台和深度学习模型，进一步提升其通用性和实用性。

📄 摘要（原文）

Deep learning (DL) compilers rely on cost models and auto-tuning to optimize tensor programs for target hardware. However, existing approaches depend on large offline datasets, incurring high collection costs and offering suboptimal transferability across platforms. In this paper, we introduce TCL, a novel efficient and transferable compiler framework for fast tensor program optimization across diverse hardware platforms to address these challenges. Specifically, TCL is built on three core enablers: (1) the RDU Sampler, a data-efficient active learning strategy that selects only 10% of tensor programs by jointly optimizing Representativeness, Diversity, and Uncertainty, substantially reducing data collection costs while maintaining near-original model accuracy; (2) a new Mamba-based cost model that efficiently captures long-range schedule dependencies while achieving a favorable trade-off between prediction accuracy and computational cost through reduced parameterization and lightweight sequence modeling; and (3) a continuous knowledge distillation framework that effectively and progressively transfers knowledge across multiple hardware platforms while avoiding the parameter explosion and data dependency issues typically caused by traditional multi-task learning. Extensive experiments validate the effectiveness of each individual enabler and the holistic TCL framework. When optimizing a range of mainstream DL models on both CPU and GPU platforms, TCL achieves, on average, 16.8x and 12.48x faster tuning time, and 1.20x and 1.13x lower inference latency, respectively, compared to Tenset-MLP.