Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Junjie Ye, Yilong Wu, Sixian Li, Yuming Yang, Zhiheng Xi, Tao Gui, Qi Zhang, Xuanjing Huang, Peng Wang, Zhongchao Shi, Jianping Fan, Zhengyin Du

TL;DR

TL-Training presents a data-efficient framework to enhance tool-use in LLMs by addressing three core issues: adverse training data effects, uneven token importance, and a constrained error taxonomy. It combines MAE-based adverse-effect mitigation, adaptive key-token weighting (PKT), and a PPO-based reward mechanism to optimize tool invocation. Evaluated on a compact dataset of $1{,}217$ trajectories, the approach matches or surpasses both open- and closed-source baselines across four open-source test sets, while also improving robustness to noisy inputs and preserving general task performance. This work offers a scalable paradigm for tool-use training that enables smaller models to achieve strong tool interaction capabilities with limited data.

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.

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

TL;DR

TL-Training presents a data-efficient framework to enhance tool-use in LLMs by addressing three core issues: adverse training data effects, uneven token importance, and a constrained error taxonomy. It combines MAE-based adverse-effect mitigation, adaptive key-token weighting (PKT), and a PPO-based reward mechanism to optimize tool invocation. Evaluated on a compact dataset of trajectories, the approach matches or surpasses both open- and closed-source baselines across four open-source test sets, while also improving robustness to noisy inputs and preserving general task performance. This work offers a scalable paradigm for tool-use training that enables smaller models to achieve strong tool interaction capabilities with limited data.

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.

Paper Structure

This paper contains 37 sections, 2 theorems, 22 equations, 5 figures, 14 tables.

Table of Contents

  1. Introduction
  2. Related Works
  3. Training LLMs for Tool Use
  4. Evaluating LLMs in Tool Use
  5. Preliminaries
  6. Task Formulation
  7. Data Analysis
  8. Performance Analysis
  9. Approaches
  10. Mitigating Adverse Effects
  11. Prioritizing Key Tokens
  12. Introducing a Reward Mechanism
  13. Experimental Setup
  14. Dataset
  15. Baselines
  16. ...and 22 more sections

Key Result

Theorem A.1

During the SFT stage for LLM in tool use, gradient updates resulting from incorrect interaction paths in the training data can adversely impact the model's ability to choose the appropriate tool.

Figures (5)

  • Figure 1: Error statistics for various tool calls in RoTLLaMA's training data.
  • Figure 2: Framework of TL-Training. TL-Training comprises three main components: (Left) mitigating the adverse effects of suboptimal data by identifying erroneous interaction trajectories through tool feedback and blocking their gradient updates; (Middle) optimizing key tokens by dynamically adjusting token weights during the SFT process; and (Right) enhancing tool call performance through a reward mechanism tailored to tool invocation error types, using the PPO algorithm for reinforcement learning.
  • Figure 3: Types of errors encountered by LLMs during tool use and their corresponding feedback messages.
  • Figure 4: Performance comparison of RoTLLaMA and TL-CodeLLaMA-2 in different noise environments. RoTLLaMA's results are from RoTBench.
  • Figure 5: Performance comparison of CodeLLaMA-2 and TL-CodeLLaMA-2 across various general tasks.

Theorems & Definitions (4)

  • Theorem A.1
  • proof
  • Theorem A.2
  • proof