MICE for CAT s: Model-Internal Confidence Estimation for Calibrating Agents with Tools (2025)

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 logit lens 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.

## introduction language models are increasingly being used as tool-using agents, where they can generate executable api calls that can change external environments @xcite @xcite . sometimes the generated tool calls are relatively safe, and mistakes will have minimal impact (e.g., if "how many grand slams has serena williams won?" resulted in the incorrect tool call tennis_reference_count_grand_slams( @xmath0, then the user would just be misinformed). but other times, incorrect tool calls can be more harmful (e.g., if "please remove slash.txt" resulted in the incorrect tool call cli(@xmath1, then the user would lose the contents of their filesystem). a confidence estimator estimates the probability that another model's output is correct. a simple confidence estimator for a language model would be based on the probability that the model itself assigns to its output (i.e., the product of token probabilities) or to its output's semantic equivalence class @xcite ). yet prior work has shown that this method can be poorly calibrated @xcite @xcite . a probabilistic classifier is well calibrated if on an unseen test distribution, it is correct about as often as it thinks it is @xcite @xcite @xcite . for example, of those unseen examples that it predicts to be positive with ≈ 25% probability, ≈ 25% really are positive. well-calibrated probabilities can be used to guide downstream decisions, but calibration should never be one's only engineering target, as even a highly unsure classifier may be well-calibrated (see §3.1). to that end, we introduce a class of modelinternal confidence estimators (mice) and an end-to-end metric, expected tool-calling utility (etcu), to evaluate a tool-calling agent that consults a confidence estimator to decide when to launch the predicted tool call. foot_0 mice extracts features by decoding from the intermediate layers of a transformer-based large language model (llm) and computes the similarities of those generations to the output of the final layer. based on these features and the llm's raw confidence, it learns a model that outputs a confidence score. mice excels on etcu, increasingly outperforming two strong baselines as the cost of incorrect tool calls increases, without increasing calibration error. this paper makes the following contributions: we propose a class of model-internal confidence estimators (mice) that are empirically wellcalibrated on the task of assessing generated tool calls ( §2). we introduce a new metric, expected tool-calling utility, that combines accuracy and calibration to better evaluate tool-calling agents ( §3). finally, we show that mice is sample-efficient and can generalize to new tools, even in a zero-shot setting ( §5). ## model-internal confidence estimators mice is a simple learned probabilistic classifier whose features are derived from model-internal activations. prior work on understanding the internals of transformer language models has shown that intermediate layers at different depths encode different types of information, and that the activation spaces at various layers of these models can be nudged to generate sequences in targeted ways @xcite @xcite @xcite . decoding from the layers of a transformer language model has provided insight into the underlying mechanisms and has been used in early-exit algorithms for faster generation @xcite @xcite . for question answering tasks, decoding from roughly the first half of the layers of the language model produces unintelligible results, but in later layers the model's predictions slowly refine into a plausible answer @xcite . we hypothesize that features from intermediate layers' hidden states could provide useful signal for calibration. drastic changes in the final few layers could indicate the inability for the llm to pinpoint a tool call. as a result, we may trust a prediction that was slowly refined into an answer over the final 50% of layers more than one that drastically changed in the final few layers, even if they had the same distribution at the end. bertscore features since we hypothesize that intermediate layers' hidden states could be a useful signal for calibration beyond the confidences derived from the final layer, we decode from each layer, much as in logit lens (nostalgebraist, 2020), a model interpretation technique. we first decode the output string y at temperature 0. this is the usual way to obtain model output in a task like tool calling. then at each layer @xmath2, we obtain a preliminary output string y (i) of the same length by per-token argmax decoding: foot_2 we then compute the bertscore @xcite between y and each y (i) . these become the main input features to the mice model. 3 raw confidence feature we also integrate the raw confidence of the language model in generating the tool call as a feature to the mice model. we calculate this by computing the product of the probabilities of the tokens in the generated tool call. we figure 3 : bertscore similarities between the generated string y and the preliminary strings y (i) from earlier layers, for llama3-8b-instruct on the ste validation set @xcite . see also figure 8 in appendix a. ## metrics perhaps the most widely reported calibration metric is expected calibration error ( §3.1). as mentioned in the introduction, however, minimizing ece should not be our only goal. we also introduce a utility metric, expected tool-calling utility ( §3.2), to assess the performance of a simple agent that makes call/no-call decisions by using our well-calibrated confidence estimates. this metric is parameterized by the cost of false positives relative to the reward of true positives. ## experiments we now look at training mice and using it at test time to measure both smooth expected calibration error (smece) and expected tool-calling utility. ## analysis what is generated by decoding from intermediate layers? here, we look at what the llm generates from intermediary layers. using the logit lens, we find that models slowly evolve their predictions throughout the layers to get closer to the final output generation. figure 2 shows sample generations. qualitatively, the first two-thirds of the layers tend to generate seemingly random strings. after this point, the generations get increasingly closer to the final generation, but significant refinement still occurs in the final layer. the box-and-whisker plot in figure 3 shows that bertscore tends to increase with layer number. figure 8 in appendix a shows that at some layers, the distribution of bertscores tends to be shifted slightly higher on correct outputs, providing signal to the classifier. what is learned? to better understand how the mice features are used, we examine our mice models trained on the ste dataset with llama3-8b-instruct. for mice rf we calculate gini coefficients, and for mice lr we analyze the feature weights, as suggested by reviewers. figures 5 and 6 indicate that confidence is the most important feature in both mice models: roughly 3 times as important as other features in mice rf and 2 times as important as other features in mice lr. 10 there is no obvious other pattern in the estimated weights, and it is possible that they are underdetermined. feature ablations to better understand which features contribute most to confidence estimation, we performed feature ablations for both mice models for the three llms in our study. in addition to the original setting with all features, we tested four new settings: confidence only; 11 first 10 perhaps calibrated confidence would have worked even better as a feature. ## related work model internals @xcite show that different layers of models encode different aspects of the classical nlp pipeline. moreover, intermediate layer activations can be nudged via steering vectors to control output generations @xcite @xcite . the activation spaces of models are relatively well-formed and there exist directions in these latent spaces that correlate with interpretable properties @xcite . these act as part of the basis for our hypothesis that the model internals could contain a trustworthiness signal, although we did not attempt to discover specific directions in these spaces. intermediate decoding we can view language models with multiple layers as doing iterative inference, where each successive layer refines the predictions of the previous layer. with this lens, decoding from intermediate layers provides signal albeit noisy: the first half of the layers generate uninterpretable text, but after this predictions refine toward a plausible answer @xcite @xcite . other work has focused on inference efficiency by early exiting from transformers @xcite @xcite . our work decodes from intermediary layers as a signal for better calibration. calibration prior work has measured the calibration of off-the-shelf models, including neural networks @xcite , large language models @xcite , and semantic parsers @xcite . a line of machine learning work focuses on calibrating binary classifiers while conditioning only on their predicted confidence. platt scaling transforms a real-valued output (like that of an svm classifier) into probabilities using logistic regression @xcite , which is proven to be equivalent to beta calibration up to preprocessing @xcite . isotonic regression @xcite is a non-parametric approach that learns a best fit to data making only a monotonic non-decreasing assumption. hres are popular, and there has been work on adaptive binning strategies @xcite . we chose hre and nwkr as strong baselines from this class of models. mice lr could be viewed as an extension to platt scaling because mice conditions on model internals in addition to the original confidence. ## conclusion in this work, we introduce model-internal confidence estimators (mice), which improve the trustworthiness and safety of language models as toolcalling agents. we introduce a new metric, expected tool-calling utility, that combines calibration and usefulness to better evaluate the safety and utility of tool calls. we show that mice matches or beats both regression baselines (hre and nwkr) when measured by smooth ece, and significantly improves expected tool-calling utility, especially in medium and high-risk regimes. finally, we find that mice is sample efficient and can generalize to unseen apis in a zero-shot setting.