COLT: Lightweight Multi-LLM Collaboration through Shared MCTS Reasoning for Model Compilation

📄 arXiv: 2602.01935v1 📥 PDF

作者: Annabelle Sujun Tang, Christopher Priebe, Lianhui Qin, Hadi Esmaeilzadeh

分类: cs.LG, cs.AI, cs.PL

发布日期: 2026-02-02


💡 一句话要点

COLT:通过共享MCTS推理实现轻量级多LLM协作,用于模型编译优化

🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)

关键词: 多LLM协作 模型编译优化 蒙特卡洛树搜索 轻量级框架 模型选择

📋 核心要点

  1. 现有编译器优化依赖单一大型LLM成本高,小型LLM效果差,缺乏高效的多LLM协作方案。
  2. COLT通过共享MCTS树实现多LLM协作,利用轻量级模型为主,大型模型辅助,降低计算成本。
  3. COLT引入模型感知树策略和路线改变机制,平衡探索与利用,提升搜索效率和优化效果。

📝 摘要(中文)

模型服务成本在AI系统中占据主导地位,因此编译器优化对于可扩展部署至关重要。最近的研究表明,大型语言模型(LLM)可以通过推理程序结构和优化历史来指导编译器搜索。然而,在整个搜索过程中使用单个大型模型成本高昂,而单独使用较小模型则不太可靠。因此,本文旨在探讨主要依赖于小型LLM的多LLM协作推理是否能够匹配或超过单个大型模型的性能。为此,我们提出了一种轻量级协作多LLM框架COLT,用于编译器优化,该框架支持单个蒙特卡洛树搜索(MCTS)过程中多个模型之间的协调推理。一个关键贡献是使用单个共享MCTS树作为LLM之间的协作基底,从而实现转换前缀的重用和跨模型的值传播。因此,我们通过简单地将模型选择内生于轻量级MCTS优化循环中,规避了繁重的内部推理机制和依赖于外部规划器、多个并发LLM、数据库、中间结果的外部存储/版本控制以及控制器的传统代理机制。每次迭代时,执行LLM都会提出一个联合动作:(编译器转换,下一个要查询的模型)。我们还引入了一种模型感知的树策略,该策略偏向于搜索较小的模型,同时保留探索能力,以及一种路线改变机制,当搜索表现出可归因于较小模型的持续回归时,该机制会升级到最大的模型。

🔬 方法详解

问题定义:论文旨在解决模型编译优化过程中,使用单一大型LLM成本过高,而小型LLM效果不佳的问题。现有方法要么依赖于昂贵的单一大型模型,要么无法充分利用多个小型模型的协同能力,导致优化效率和效果受限。

核心思路:论文的核心思路是利用多个LLM进行协作,通过共享的蒙特卡洛树搜索(MCTS)框架,将模型选择内生于优化循环中。优先使用小型LLM进行探索,并在必要时升级到大型LLM进行更精确的评估,从而在成本和性能之间取得平衡。

技术框架:COLT框架包含以下主要模块:1) 共享MCTS树:所有LLM共享同一棵MCTS树,用于存储和传播优化过程中的信息。2) 模型感知树策略:该策略根据模型的性能和成本,动态调整搜索方向,偏向于选择较小的模型。3) 联合动作提议:每个LLM在每次迭代时,提出一个联合动作,包括编译器转换和下一个要查询的模型。4) 路线改变机制:当搜索出现持续的性能下降时,该机制会将查询的模型升级到更大的模型,以避免陷入局部最优。

关键创新:COLT的关键创新在于其轻量级的多LLM协作方式。它避免了复杂的代理机制和外部存储,而是通过共享MCTS树和内生的模型选择,实现了高效的模型协作。此外,模型感知的树策略和路线改变机制进一步提升了搜索效率和优化效果。

关键设计:COLT的关键设计包括:1) 共享MCTS树的结构和更新方式,确保所有LLM能够共享和利用优化信息。2) 模型感知树策略的具体实现,包括如何根据模型的性能和成本来调整搜索概率。3) 路线改变机制的触发条件和升级策略,确保在必要时能够切换到更大的模型。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

论文提出的COLT框架在编译器优化任务上取得了显著的性能提升。实验结果表明,COLT能够以较低的计算成本,达到甚至超过使用单一大型LLM的优化效果。具体性能数据和对比基线在论文中有详细展示。

🎯 应用场景

COLT框架可应用于各种模型编译优化场景,例如深度学习模型的部署优化、硬件加速器的代码生成等。通过降低优化成本,COLT能够加速AI模型的部署,并提升模型在资源受限设备上的性能,具有广泛的应用前景。

📄 摘要(原文)

Model serving costs dominate AI systems, making compiler optimization essential for scalable deployment. Recent works show that a large language model (LLM) can guide compiler search by reasoning over program structure and optimization history. However, using a single large model throughout the search is expensive, while smaller models are less reliable when used alone. Thus, this paper seeks to answer whether multi-LLM collaborative reasoning relying primarily on small LLMs can match or exceed the performance of a single large model. As such, we propose a lightweight collaborative multi-LLM framework, dubbed COLT, for compiler optimization that enables coordinated reasoning across multiple models within a single Monte Carlo tree search (MCTS) process. A key contribution is the use of a single shared MCTS tree as the collaboration substrate across LLMs, enabling the reuse of transformation prefixes and cross-model value propagation. Hence, we circumvent both heavy internal reasoning mechanisms and conventional agentic machinery that relies on external planners, multiple concurrent LLMs, databases, external memory/versioning of intermediate results, and controllers by simply endogenizing model selection within the lightweight MCTS optimization loop. Every iteration, the acting LLM proposes a joint action: (compiler transformation, model to be queried next). We also introduce a model-aware tree policy that biases search toward smaller models while preserving exploration, and a course-alteration mechanism that escalates to the largest model when the search exhibits persistent regressions attributable to smaller models.