Uncertainty Calibration for Tool-Using Language Agents (2024)

Abstract: There is increasing interest in equipping language models with the ability to leverage external tools for complex, goal-oriented tasks. However, interacting with external tools introduces inherent uncertainties due to imperfections and misalignments between the tools’ outputs and the agents’ internal models, often leading to suboptimal outcomes. We thus study the problem of tool-use calibration in language agents, and identify prompt design and execution trace selection as two primary areas that suffer from miscalibration. We then propose ProbeCal, which recalibrates the internal probabilities of tool-using language agents to better reflect the actual effectiveness of tool, and enables a more appropriate selection of prompts and execution paths. We empirically show that ProbeCal can significantly and consistently improve off-the-shelf language models in tool-using applications.

## introduction language agents are increasingly used to perform tasks and interact with a variety of external tools to achieve specific, goal-oriented objectives @xcite . recent advancements have focused on designing agents to more effectively utilize a diverse set of tools for different tasks @xcite @xcite . tool-using agents have also been applied to many domains, such as employing retrieval tools for multi-hop reasoning @xcite , leveraging visual perception models for visual question answering @xcite , and utilizing calculators for mathematical reasoning @xcite . despite these advances, tool use poses inherent uncertainties for language agents. tools may be imperfect and produce erroneous outputs, leading to suboptimal interactions and outcomes. for instance, in tasks requiring knowledge retrieval, the quality of the retrieval tool can significantly affect the agent's performance. if the retrieval tool fails to fetch accurate or relevant information, the agent's ability to reason and generate correct responses is compromised. in visual question answering tasks, @xcite demonstrated that even with the correct algorithms, the performance bottleneck often lies in imperfect visual perception. the above challenges highlight the need for systematically modeling and managing the uncertainties associated with tool-use in language agents. existing approaches generally rely on the agent's internal probabilities to choose which tools to use and how to use them (e.g., @xcite ; @xcite ). however, because these tools are typically not involved during the agent's training, their operation and outputs can be misaligned with how the agent reasons using the tools' outputs. in this paper, we study this problem of tooluse calibration, where the goal is to calibrate a language agent's internal probability estimates to more accurately reflect the actual success rates of tool-using decisions. as illustrated in figure 1 , we identify two primary scenarios where such misalignment can be significantly mitigated through uncertainty calibration: language agent prompts and tool execution traces. our key technical contribution is a set of calibration methods called probing calibration (probcal), which draws inspiration from probing literature @xcite , involves training a classifier using large language model embeddings to predict the expected reward for a given prompt or execution trace. we also explore other design choices such as reweighting during training and temperature scaling @xcite during inference to further refine the calibration of language agents, and provide a systematic analysis of different techniques. we conduct experiments using various types of tool-using language agents, including both static and dynamic agents, depending on whether a fixed figure 1 : overview of our probecal framework for calibrating tool-using language agents. the process begins with a given task, exemplified by calculating pay after taxes from a pay stub. the agent is provided with a user prompt to utilize its tool-using capabilities, generating multiple execution traces. these traces involve different combinations of tool applications, such as row lookup, column lookup, and calculations. however, integrating user prompts and external tools introduces uncertainties, often resulting in misalignment between the agent's internal confidence levels and actual success rates. the framework addresses these issues by calibrating the agent's confidence levels with actual performance, based on which we select the most suitable prompt and execution trace. the calibration model is implemented as an mlp with llm embeddings as inputs and calibrated probabilities as outputs, and is trained with the execution result supervision. the figure is adapted from @xcite . or dynamic set of tools is available. our results on the math dataset @xcite and the tabmwp dataset @xcite demonstrate that off-the-shelf language models are not well-calibrated for tool-using applications. our probecal approach significantly improves both the calibration and performance of these models consistently across different tasks. for example, applying probecal results in over 7% improvement on tabmwp in dynamic tool-using settings compared with strong baselines. furthermore, our analyses reveal that our approach can generally perform well without the need for the llm logit or posthoc calibration, such as temperature scaling, which can translate into simpler overall frameworks for tool-using language agents. ## tool-use calibration for language agents tool-using language agents. we define a toolusing language model as an agent that can interact with external environment such as program compilers via a defined set of function interfaces @xmath0, . . . , f k }. each function f i in f represents a callable program that the lm can utilize to perform specific tasks or computations external to the lm itself. given a textual query q and the toolbox f, the language model synthesizes a program @xmath1, f)), where π(q, f) is the prompt generated from both q and f, and fed to the lan-guage agent; lm is the language agent being able to interpret natural language queries and generate corresponding executable code. the generated program z is then executed in the given environment e by an execution engine ϕ, such that the result of the execution is given by @xmath2, z). following previous work @xcite , the programs @xmath3. the execution engine ϕ is capable of handling various types of inputs x (e.g., images, videos) and producing outputs r of different types (e.g., text, image crops). in practice, multiple plausible candidate programs {z 1 , . . . , z n } are generated from the language agent, with each program z i being a tool execution trace, and decoding strategies such as beam search and self-consistency @xcite are used to select the best program from this candidate set. ## probing calibration (probecal) as discussed, the performance of tool-using language agents is influenced by the variability in tool information prompts and the complexity of determining the correct tool execution traces. in this part, we introduce our proposed methods to calibrate these uncertainties and use the calibrated scores to enhance the performance of language agents. inspired by the literature on probing, where a separate diagnostic model is trained on top of a pretrained language model's internal representations to study linguistically-informed properties @xcite , we employ the training of a classification model that takes inputs as the embeddings from a large language model and to predict the expected reward given a prompt or execution trace. reward estimation. given a dataset containing questions and their corresponding answers, we feed models with different candidate prompts or sample multiple execution traces to collect the corresponding embeddings from a language agent and the final reward r for each prompt or trace. for example, in question-answering settings, the reward r equals 1 if the final result matches the ground truth answer and 0 otherwise. formally, given the embedding ϕ from the language model and the corresponding reward r, reward calibration involves training a reward classification calibration model to learn the mapping p (r|ϕ) from ϕ to r. in our setting, we use a three-layer mlp as the model. we propose to use the embedding-based reward estimation to calibrate the language model, which we call probing calibration (probcal). we apply the embedding-based reward estimation to the prompt design setting, (prompt calibration (probecal-prompt)), the trace selection setting (trace calibration (probecal-trace)), and both of the settings (prompt and trace calibration (probecal-prompt&trace)) as follows. prompt selection calibration. under our framework, we consider the problem of calibrating the probability of selecting which prompt to choose for different questions. formally, for each question q, we want to select each prompt π ∈ p proportional to the estimated reward p (r|ϕ) obtained from a calibration method, which we will illustrate in the next section. this is related to a classic method for decision-making under uncertainty called boltzmann exploration, where each action is selected proportional to the exponential of the expected reward divided by a temperature parameter @xcite . execution trace selection calibration. many existing strategies rely on language agents to estimate the success probability, which can be unreliable since these tools are external to the agents. we explicitly calibrate the success probability and compute the answer by estimating the expected reward p (r|a, z i , q) using a calibration model: arg max a p(a|q) = arg max a z i p(r|a, z i , q). this approach is related to the literature on confidence-weighted majority voting @xcite @xcite , which posits that the theoretically optimal method for aggregating individual confidences is confidence-weighted majority voting, where more reliable individual responses are given greater weight. ## conclusion this paper addresses the inherent uncertainties in tool-using language agents, crucial for performing complex tasks. despite advancements, challenges persist due to variability and errors from integrating external tools. our framework recalibrates probability estimates for tool-using decisions to improve robustness and effectiveness. we focused on two primary areas of miscalibration: language agent prompts and tool execution traces. we systematically addressed these uncertainties by training a classification model using embeddings from a large language model and employing complementary methods like reweighting during training and temperature scaling during inference. experiments on the math and tabmwp datasets demonstrate improvements in calibration and performance, underscoring the effectiveness of our approach. our work fills a critical research gap, enhancing the reliability and accuracy of tool-using language agents.