We think your primary concern is an important one, and we would like to point out some ways that we directly address the question: "Are probes simply identifying features of the input, or are they truly predicting the future behavior of the LM?"

Several tasks in our study are impossible to perform by simply identifying features of the input:

We point to the tasks of format following (specifically the 3 bullet points task, which predicts whether a LM will output an answer in exactly 3 bullet points) and the jailbreak avoidance task (which predicts not whether an input text is malicious, but whether or not a language model will fail to abstain on a given malicious request). On the jailbreak avoidance datasets, LLama3-8B, Mistral-7B, and DeepSeek-R1 all have different abstention rates:(75%, 78% and % on the SelfAware dataset) and, in general, do not abstain/fail to abstain on the same input instances. We also see that probes trained on one LM do not transfer to unseen LMs (on the same dataset), i.e. a conformal probe trained and calibrated to provide 90% accuracy on the SelfAware dataset with LLama3-8B achieves only 60% accuracy on the Mistral-7B model. This shows that the probes are not identifying general, fixed features of the input text, but rather identifying model-specific patterns.

Pretrained LMs are incapable of predicting each others behavior in this way (without tuning):

We stick to the jailbreak avoidance task and perform the following experiment:

Llama3-8B is used to generate hidden states, and we then annotate the instance for whether or not Llama3 failed to abstain from a malicious input. We then measure the performance of the following methods: Probe trained on the hidden states to predict whether or not LLama3 will fail to abstain FewShot Prompting of Mistral7B to predict whether or not LLama3 will fail to abstain i.e. without any retraining / finetuning

The performance is as shown below:

As we can see, when we do not fine-tune Mistral for the task, it performs worse than random, and significantly worse than the probes. This suggests that the property we are trying to predict is not one that can be inferred from just the text alone and requires some model-specific information to reliably ascertain.

Directly addressing the suggested experiment:

We finally address the experiment suggested by the reviewer: "A more reasonable baseline would be using one causal LM of similar parameter counts to predict the behavior of another. If such cross-prediction fails completely, I would be more convinced that the LLM is genuinely predicting its own behavior before token generation, and linear probing does capture the exact representations for prediction."

We have shown above that such cross-prediction fails when we are not allowed to train the other causal LM (Mistral); however, we suspect the reviewer was more curious about the case where you train the other causal LM to predict this. We would maintain that in such an experiment, you should not expect the cross-prediction to fail.

This is because, when provided with sufficient training data in the form of (input text, output behaviour), the other causal language model will eventually learn the patterns present in the data, and hence learn the behaviour of the LLama model. This would not be evidence that the probes are simply identifying input features, but rather evidence that with sufficient training data and a sufficiently powerful text classification system (in this case, Mistral with a classification head), you will be able to model the patterns in the data and learn the Language Model's behavior.

The point of our paper is to show that you do not need to distil the Llama model in this way to get a reliable signal on its behavior, because the hidden states provide you with helpful information.

However, we understand the core concern here, and so we perform a modified experiment to demonstrate our point. We compare the probes to Lora-finetuning Mistral in a data scarce scenario: i.e. we train probes and the Mistral model to identify whether or not LLama3 will fail to abstain using only 50 datapoints. The performance is as follows:

While the Mistral finetune is better than few-shot prompting, showing that some pattern is being established, it is worse than the probe. This data-efficient learning ability from the probe shows that the pattern is more readily readable from the hidden states, which shows the core claim of our paper: that hidden states of LMs contain information on how the LM will behave.

On weakness: discuss concurrent work

We thank the author for bringing this related work to our attention. We have added writing for both this and the work that it cites into the related work section of our paper.

On weakness: Probing final input token with unclear motivation

