Title: Language Models Can Predict Their Own Behavior
Venue: NeurIPS 2025
Accepted: yes
Keywords: foundation models

Large Language Models

LLMs

Natural Language Processing

NLP
Abstract: The text produced by language models (LMs) can exhibit specific `behaviors,' such as a failure to follow alignment training, that we hope to detect and react to during deployment. Identifying these behaviors can often only be done post facto, i.e., after the entire text of the output has been generated. We provide evidence that there are times when we can predict how an LM will behave early in computation, before even a single token is generated. We show that probes trained on the internal representation of input tokens alone can predict a wide range of eventual behaviors over the entire output sequence. Using methods from conformal prediction, we provide provable bounds on the estimation error of our probes, creating precise early warning systems for these behaviors. The conformal probes can identify instances that will trigger alignment failures (jailbreaking) and instruction-following failures, without requiring a single token to be generated. An early warning system built on the probes reduces jailbreaking by 91%. Our probes also show promise in pre-emptively estimating how confident the model will be in its response, a behavior that cannot be detected using the output text alone. Conformal probes can preemptively estimate the final prediction of an LM that uses Chain-of-Thought (CoT)  prompting, hence accelerating inference. When applied to an LM that uses CoT to perform text classification, the probes drastically reduce inference costs (65% on average across 27  datasets), with negligible accuracy loss. Encouragingly, probes generalize to unseen datasets and perform better on larger models, suggesting applicability to the largest of models in real-world settings.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper shows that large language models (LLMs) often encode signals in their internal states that can predict their future behavior before generating any output. By training simple linear probes on the hidden representations of input tokens and calibrating them with conformal prediction, the authors build early warning systems for issues like instruction-following failures, safety alignment violations (e.g., jailbreaking), and low-confidence responses. These probes can reduce harmful outputs and improve efficiency, such as cutting inference costs by 65% in Chain-of-Thought prompting with minimal accuracy loss. The approach is lightweight, data-efficient, and scales well with larger models. It generalizes to unseen tasks and helps reveal what LLMs implicitly \"know\" about their own future actions.\n\nSTRENGTHS AND WEAKNESSES:\n## Strengths:\n1. The paper is well-written and easy to read.\n2. The research topic is compelling—detecting when an LLM is \"lying\" or misbehaving by analyzing its internal states before token generation.\n3. The paper effectively operationalizes the concept of \"LLM predicting its own behavior\" into practical techniques for improving safety and inference efficiency.\n\n## Weakness:\n1. My main concern is that high probing accuracy doesn't necessarily mean the LLM is predicting its own behavior. It could instead be classifying the input prompts, with the latent representations capturing prompt features rather than future behavior. While the BERT baseline partially addresses this, the performance gap between probing and BERT could also stem from architectural differences or parameter counts. **Although BERT is larger than the probing layer, the authors should take into account the parameter counts of the frozon LLM layers generating the latent representations. 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. I leave more questions about this in the Question section.\n\n2. It would be helpful for the authors to discuss related concurrent work. For example, [1] also uses linear probing to predict LLM response features (e.g., token count, reasoning steps) before generation begins.\n\n[1] Dong, Zhichen, et al. \"Emergent Response Planning in LLM.\" arXiv preprint arXiv:2502.06258 (2025).\n\nQUESTIONS:\n1. In Figure 2, the performance gain of probing over BERT appears marginal. While conformal prediction improves probe accuracy, it’s unclear whether similar techniques could also enhance BERT performance.\n\n2. In Section 4.1, the authors state they filter the training set to only include cases where the LM answers or complies. This risks making responses too similar, effectively turning the task into prompt classification—an issue discussed in the Weakness section. A better setup may involve using similar prompts that lead to different behaviors, training probes to predict those outcomes.\n\n3. In Section 5, it would be useful to report the LLM’s accuracy without CoT. If it's comparable to CoT accuracy, the benefit of using probes to skip CoT may be limited.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 4, "ORIGINALITY": 2, "DATE": "2025-06-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces a novel and practical method for predicting the future behavior of a language model by probing the internal representations of its input tokens, before any text is generated. By integrating simple linear probes with conformal prediction, the authors create a reliable early-warning system with provable error bounds. The method is validated across a wide range of applications, including preemptively detecting safety failures (e.g., jailbreaking) and accelerating Chain-of-Thought inference, demonstrating significant performance gains (e.g., 91% jailbreak reduction, 65% inference cost reduction) with minimal overhead and strong generalization.\n\nSTRENGTHS AND WEAKNESSES:\n### Strenghts\n- The core idea of preemptive behavior prediction is innovative and provides insights into the critical, real-world problems of high inference costs and the need for proactive safety guardrails.\n- The claims are substantiated by extensive experiments across a diverse set of 27 datasets with OOD generalization tests.\n\n### Weaknesses\n- The paper primarily probes the hidden state of the final input token. The motivation and how does the attention mechanism effectively compress the \"intent\" or \"future behavioral plan\" of the entire prompt into this single token's state are unclear.\n- Apart from the out-of-distribution (OOD) generalization tested across different datasets of the same task type, it is necessary to discuss this generalization against more fundamental distribution shifts (e.g., code vs. text classification).\n- The authors chose simple linear probes, which ensures interpretability. However, this might limit their ability to capture more complex, non-linear behavioral patterns. Could a small, non-linear probe (e.g., a shallow MLP) unlock predictions for more nuanced behaviors (e.g., sarcasm, humor) without excessively increasing computational overhead?\n- Lack analysis of the overhead introduced by collecting hidden states for probe training.\n\nQUESTIONS:\n- The finding that \"bigger models lead to better probe performance\" is a fascinating one. Is it because larger models encode their future behavioral plans in a more \"linearly separable\" or \"redundant\" manner? Does this performance improvement trend continue indefinitely, or does it plateau?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-24"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes conformal probes, simple linear classifiers trained on a language model’s input-token hidden states and calibrated with conformal prediction. Once calibrated, the probes fire only when they can guarantee a user-specified error rate, giving an early-warning or early-exit signal before the model generates a single output token. \n- Safety & alignment: Detect prompts that will cause refusal-policy failures, cutting jailbreak success from 30 % to 2.7 % (≈ 91 % relative reduction) on the WildJailbreak set.\n- Format & instruction following: Predict whether an answer will violate output-format constraints with higher precision than a 340 M-parameter fine-tuned BERT baseline while using only ~0.0025 % as many trainable parameters.\n- Efficiency: Interrupt Chain-of-Thought (CoT) generation for text-classification tasks when the probe is confident, yielding a 65 % average reduction in forward passes across 27 datasets with ≤ 1.4 % worst-case accuracy loss and often a small accuracy gain\n\nSTRENGTHS AND WEAKNESSES:\n## Strengths\n- Writing is concise.\n- A lightweight, model-agnostic safety-efficiency layer that can sit atop any inferene stack and composes with speculative decoding; addresses growing cost of long generations.\n\n## Weakness\n- Experiments are somewhat limited to the Llama3 families.\n- No latency measures for probe inference itself.\n- Linear probing are conceptually simple; early-exit ideas exist in speculative decoding (or adative inference) over vision/nlp. So, somewhat incremental novelty may hinge on the conformal guarantee.\n\nQUESTIONS:\n1. Scope of detectable behaviours \n- You report that probes cannot predict outright answer correctness in MCQ tasks. Could you quantify this ceiling (e.g., accuracy vs. random) and elaborate on what differentiates behaviours the probe can vs. cannot capture? Clarifying this boundary would strengthen Section 6.\n\n2. Latency & resource overhead\n- While the method reduces generation compute, each query still requires a forward pass to produce hidden states for the probe. Please provide wall-clock latency and energy measurements relative to vanilla inference to validate net savings.\n\n3. Calibration stability across domains\n- How sensitive are the conformal thresholds to domain shift? The OOD CoT results are encouraging, but would thresholds learned on, say, conversational data hold for code or medical text? A small experiment or guidance on recalibration frequency would help practitioners.\n\n4. What are the practical advantages and potential drawbacks of turning conformal probes into a proprietary, in-house safety and efficiency layer?\n\n5. How would conformal probes perform differently on reasoning‐focused LMs versus instruction‐tuned LMs, and what model-specific considerations should be taken into account?\n\nLIMITATIONS:\nyes.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 2, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper explores the concept of enabling models to predict their own behavior in advance. The proposed method aims to identify potentially problematic behaviors a priori within a given confidence level, thereby facilitating timely intervention. The authors demonstrate the applicability of their approach across multiple tasks, including instruction following, safety alignment, uncertainty estimation, and inference acceleration. While the idea of a priori behavior prediction is promising, the paper would benefit from a stronger focus on a specific domain or core problem. The current broad coverage makes it challenging to evaluate the depth and significance of the technical contribution.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n\n-\tThe idea of a priori behavior prediction has strong practical and application value.\n-\tThe paper is well-written and easy to understand.\n\n**Weaknesses**\n\n-\tThe paper feels disorganized and lacks focus. It covers a wide range of topics—instruction following, safety alignment, uncertainty estimation, and inference acceleration—without clearly defining its main objective. These topics should be explored in separate studies rather than being briefly touched upon in a single paper.\n-\tFor most of the domains discussed, there is little technical distinction between attaching a classifier to a large model and using BERT-based classification; the experimental results also reflect this lack of differentiation.\n-\tThe paper does not compare or discuss its approach against post-hoc decision-making methods, which limits the evaluation of the proposed method's advantage.\n\n**Minor Issues and Suggestions**\n\n-\tFigure 2 may be misleading in its current presentation, even though the coverage percentages are explained in the caption. The term Estimation consistency is also inaccurately used — it should refer specifically to the estimation accuracy among the covered samples , not overall.\n-\tIn the discussion of inference acceleration, the cost of inference is not solely determined by the number of generated tokens; the number of input tokens should also be taken into account. Therefore, the relevant claims or reported metrics should be revised accordingly to reflect a more accurate cost analysis.\n\nQUESTIONS:\nIt is unclear from the paper whether the backbone model is updated during the training of the classifiers. If the backbone is fine-tuned, it might lead to better probing performance, but could also affect the language model’s generation quality. What is the trade-off between probe effectiveness and model behavior stability in such a setting?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-04"}