arXiv:2502.13575cs.LG2025-02被引 17

通过优化搜索路径减少内存占用,让模型推理更快更省资源。

ETS: Efficient Tree Search for Inference-Time Scaling

  • 用线性规划模型修剪重复路径,提升缓存共享效率。
  • 搜索过程平均缓存大小降低1.8倍,吞吐量提升1.4倍。
  • 适合追求高效推理的开发者,无需改写底层代码。

测试时计算扩展已成为提升模型准确率的新维度,通过在推理阶段增加计算量,使模型有更多时间处理复杂问题。一种有前景的方法是基于过程奖励模型的搜索:模型在每一步生成多个候选路径,由独立的奖励模型对这些部分轨迹进行评分以指导搜索。然而,轨迹多样性会降低键值(KV)缓存共享,导致内存消耗增加、速度变慢。现有方法要么探索不足,要么虽多样但延迟高。本文提出高效树搜索(ETS),通过剪枝冗余路径,在保持必要多样性的同时促进KV共享。ETS引入线性规划成本模型,惩罚保留节点数,同时加入语义覆盖项,确保保留语义不同的路径。实验表明,ETS可使搜索过程中平均KV缓存大小减少1.8倍,吞吐量提升1.4倍,且精度损失极小,无需定制内核实现。代码已开源:https://github.com/SqueezeAILab/ETS。

原文摘要 · Abstract (English)

Test-time compute scaling has emerged as a new axis along which to improve model accuracy, where additional computation is used at inference time to allow the model to think longer for more challenging problems. One promising approach for test-time compute scaling is search against a process reward model, where a model generates multiple potential candidates at each step of the search, and these partial trajectories are then scored by a separate reward model in order to guide the search process. The diversity of trajectories in the tree search process affects the accuracy of the search, since increasing diversity promotes more exploration. However, this diversity comes at a cost, as divergent trajectories have less KV sharing, which means they consume more memory and slow down the search process. Previous search methods either do not perform sufficient exploration, or else explore diverse trajectories but have high latency. We address this challenge by proposing Efficient Tree Search (ETS), which promotes KV sharing by pruning redundant trajectories while maintaining necessary diverse trajectories. ETS incorporates a linear programming cost model to promote KV cache sharing by penalizing the number of nodes retained, while incorporating a semantic coverage term into the cost model to ensure that we retain trajectories which are semantically different. We demonstrate how ETS can achieve 1.8$\times$ reduction in average KV cache size during the search process, leading to 1.4$\times$ increased throughput relative to prior state-of-the-art methods, with minimal accuracy degradation and without requiring any custom kernel implementation. Code is available at: https://github.com/SqueezeAILab/ETS.

推理加速树搜索缓存优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。