• Our decision to use the final token is motivated by a large body of work[1, 2, 3], which shows that the final token encodes useful information that is relevant for the prediction of a variety of properties of the input and output. Hence, we consider our choice of using the final input token as well motivated.
• We also provide an analysis on the performance of using tokens at different positions, and show that the final token does contain more useful representations for behavior prediction, further solidifying this growing observation in the field
• We understand that there are potentially methods of mechanistic interpretability that could be deployed to understand why the final token encodes the most information, however that is beyond the scope of this paper. Given the page limit, we would ask the reviewer to consider that the other sections we have added are of a more direct relevance to the paper, and we leave a deep understanding of the mechanism of attention, compressing the intent to future work.

On weakness: Lack of analysis of the overhead introduced by collecting hidden states for probe training

1. We would first point out that any method which hopes to learn the patterns of the LMs behaviors require some source of training data, and the cost of generating that data is roughly equal. For example, in order to generate the probe training data, we must run our LM on a training set, collect the hidden states and collect the output. In order to train our BERT baseline, we must collect the input text and the output, a process which naturally generates the hidden states as a result of the forward pass. Hence, the only extra cost of the probe (when compared to the BERT baseline), is that the hidden states must be stored in memory. This is not a significant cost, and is more than offset by the cost of fine-tuning BERT. Additionally, considering that once the probe has been trained, the hidden states can then be deleted, we do not consider this to be a significant overhead.
2. Furthermore, we show (Figure 6) that within 500 training datapoints, probes often approach the consistency and coverage they will attain when trained on the entire dataset. This suggests that the amount of training data required for good probe performance is small, minimizing the cost of training.

On weakness: The authors chose simple linear probes...

This is an interesting question. We experiment with a variety of probe models, comparing Differences in Means, Linear Models, shallow MLPs, and RandomForests.

The Linear and RandomForest approaches have comparable performance, while a shallow MLP proves consistently more capable. We believe that on KnownUnknown and SelfAware, the difference is not drastic enough to enable nuanced behavior detection, however, on MMLU, the improvement is significant. As you mention, even more sophisticated probes have the potential to identify more nuanced behavior, but such a decision would increase computational overhead significantly. In our work, we focus on linear probes primarily for their simplicity, however we do agree that there are several options that strike different points on the compute-performance curve, and there are multiple options in the low-compute regime that can give good performance.

On weakness: It is necessary to discuss generalization against more fundamental distribution shifts (e.g., code vs. text classification).

We would first remember the overall goal of the probe: to preemptively predict a specific behaviour of the LM. When the distribution is shifted so fundamentally that the behaviours you are trying to predict is not well defined, it does not make sense to even try to apply a probe trained on one distribution on the other. For example, if you have trained a probe to preemptively predict which text classification option a LM will choose, there is no sense in applying it to a LM that is trying to synthesise a Python program (as the behaviour of interest, selecting a specific option, does not occur in this distribution). We do not see this as a limitation of our method, as no method that seeks to predict a behaviour can generalise to distributions where that behaviour itself is not well defined. What we have shown is that when this behaviour is well defined (ensured by taking multiple datasets spanning different domains, but having some common task properties), then we see signs of OOD generalisation.

We understand the importance of showing that our probes can generalize to fundamentally unseen domains, and conduct a shifting domain experiment to verify this. We stick within the sentiment prediction task, and use only news-based, generic sentiment datasets to train the probes (twitter-mteb and news-mtc). Then, we test our probes on financial sentiment analysis datasets, a shifted domain. The results show that despite never explicitly training on such financial data, the probes are able to achieve significantly higher than random performance:

[1]: Li, Kenneth, et al. "Inference-time intervention: Eliciting truthful answers from a language model." Advances in Neural Information Processing Systems 36 (2023): 41451-41530.

[2]: Azaria, Amos, and Tom Mitchell. "The internal state of an LLM knows when it's lying." arXiv preprint arXiv:2304.13734 (2023).

[3]: Von Rütte, Dimitri, et al. "A language Model's guide through latent space." arXiv preprint arXiv:2402.14433 (2024).

On question: Practical advantages / potential drawbacks

