MICE for CATs: Model-Internal Confidence Estimation for Calibrating Agents with Tools

TL;DR

The paper addresses the safety of tool-using language models by proposing MICE, a framework that derives model-internal confidence signals from intermediate transformer layers. By decoding each layer and computing BERTScore-based similarities to the final output, combined with raw token-probability confidence, MICE trains a probabilistic classifier to estimate the correctness of tool calls. It also introduces Expected Tool-Calling Utility (ETCU), a metric that integrates calibration with decision-making across risk levels. Empirical results on the STE dataset with Llama3 variants show MICE, especially MICE RF, achieves better ETCU and competitive smECE than strong baselines, with strong zero-shot generalization and notable sample efficiency. The work advances practical, calibrated decision-making for tool use in LLMs and provides open-source code for replication and extension.

Abstract

Tool-using agents that act in the world need to be both useful and safe. Well-calibrated model confidences can be used to weigh the risk versus reward of potential actions, but prior work shows that many models are poorly calibrated. Inspired by interpretability literature exploring the internals of models, we propose a novel class of model-internal confidence estimators (MICE) to better assess confidence when calling tools. MICE first decodes from each intermediate layer of the language model using logitLens and then computes similarity scores between each layer's generation and the final output. These features are fed into a learned probabilistic classifier to assess confidence in the decoded output. On the simulated trial and error (STE) tool-calling dataset using Llama3 models, we find that MICE beats or matches the baselines on smoothed expected calibration error. Using MICE confidences to determine whether to call a tool significantly improves over strong baselines on a new metric, expected tool-calling utility. Further experiments show that MICE is sample-efficient, can generalize zero-shot to unseen APIs, and results in higher tool-calling utility in scenarios with varying risk levels. Our code is open source, available at https://github.com/microsoft/mice_for_cats.

Figures (8)

• Figure 1: The MICE architecture.
• Figure 2: Example generations from the validation set across layers of the Llama3-8B-Instruct model. Generations from early layers (5, 15) are seemingly random, but later layers (25, 31) generate thematically relevant tokens. Layer 32 is the final layer.
• Figure 3: BERTScore similarities between the generated string $\bf y$ and the preliminary strings ${\bf y}^{(i)}$ from earlier layers, for Llama3-8B-Instruct on the STE validation set wang-etal-2024-llms-imaginarium. See also \ref{['fig:bertscores-posneg']} in \ref{['sec:experimental-details']}.
• Figure 4: Expected tool-calling utility on the test set at varying risk levels. We include four trivial policies for reference: oracle executes only when the underlying model is correct (an upper bound); always abstain never executes, getting reward 0; always execute never abstains; and the base rate policy switches from always execute to always abstain when the risk level exceeds the base accuracy. All policies perform similarly at low risk levels, where always execute is close to optimal and hard to improve on. MICE models show clear improvements in the medium and high risk regimes.
• Figure 5: Feature importance for BERTScore features and confidence on the trained MICE RF model on the STE dataset for the Llama3 LLM.