TL-Training: A Task-Feature-Based Framework for Training Large Language Models in Tool Use

TL;DR

This paper introduces TL-Training, a novel task-feature-based framework that improves large language models' tool use by dynamically adjusting training focus and incorporating error-aware reward mechanisms, achieving high performance with limited data.

Contribution

The paper presents TL-Training, a new framework that enhances LLM tool use by addressing data issues, emphasizing key tokens, and optimizing error handling, outperforming existing methods with minimal data.

Findings

Achieves comparable or better tool-use performance with only 1,217 training points.

Enhances robustness in noisy environments.

Improves general task performance.

Abstract

Large language models (LLMs) achieve remarkable advancements by leveraging tools to interact with environments, a critical step toward generalized AI. However, the standard supervised fine-tuning (SFT) approach, which relies on large-scale datasets, often overlooks task-specific characteristics in tool use, leading to performance bottlenecks. To address this issue, we analyze three existing LLMs and uncover key insights: training data can inadvertently impede tool-use behavior, token importance is distributed unevenly, and errors in tool calls fall into a small set of categories. Building on these findings, we propose~\emph{TL-Training}, a task-feature-based framework that mitigates the effects of suboptimal training data, dynamically adjusts token weights to prioritize key tokens during SFT, and incorporates a robust reward mechanism tailored to error categories, optimized through proximal policy optimization. We validate TL-Training by training CodeLLaMA-2-7B and evaluating it on four open-source test sets. Our results demonstrate that the LLM trained by our method matches or surpasses both open- and closed-source LLMs in tool-use performance using only 1,217 training data points. Additionally, our method enhances robustness in noisy environments and improves general task performance, offering a scalable and efficient paradigm for tool-use training in LLMs. Code and data are available at https://github.com/Junjie-Ye/TL-Training.