We think that conformal probes can be a significant boon to organisations hoping to limit the computational cost of their LM inference. By identifying easy instances, catching malicious inputs early and exiting when they have precisely anticipated the LMs output, the probes can greatly reduce time and cost, while giving users exact control over the allowable error rate.

Conformal probes are not without their limitations. Specifically, the mathematical guarantee on the error rate only holds when the test data is exchangeable (similar to IID), and hence if deployed on data that is too different from the training distribution, it may fall below the user's stated performance threshold without notice. We speak more about this below, in an answer to your question on conformal calibration, where we suggest ways to use this method with confidence even in the OOD setting.

On question: latency and overhead of inference efficiency measurement

We thank you for your attention to this detail. We would note our initial motivation for measure % of tokens generated as opposed to wallclock time: we wanted a measure of efficiency that is agnostic to specific hardware implementations, and that captures the reason why the efficiency of the probe-based generation is superior. However, we understand your concerns that the empirical reality may not match up with this more conceptual metric, hence we measure the wallclock time of normal inference v.s. Inference with conformal probes guided early exiting on Llama3-8B using alpha=0.8, and present the results in the table below:

As we can see, the inference cost reductions are significant even when measured by wallclock time. Importantly, the inference time reduction roughly matches the % reduction in tokens generated, which is the original metric we report. This confirms that empirical results match our more hardware-agnostic measures of inference efficiency.

On question: Calibration stability across domains

This is an important question, and we would like to answer it from two perspectives:

1. Theoretically, no, the learned conformal thresholds do not apply under shifting domains. This is intuitive as the formal guarantee that the threshold will bound the error rate depends on the test dataset being exchangeable (similar to IID) with the calibration set. This does not hold in the OOD case, and hence, unless we add assumptions on the statistical nature of the distribution shift (which is hard to do for abstract concept shifts like code v.s. medical text) we cannot guarantee a bound on the error rate.
2. Emperically, we see optimistic results. While there is significant variance across tasks and datasets, we observe that the calibration threshold still limits the error rate when used in the ood setting. To see this, refer to Figure, which uses alpha=0.9 and hence hopes to limit the error rate to 10%. We see that across 18 datasets, this threshold is held even in the OOD setting. However, from our experiments we are able to give the following recommendations for practitioners who hope to use conformal probes in an OOD setting: a. As usual, the training set that is closest in approximation to the testing set will give the most secure guarantees for the threshold. This is seen in the case of the sentiment datasets: we use only news-based, generic sentiment datasets to train the probes (twitter-mteb and news-mtc). Then, we test our probes on both generic (IMDB) and financial (TwitterFinance) sentiment analysis datasets, a shifted domain. The results show that when we use news based hold out datasets, the calibration is more precise | Dataset | Conformal Consistency with alpha=0.8 | Conformal Consistency with alpha=0.9 | |-----------------|--------------------------------------|--------------------------------------| | Twitter-Finance | 77.69 | 81.21 | | IMDB | 83.72 | 86.81 |

b. When in doubt, having more datasets from a particular task is helpful. We see this in the MCQA task, where calibration is consistently better when we add more OOD datasets into the training set (for a single, fixed test dataset). This suggests that include a diverse range of training data helps improve calibration c. Expect larger than specified errors, and if the deployment scenario is safety critical, then use a small hold out test set to validate the conformal probes performance before using it in practice.

We believe that the conformal probes show promising signs of generalizing to unseen distributions, and that a careful selection of training datasets can provide highly performant and well calibrated probes.

On question: scope of detectable behaviors

We are quite interested in better understanding the kinds of behaviors that LM hidden states struggle to predict. One such behavior is whether or not the LM will incorrectly answer an MCQA question. To quantify this, we experimented with 8 MCQ datasets, and trained hidden state probes to predict whether or not the LM will output an incorrect answer to a question. On 5/8 of these datasets, the hidden probes (when used without any conformal calibration) perform worse than, or equivalent to random performance. On the other 3 datasets, the probes perform slightly better than random, however any attempts at calibration show that no reliable threshold can guarantee an error rate of less than 60%, which is quite poor.

