Towards Uncertainty-Aware Language Agent (2024)

Abstract: While Language Agents have achieved promising success by placing Large Language Models at the core of a more versatile design that dynamically interacts with the external world, the existing approaches neglect the notion of uncertainty during these interactions. We present the Uncertainty-Aware Language Agent (UALA), a framework that orchestrates the interaction between the agent and the external world using uncertainty quantification. Compared with other well-known counterparts like ReAct, our extensive experiments across 3 representative tasks (HotpotQA, StrategyQA, MMLU) and various LLM sizes demonstrate that UALA brings a significant improvement of performance, while having a substantially lower reliance on the external world (i.e., reduced number of tool calls and tokens). Our analyses provide various insights including the great potential of UALA compared with agent fine-tuning, and underscore the unreliability of verbalised confidence of LLMs as a proxy for uncertainty.

## introduction language agents @xcite @xcite utilise large language models @xcite @xcite to interact with the external world (e.g., through tools) and to process collected observations towards solving a task. having achieved improvements in previously challenging reasoning tasks @xcite @xcite , these agents autonomously obtain new knowledge from the world, and leverage memory or self-refinement mechanisms to iteratively improve their reasoning trajectories @xcite @xcite . current language agent designs primarily define the role of an llm as the planner, moderator, or aggregator of observations and knowledge collected from the external world (i.e., via tools and apis), overlooking the potential of llms to tackle numerous tasks without depending on any external tool and resource @xcite @xcite . while exploiting external mechanisms has an advantage in terms of verifiability of outputs (e.g., output extracted from a retrieved wikipedia article provides a verifiable reference) or improvement on accuracy (e.g., using a python program to do mathematical calculation instead of relying on the llm), this design is inherently wasteful in how it utilises the llm's internal implicit knowledge. in fact, the field of nlp has shown great interest in probing pretrained language models to highlight various types of knowledge they embody directly in their weights (e.g., see @xcite ; @xcite ; @xcite and references therein). we postulate that a more effective design for language agents should have a better interplay between the implicit knowledge encoded in llm's weight and the explicit knowledge of the external world. drawing upon principles from decision theory, an intelligent agent should possess a reliable mechanism for measuring uncertainty of its own proposition or the observations from the environment @xcite . we build on this, and move towards equipping language agents (as a special case of intelligent agents) with the means of measuring uncertainty to efficiently regulate their use of external help (e.g., tools, knowledge bases). in nlp, in general, there has been a great amount of research on uncertainty estimation for generative tasks such as machine translation @xcite , summarisation @xcite , and data-to-text generation @xcite . recently, a dedicated space emerged to study the intersection of llms and uncertainty. this space can be divided into approaches applicable to open-box llms which rely on token-level logits @xcite , and those applicable to black-box llms leveraging multiple response generation as a proxy for model's confidence @xcite , or relying on llm's verbal self-awareness of confidence @xcite . nonetheless, uncertainty in the context of language agents is an under-explored space, with our work being the first in this direction. we propose an uncertainty-aware language agent framework which utilises various existing uncertainty measurement methods as a dynamic switch between llm's own trajectory or resorting to external resources (i.e., search engine, wikipedia knowledge base) during the course of a reasoning task. we study the properties of this framework on a large space of llms (chat-gpt; openai (2023), and llama2 70b, 13b, 7b; @xcite ), and language agent frameworks such as react @xcite , and critic @xcite on several representative tasks: hotpotqa @xcite a free-form qa task; strategyqa @xcite a representative binary qa task; and 57 multiple-choice qa tasks of mmlu @xcite . our key findings and contributions are as follows: we propose the first language agent framework which integrates uncertainty in the reasoning trajectories. the proposed design allows for a plugand-play use of various uncertainty estimation techniques. we highlight across our wide spectrum of experiments that integrating uncertainty not only leads to a significant performance improvement, it also has several practical benefits such as a substantial reduction in external calls and number of output tokens (i.e., less api charge). our analysis reveals several additional findings: (1) the gain is significant across different llm sizes and tasks. (2) uncertainties have a degree of calibration because there exists a divergence of uncertainty between correct and incorrect answers. (3) the verbalised confidence of llms is an unreliable proxy for uncertainty quantification. (4) given a limited amount of data, leveraging uncertainty leads to a much higher performance improvement compared to fine-tuning language agents. ## uncertainty-aware language agent we present uncertainty-aware language agent, uala, a framework that integrates uncertainty in language agent's cycle of thought, action, and observation. the uncertainty moderates the interaction between the llm and the external world, facilitating a more effective and efficient dynamic. given an input question, the agent's potentially goes through 3 modes of attempt, each involving a thought, an action, and an observation (similar to react). figure 1 presents real trajectories for each mode of attempt. in our approach, as illustrated in example (a), we initially generate a response to a query using standard or chain-of-thought (cot; @xcite ) prompting. this response is then assessed by an uncertainty estimation module, which determines the subsequent action between two courses: either accepting the response as-is, or employing external tools for a more certain solution, as shown in example (b). this process is similarly applied when using tool-activated responses. if the uncertainty remains high post-tool engagement, the agent seek assistance from an external source such as a human, as depicted in example (c). this protocol ensures a more robust and reliable response generation by actively addressing uncertainty at each step. next, we describe our uncertainty estimation mechanisms. ## experiments we explore three types of qa tasks and utilise wikipedia and google search and external tools: hotpotqa @xcite is a challenging qa dataset which requires multi-hop reasoning over some wikipedia passages. the answer of the questions in hotpotqa is free-form with an arbitrary length. we randomly select 500 questions from the training set for the creation of calibration set, and 500 questions from the dev set for evaluation using the same random seed as react @xcite . strategyqa @xcite is an opendomain qa dataset which requires implicit reasoning steps. it only contains binary answers (yes or no). we randomly select 500 questions from the training set to create the calibration set. for the evaluation, we use the same dev set used in fire-act @xcite containing 229 questions. mmlu @xcite is a multi-choice qa dataset in various domains which covers 57 tasks including elementary mathematics, us history, computer science, law, and more. this dataset is to evaluate model's academic and professional understanding. each question contains four options and only one option is correct. to cover all different topics, we randomly sample 10 questions for each task from dev set to create the calibration set, and similarly sampled 10 question for each task (total of 570 questions) from test set for evaluation. wikipedia. following the wikipedia web api setting used in react @xcite , the tool prompt contains three actions. (1) search[entity], which returns the initial five sentences from the entity wiki page if it exists. otherwise, it suggests the top five similar entities from the wikipedia search engine. (2) lookup[string], which functions akin to the ctrl+f in a browser, returning the subsequent sentence on the page containing the specified string. ( 3 ) finish[answer], which contains an answer and finishes the current task. this tool is used to retrieve some relevant knowledge of the entities in the query, which helps the agent to do the reasoning based on the external knowledge and its internal knowledge. the wikipedia tool is used for hotpotqa and strategyqa datasets. the universality of google search. is better than wikipedia. it can search any query which is suitable for more complex reasoning tasks. therefore, we adopt google search tool for mmlu task. following fireact @xcite , the google search tool we use is based on serpapi foot_3 . the tool prompt contains two actions: (1) search[query], which searches a query on google and returns a short snippet containing the answer of the query. the returned snippet is the first existent item from "answer box", "answer snippet", "highlight words", or "first result snippet", which ensures the response is short and relevant. ( 2 ) finish[answer], which returns the answer and finishes the episode. ## conclusion we present the uncertainty-aware language agent (uala), which improves a language agent's efficiency in interaction with the external world to answer various qa problems. we first measure the uncertainty of an agent on a small set of examples and establish an acceptable uncertainty threshold. then the uncertainty estimation function is integrated into a language agent's trajectory of think-act-observe by measuring the uncertainty of generated answers, and then choosing to accept the answer or to resort to external resources or human. this approach does not require any training and the language agent is prompted by few-shot learning. we find that this technique significantly improves the language agent's performance on various qa tasks (regardless of the backbone llm size), and offers a much more resource-efficient solution (substantially lower tool calling, and less number of output tokens). we also show that uala outperforms fine-tuning language agents in the presence of small amount of data.