More generally, we find that there are two types of tasks that LM hidden states cannot predict well:

1. Lexical quantification tasks: We experimented with tasks such as, predict the exact number of words that will be in the output, predict the number of verbs that will appear in the output etc. While the probes are able to perform reasonably (considerably higher than random) on the binarized version of these tasks (i.e. will the output be long or short, will the output contain a verb or not), it performs poorly on the quantification version. This is consistent with prior works that show LMs struggling with counting based problems[1], and suggests that the problem transfers to hidden state representations.
2. Tasks which require external knowledge: Our hypothesis as to why the MCQ tasks fail, is that unlike the other tasks we explore of format following, jailbreaking, etc, external knowledge is required to verify an answer to the MCQ questions. We can be sure that this knowledge is either not present, or hard to access in the LM, as it incorrectly answers many of these MCQ questions, which suggests that the hidden probes will always struggle to detect when the LM is going to incorrectly answer.

On Weakness: Experiments limited to LLama3

We have made sure to include ablations on other models in the Appendix, showing that the trends discussed in the main text are replicable on the Mistral-7B and DeepSeek-R1-Qwen-14B models. We follow this up with additional experiments below, we train hidden state probes on the following model families: Mistral (Mistral-7B), Qwen3-8B, DeepSeek-R1

The probe results (when used without any conformal calibration) shows that we achieve above random performance with all LMs, showing the central claim of our work: that the hidden states of LMs can predict the future behavior of the output.

We observe that when using conformal alpha = 0.90, the probes are able to maintain high estimation consistency regardless of the underlying LM used, showing that the method is able to properly calibrate probes for a variety of LMs.

Our motivation for focusing on the Llama3 family in the main text is primarily to keep the focus on the methods and tasks, as opposed to the specific language models used. We would posit that we have shown sufficient evidence to believe that none of our results are tied to any particular feature of the Llama3 LMs, and would hence disagree that the paper's use of Llama3 is a significant weakness.

[1]: Fu, Tairan, et al. "Why Do Large Language Models (LLMs) Struggle to Count Letters?." arXiv preprint arXiv:2412.18626 (2024).

On weakness: Paper covers a wide range of topics

We thank you for your feedback and would like to ask you for suggestions on how we may rewrite certain sections to make the paper feel more organised. However, we would expressly disagree that the wide range of topics covered is a weakness. This paper is not attempting to set the state-of-the-art in a particular task, but rather is trying to demonstrate a phenomenon that occurs in LMs across a wide range of inputs. That is, we are trying to firmly establish that across a wide range of inputs and behaviours, the hidden states of LMs contain valuable information on how the model will behave. We additionally hope to show that by leveraging this information, we can craft efficient and precise early warning and exit systems for these behaviours. The fact that the same method of conformal probes is applicable to such a diverse range of topics without any modification is one of its greatest strengths, and shows the generality of our findings. We do not agree that the work should have instead focused on a single set of tasks.

On issue: cost of inference measurement

We thank you for your attention to this detail. We would note our initial motivation for measure % of tokens generated as opposed to wallclock time: we wanted a measure of efficiency that is agnostic to specific hardware implementations, and that captures the reason why the efficiency of the probe-based generation is superior. However, we understand your concerns that the empirical reality may not match up with this more conceptual metric, hence we measure the wallclock time of normal inference v.s. Inference with conformal probes guided early exiting on Llama3-8B using alpha=0.8, and present the results in the table below:

As we can see, the inference cost reductions are significant even when measured by wallclock time. Importantly, the inference time reduction roughly matches the % reduction in tokens generated, which is the original metric we report. This confirms that empirical results match our more hardware-agnostic measures of inference efficiency.