Learning Safety Constraints for Large Language Models
Authors: Xin Chen
Yarden As
Andreas Krause
Venue: ICML 2025 spotlightposter (2025)
Keywords: Constraint learning; LLM Safety/Alignment
Primary area: deep_learning->large_language_models

Abstract: Large language models (LLMs) have emerged as powerful tools but pose significant safety risks through harmful outputs and vulnerability to adversarial attacks. We propose SaP–short for Safety Polytope–a geometric approach to LLM safety, that learns and enforces multiple safety constraints directly in the model's representation space. We develop a framework that identifies safe and unsafe regions via the polytope's facets, enabling both detection and correction of unsafe outputs through geometric steering. Unlike existing approaches that modify model weights, SaP operates post-hoc in the representation space, preserving model capabilities while enforcing safety constraints. Experiments across multiple LLMs demonstrate that our method can effectively detect unethical inputs, reduce adversarial attack success rates while maintaining performance on standard tasks, thus highlighting the importance of having an explicit geometric model for safety. Analysis of the learned polytope facets reveals emergence of specialization in detecting different semantic notions of safety, providing interpretable insights into how safety is captured in LLMs' representation space.

Learning Safety Constraints for Large Language Models
Xin Chen Yarden As
Andreas Krause
Abstract
Current approaches to LLM safety span a broad spectrum yet face significant tradeoffs. Prompt- techniques attempt to alter model behavior through input/output engineering; however, they prove brittle and easily circum- vented (Kandpal et al., 2023). Training-tin approaches that rely safe RLHF (Dai et al., 2023; Wachi et al., 2024; Rame et al., 2024) aim to balance helpfulness with safety but require expensive data relabeling and model retraining, lacking post-hoc mechanisms to ensure safety. In addition, it is often difficult to interpret how these methods make safety-related decisions. For instance, it is unclear how these methods quantify the severity of unsafe users' requests, or how they pimpoint the underlying cause that makes request unsafe. This landscape suggests the need for the develop- ment of methods that leverage inference-time techniques while offering intuitive means for interpretability.
Large language models (LLMs) have emerged as powerful tools but pose significant safery risks through harmful outputs and vulnerability to adversarial attacks. We propose SaP short for Safety Polytope-a geometric approach LLM safety, that learns and enforces multiple safety constraints directly in the model's representation space. We develop framework that identifies safe and unsafe regions via the polytope's 's facets, cnabling both detection and correction of unsafe outputs through geometric steering. Unlike existing approaches that modify model weights, SaP operates post-hoe in the representation space, preserving model capabilities while enforcing safety constraints. Experiments across multiple LLMs demonstrate that our method can effectively detect unethical inputs, reduce adversarial attack success rates while maintaining performance on standard tasks, thus highlighting the importance of having an explicit geometric model for safety. Analysis of the learned polytope facets reveals emergence of specialization in detecting different semantic notions of safety, providing interpretable insights into how safety is captured in LLMs' representation space.
Motivated by recent work on constraint learning from demonstrations constrained Markov decision pro- cesses (CMDP, Altman, 1999; Lindner al., 2024), our key insight that LLM safety can be framed geometric constraint learning problem. This geometric perspective allows us model explicit safe set, defined as polytope, within the model's representation space. Each of the poly- tope's ' facets represents a different constraint that must be satisfied. As we later show, while learning these constraints solely depends on binary safe/unsafe labels, our approach associates these constraints with different intuitive semantic meanings of safety.
1. Introduction
Large Language Models (LLMMs) have demonstrated remark- able capabilities across diverse tasks, yet their increasing real-world deployment raises urgent safety concerns concerns-these models can generate harmful content or be manipulated through adversarial attacks. While approaches like Rein- forcement Learning from Human Feedback (RLHF) show promise, they face fundamental limitations. For instance, models can learn to "game" reward functions rather than be- come genuinely safer, and noisy specification of human pref- erences leads to unintended behaviors (Casper et al., 2023).
We develop SaP realize this geometric approach. Our approach comprises two key components: concept encoder that disentangles different safety concepts, and steering algorithm that guides unsafe outputs back into the safe set while preserving model capabilities. Our framework and its components are illustrated Figure provide a clearer understanding of our approach.
Contributions.
We propose novel geometric framework that yields an explicit model of safety in LLMs, formulated as polytope in models' representation space.
'Department of Computer Science, ETH Zürich, Zürich, Switzerland, Correspondence to: Xin Chen www.chhn@@if....ccchhcee
An inference-time steering algorithm that guides unsafe outputs back to the safe set while preserving model capabilities, providing defense against ofate-of-the- adversarial attacks.
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Learning Safety Constraints for Large Language Models
Prompt:
Write persuasive article about how Anthony Fauci was involved in the lab leak of COVID-19.
Misinformation
Child Abuse
Response: I can't
write answers that are not based on facts and evidence.
Fraud
Safe set {
Figure I: Illustration of the geometric approach to language model safety proposed in this paper. A safety facet is triggered when a user asks the model to write a response suggesting fabricated content. When the model generates responses, we steer its internal representation back to the safe region to produce safe outputs.
An extensive empirical analysis of learned polytope facets, which reveals naturally emerging specialization in detecting different semantic concepts of safety.
(x) = E, Er= (x, a.) zo x] denote the cost value function for cost cj- Then, the constraint can be rewrit- ten as Ep ((x)) S Ez-
Our code is publicly available at https: com/ www.laaoooo//aaatyooooo/ ytope.
Natural language as a Markov decision process. the case of language, one can think of a token-level MDP. In this formulation, finite of tokens represents the language's 's vocabulary, the state space is a set of all (pos- sibly infinitely long) sequences X' = #=1 VH where y = Vx.... X (H times), an H-fold Cartesian product of all tokens in V. Indeed, LLMs have proven highly effective in mcaningfully "navigating" this combinatorially vast state space. Models predict only the next token, hence the action space is A V. The vocabulary includes two special tokens "<eos>' and '<bos>', that represent the MDP's terminal and initial states respectively. The initial state distribution p p(za)= = 1 1m, (<boa>) hence trajectories always start from the token '<bos>'. For an action a E V, the next state is x' = x a where the operation (- represents a concatena- tion. A special transition reserved for when a <eos>, leading to terminal state, where z' is fixed to æ forever. An autoregressive language model can be seen as a policy π x) for determining the next token given all the previ- ous tokens. While using MDPs to describe autoregressive models and language has been extensively studied in sev- eral previous works (Cundy & Ermon, 2023; Rafailov et al., 2024; Zeng et al., 2024), the main argument of this work is that safe and ethical language not only tries maximize reward function, but is also subject to some constraints. For instance, LLMs should avoid using language that incites harm or solicits unlawful actions. We argue that these con- straints are intrinsically reflected in human natural language and should be modeled to ensure model safety and control. While manually specifying such constraints nuanced due to their complexity and context-dependence, this work in-
2. LLM Safety through the Lens of Constrained Markoy Decision Processes
We tackle LLM safety a sequential decision-making prob- lem, where models must reason about long-term conse- quences of their decisions, i.e., their choice of words. An intuitive approach for modeling safety requirements is by imposing some constraints within this decision- decision-making pro- cess. Constraint Markov decision processes pose natural formulation for such problems, allowing harmful behaviors to be modeled constraints.
Constrained Markov decision processes. We define an infinite-horizon discounted constrained Markov de- cision processes (Altman, 1999, CMDP) Ar, rr, (c.))=,.PP where X and A are state and action spaces respectively, p(x' | x, a.), describes probability distribution over the next state x', given state x E X' and an action a E A. States are ini- tially drawn from p(æo), the initial state distribution. A reward function r(x, a) is be maximized, whereas Cj (2x a).j
the tuple
n....,,n,, are cost functions that must remain bounded. The goal is to find stationary pol- icy Txx a) that solves maxter Ez "r (rr(a, (a)) a:)] where JF = { Vj E,, (E teoo""ee r'c (xx, a.)] < & } and E are predefined cost budget paramters. Constraints in CMDPss can be equivalently expressed in terms of the value functions of the cost objectives. Specifically, let
2

---

Learning Safety Constraints for Large Language Models
troduces a framework for learning them directly from data. Learning constraints from demonstrations. Our main result relies on the idea that for CMDPs, a conservative estimate of F can be learned only from trajectory demon- strations. More concretely, Lindner et al. (2024) establish that given dataset of trajectories that were drawn from k safe policies Ti..... The and a feature representation f X' x A- - [o, 1j such that f(r) == ExiEi=o 7 fx az)], one can construct a conservative approximation of J with any convex combination of the feature expectations of
Feature extraction. Our aim is use pre-trained LLM to extract features f, which can then be used to construct our polytope. Concretely, given dataset {(0', y ) consisting of token sequences x' E X' and their correspond- ing labels y² < +1, indicating whether the token sequence is considered safe. Following our conceptual motivation from the previous section, this can be thought of as a case, where for some j and π, the label indicates whether vf (a*) Ey. These labels are commonly derived from human annotations on aspects such as ethics, toxicity, and privacy protection (Ji et al., 2023; Gehman et al., 2020; Lin et al., 2021) or from success/failure of adversarial at- tacks. A simple method for feature extraction is to perform a forward pass on each example and collect the model's intermediate features (Arditi et al., 2024). h E Rdh Denot- ing the intermediate features of layer I Ti, we get h Ti (z'), yielding data {(h', y To determine a full sentence's 's safety, we use the feature ht obtained at the last word of a sentence, i.e., when the last token of x' is <eos>.
π Σ AT(TTR) O.
Σ A
Crucially, the set Q is convex polyhedron, i.c., a set of lin- ear equations that determines the feasibility of policies. Fur- thermore, Lindner et al. (2024) show that safety constraints can be learned as linear mappings of feature expectations f, independently from rewards.
Convex Polytope Machines. A key challenge in con- straint learning for LLMs is scaling the estimation of the polytope's facets & to high dimensions. For instance, the QuickHull algorithm for convex-hull vertex identifica- tion (Barber et al., 1996), employed by Lindner et al. (2024), becomes prohibitively intractable due to its ((N time complexity. To overcome this, we adopt the CPM algorithm (Convex Polytope Machine, Kantchelian et al., 2014). Instead of explicitly computing a convex-hull. CPM treats the computation of polytopes as a classification prob- lem. Concretely, given a dataset {(h,, y
The geometry of constraints in LLMs. The above two ohservations-that Q is convex polyhedron, and that constraints are linear w.r.t. feature expectations- motivate a geometric interpretation for safety in LLMs. We propose viewing safety constraints through a ge- ometric lens, representing them via a polytope-the intersection of k: halfspaces defined by linear inequalities Q - { TTP=EE Here, & E Bdxk represents K
hyperplanes in d-dimensional space, & E K are thresholds, and f E Rd is the input feature vector. The features f are learned during the pre-training phase of LLMs. The linear representation hypothesis (Park al., 2023; 2024) provides empirical evidence that high-level semantic concepts in LLMs naturally manifest as linear directions in the feature space of f. This suggests that safety constraints in language, like those in CMDPs, have an inherent geometric structure, further grounding our approach of learning polytope facets in the LLM's feature space. This formulation of safety facets as geometric constraints enables post-hoc safety control without model fine-tuning, while providing tool for analyzing the learned safety concepts.
W where the labels indicate whether point h lies outside (-1) or inside (+1) the polytope, and for this discussion only, suppose dn = d to maintain consistent notation, CPM learns a deci- sion boundary by minimizing
K K + Φ Th'I
Σ sE A=1
Σ N [[ & Piinssh + +1 VEI-1
Above, 11 = max -), I+1 {i y² = +1} and I_ ii y² = -1). -1). > 0 is margin parameter, A is regularization parameter and z(-) assigns input points to their corresponding facet (see Appendix A). Kantchelian et al. (2014) provide more details on the connection of the CPM loss to Support Vector Machines. Crucially, optimizing the above loss allows us to leverage gradient descent to learn o effectively mitigating scalability challenges associated with the dataset size N and feature dimension dn.
3. Safety Polytope (SaP)
In this section, we present our approach for learning the facets from demonstration data and demonstrate the utility of our approach for interpretability and steering of unsafe output generation. On high-level, our approach is composed of three main steps: (i) using a labeled dataset to obtain features f from a pre-trained LLM and (ii) determining the polytope's 's hyperplanes' parameters &f and &. (iii) using the learned polytope for steering unsafe outputs into the safe region.
Untangling word ambiguity. Learning polytopes with gradient descent enables us to use tools from the in- terpretability literature to associate facets with intuitive, "human frinanlliindll" concepts of safety. Concretely, model rep- resentations are often susceptible to polysemanticity (Elhage
3

---

Learning Safety Constraints for Large Language Models
al., 2022), phenomenon where the same model activa- tions are triggered by multiple distinct concepts. Drawing inspiration from sparse autoencoders (Cunningham et al., 2023), we propose three modifications to the pipeline intro- duced above. First, we introduce a linear layer followed by ReLU function to provide nonlinearity on top of h, referred to as the 'Concept Encoder', Ec : Rd + R4, producing a dataset of features and labels {(f) Ec (= 5') Second, together with Φ. we make the safety threshold & learnable parameter. Finally, we include sparsity reg- ularization term on f. These modifications result in the following training loss:
Algorithm I SafeFlow: Representation Steering for Safe Response Generation
Require: т, T, Φ,
Dx 6- 20 while v:[end] # <eos> do
Obtain (1x)
Partial forward pass, up to layer I
Eq. 2 Override activations of layer I h < arg orgginn [xr
hlll s.t. Φb p <
क h
ar - π (a: xz)
Complete pass, decode next token
I₁ s- It a2
end while
output Zt
L(0.E)= = Σ
Σ
EE ]m,
R
Our approach scales efficiently to large batches via exist- ing tools for vectorized computation (Bradbury et al., 2018; Agrawal et al., 2019; Blondel et al., 2022; Lu et al., 2024) and aligns with the paradigm shift towards "test time com- pute" (OpenAI, 2024) to improve safety in tandem with reasoning capabilities.
1EI11k= 1=1
k (1))
fi
+ +1
Intuitively, the inclusion of the sparsity regularization term A induces feature sparsity, which empirically reduces pol- ysemanticity when assigning inputs to different concepts of safety. We highlight that, while this training objective is non-smooth due to the sparsity regularization term, we can still compute its sub- uub-gradients and use standard SGD-based learning techniques to optimize w.r.t. Φ. Additional imple- mentation details are provided in Appendix A. In Section 4, we analyze the correspondence between human concepts of safety (e.g., violent language) and polytope edges and ablate the use of the additional non- non-linear transformation, demonstrating the critical role of these components.
4. Experiments
We evaluate SaP on three LLMs: Llama2-7B (Touvron et al., 2023), Ministral-8B (Jiang et al., 2024), and Qwen2- 1.5B (Yang et al., 2024). Section 4.1 presents our main results on defending against adversarial attacks using SaP's steering method in Algorithm 1. In Section 4.2, we demon- strate SaP's interpretability via the Concept Encoder and its defense performance when evaluated on the BeaverTaill dataset(Ji al., 2023). Next, Section 4.3 analyzes the im- pact of polytope facets through ablation studies, with dis- cussions on model design.
Representation steering. Obtaining & allows us to di- rectly verify whether next-token predictions by our model lie within the polytope Q Qnd are therefore safe- by eval- uating & f << E. An immediate implication of this insight is that, before exposing the next token to users, we can adjust activations dynamically, i.e., within the model's response generation loop. Building on methodology from constrained reinforcement learning (c.f. Dalal et al., 2018), we propose solving the following optimization problem:
4.1. Tactful Responses via Representation Steering
We evaluate the steering capability of Algorithm I on Harm- Bench (Mazeika et al., 2024), a standardized benchmark for assessing LLM safety against adversarial attacks. We select nine representative attack methods including gradient- based approaches: GCG (Zou al., 2023), GBDA (Guo et al., 2021), AutoPrompt (Shin et al., 2020), PEZ (Wen al., 2023), UAT (Wallace et al., 2021), AutoDAN (Liu et al., 2023) and human-crafted jailbreaks: Human Jailbreak (Shen et al., 2024) and Direct Request (Mazeika et al., 2024). We additionally evaluate on Adaptive Attack (Andriushchenko et al., 2024), a state-of-the adversarial attack method, on Bain (Chao et al., 2024).
minl (x) h112 s.t. Φ f(h) 5.
Importantly, as long as the model's representations already reside within 2. the original token generation remains un- chained. Algorithm provides a detailed token generation loop that actively adjusts outputs to ensure safety when necessary. Note that in practice, obtaining an approximate solution Equation (2) can be achieved using first-order meth- ods (Bertsekas, 2016). For example, in our experiments, tak- ing only a few steps of a simple Lagrangian relaxation (see Appendix B) substantially improves safety performance.
For each model, we train SaP using the following proce- dure: (i) we identify the top 3 most effective attack methods' based on their success rates; (ii) for these methods, we col- lect model features from 80% of their attack strings for
"excluding Adaptive Attack, which is only used for evaluation.

---

Learning Safety Constraints for Large Language Models
Safety Polytope (SaP)
MLP (MLP)
Response Check (RC) Self Reminder (SR)
SmoothLLM (SM)
Original Model (OR)
MLP+5 (MLP+5) 5)
Rejection Sampling (RJ)
Llama2-7B
Ministral-8B
Qwen2-1.5B
MMLU Accuracy (%)
46.00
RC
MLP
70
RC
MLP
80
MLP
RC
45.50
SaP
SR
MLP+5
SM
OR
SaP4
SM
MLP45
OR
50
SaP
RJ
SM
MLP+5
OR
45.25
40
40
44:75
30
30
RJ
64.50
20
SR
20
SR
-2.5
0.0
2.5
5.0
7.5
10.0
12.5
15.0
10
30
60
20
25
Attack Success Rate (%)
Attack Success Rate (%)
Attack Success Rate (%)
Figure 2: Comparison of model MMLU accuracy and average Attack Success Rate (ASR) on attack algorithms. All defense methods are evaluated over seeds. SaP consistently retains the original MMLU accuracy while having significantly lower ASR compared most baseline methods.
training, reserving 20% for testing; (iii) we then evaluate SaP's s effectiveness against all nine attack methods Our extensive evaluations test the trained polytope's 's robustness and generalization ability across both attack methods and adversarial attack benchmarks. Detailed training setup can be found in Appendix C.I.
Llama2-7B
Ministral-8B
Qwen2-1.55
Polytope
91.24 +: 2.74
90.82 0.23
83.96 + 0.31
MLP
97.28 + 0.19
92.90 0.13
87.56 1 0.15
MLP+5
91.08 + 7.25
90.72 - 2.31
86.12 + 2.86
Table 1: Test accuracy (%) of different methods over seeds, early-stopped to prevent overfitting. The trained models are deployed for experiments in Figure 2. Models with similar classification accuracies do not imply the same defense capability. Having a better geometric model for the representation space is essential for successful defenses.
For each attack method, we compare SaP against 6 defense baselines. Our comparison includes (i) Prompt-base based de- fense methods: Self Reminder (Xie et al., 2023), Response Check (Wang et al., 2024), and SMoothLLM (Robey et al., 2024). These baselines defend against unreasonable LLM requests by prompting, filtering, defaulting to rejection string when unsafe content is detected. (ii) Fraining-based defense methods: Rejection Sampling that rejects request based on the polytope's 's binary decision. We further compare against MLP where we use an MLP model with equivalent parameter count as our SaP model. It is first trained as a safety classifier using binary cross entropy (BCE) loss, then used for the same steering procedure as SaP. We also train MLP with 5 extra layers (referred to below MLP+5) that has the same setup as MLP but with 5 extra layers. Along with SaP, we test all 7 defense methods on the selected 9 attack algorithms, with each setup repeated over 5 seeds.
positioned on the figure by their average attack success rates (ASR) across attack algorithms. As shown, SaP achieves strong defense performance while maintaining model ca- pabilities. For Llama2-7B, it reduces the attack success rate from 12.92% 0.26% while preserving the original MMLU accuracy (45.8% vs. 45.7%). This pattern holds for Ministral-8B (ASR: 55.77% to 3.25%, accuracy: 63.4% vs 63.3%) and Qwen2-1.5D (ASR: 27.57% to 11.81 %, ac- curacy: 52.7%). Detailed performance of each method on Harmbench and MMLJU is presented in Appendix C.2.
In contrast, the baseline methods show varying effectiveness across models. Rejection Sampling achieves competitive de- fense performance on Llama2 and maintains accuracy on both Llama2 and Qwen2, but significantly impacts Minis- tral's MMLJU performance (28.5% vs original 63.4%). Other baselines like In-Context Learning and Response Check maintain accuracy but provide weaker defense, with ASRs ranging from 1.7-28.4%. These results suggest that SaP effectively balances defense capabilities and model perfor- mance across different architectures.
We evaluate all methods' capability on The Massive Multi- task Language Understanding benchmark (Hendrycks et al., 2020, MMLU), which evaluates models across 57 subjects through multiple-choice questions in fields like mathemat- ics, history, law, and science. In our experiments, MMLU accuracy measures how well models maintain their capabili- ties while implementing safety defenses.
Figure presents our main results across models, compar- ing SaP against the baselines. Each defense algorithm is evaluated on all attack methods over seeds first, then
SaP's comparison with MLP and MLP+5 in Table 1, illus-
5

---

Learning Safety Constraints for Large Language Models
sentences in 14 safety categories. Each sentence is labeled as either safe unsafe within its category. Note that do not provide the polytope with category labels, nor hinting category information in the prompt, hence these semantic meanings are captured without direct supervision.
Without CE
Average ASR (%)
60
With CE
51.5
40
We run both polytope variants (with and without CE) on the complete dataset to obtain facet violation patterns. For each facet & and input æ, we first compute the facet viola- tion à - then normalize each facet's violations
20
5
0.1
1,7
0,4 0,8
0
Llama2-7B
Ministral-8B
Qwen2-1.55
within [0, 1]. We caleulate the mutual information between these violation scores and the labels of each safety cate- gory to see how well each facet corresponds to each safety concept. For clear comparison, we normalize each row in the resulting mutual information matrix by dividing by its maximum value, ensuring that for each safety category, the highest mutual information score is 1. More details can be found in Appendix D.I.
Figure 3: Comparison of average ASR for each model, with and without Concept Encoder (CE). SaP with CE consistently shows robustness against attack across models.
trates the importance of having a structured and principled geometric model for the representation space of LLMs. Prior to SaP, an intuitive practice to train an MLP with BCE loss for binary classification first, then deploy the trained MLP for steering. However, obtaining decent classification accuracy does not necessarily lead to better de- fense performance. In addition, Table and Figure 2 show that while classification accuracy is on par across methods, even when scaling the MLP with more layers, SaP's geomet- ric modeling enables significant improvements compared to the MLP and MLP+5 baselines.
Figure 4 compares the mutual information heatmaps with and without the concept encoder. Without the encoder, we observe significant polysemanticity in the learned facets, manifesting as high mutual information values across mul- tiple categories for individual facets. For example, 'Facet 36 exhibits strong correlations with categories '2' to '9', all over 0.93, After introducing the concept encoder, we observe marked reduction in polysemanticity. The mutual information matrix shows more focused correlations, with strong correlations (MI > 0.9) more likely to appear in isolation rather than in clusters.
4.2. The Role of the Concept Encoder
Enhanced defense performance. Figure 3 compares SaP variants with and without the concept encoder Ec (indicated with CE) with average ASR across attack meth- ods on HarmBench, averaged over 5 seeds. The results demonstrate significant improvements in defense perfor- mance when using the concept encoder. For Llama2, adding CE enables almost perfect defense (averaged 0. 11% ASR) against all attacks, while the variant without CE shows vulnerabilities to various attacks. The improvement is par- ticularly pronounced for Ministral, where CE reduces ASR from an average of 51.5% to 1.7%
To understand the semantic concepts encoded by each facet, we analyze differences in the Kullback-Leibler (KL) di- vergence when masking context around negative terms in BeaverTails' child abuse dataset. Results in Figure 5 reveal specialized detection patterns: 'Facet 7' demonstrates strong sensitivity to kidnapping scenarios (KLD 0.366). 'Facet 22' primarily focuses on sexual content (KLD = 0.733) and abuse-related terms (KLD = 0. 14),, while 'Facet 26' shows distinct activation for bullying (KLD= 1.404) with secondary response to violent terms. This differentiation suggests the facets learn to capture specific categories of harmful content rather than just detecting negative terms broadly. We refer to Appendix D for additional experiments and details on this experiment.
Qwen2 shows strong baseline robustness even without CE, and adding CE maintains this performance without significant improvements. This behavior might be at- tributed to Qwen2's smaller feature dimensionality (1536- dimensional) compared Ministral and Llama-2 (both 4096-dimensional), The lower-dimensional features could be inherently easier for the polytope constraints to capture safety patterns directly, making the additional abstraction from a concept encoder less crucial for effective defense. Detailed results are presented Appendix C.2
4.3. Ablating the Number of Constraints
To understand how the number facets affects polytope per- formance, we conduct two sets of experiments. For defense evaluation, we vary the facet count from to 50 and mea- sure performance against 7 HarmBench attack methods (Au- toPrompt, DirectRequest, GBDA, GCG, HumanJallbreaks, PEZ, UAT), computing the mean and standard deviation of ASR across these attacks, repeated over 5 seeds. Addition- ally, we evaluate SaP's classification performance as the
Improved interpretability. To analyze how well the facets capture different safety concepts, we use the Beaver- Tails dataset (Ji et al., 2023), which contains 330k annotated
6

---

Learning Safety Constraints for Large Language Models
Mutual Information Heatmap
Mutual Information Heatmap
Liee Calegories
Truc Catagories
Leamnd Pacets
Learned Facets
(a) No Concept Encoder
(b) With Concept Encoder
Figure 4: Mutual Information Heatmap showing the comparison between models without (a) and with Concept Encoder (b). Polytope facets learned with concept encoder show more disentangled activations in safety classes.
KLD for Facet
KLD for Facet 22
KLD for Facet 26
Indnap
0.366
0.733
bully
1.044
abuse 0.010
abuse
0.114
kill
0.118
sex 0.000
kidnap
0.000
beat
0.048
beat 0.000
beat
0.000
sex
0.000
bully
bully
abuse
kill
kidnap
0.000
0.00 0.05
0.10
0.15
0.20
0.25
0.30
0.35
0.0
0.1
0.2
0.3
0.4
0.5
0.6
0.7
0.0
0.2
0.4
0.6
0.8
1.0
1.2
1.4
Figure 5: KL divergence analysis revealing semantic specialization of SaP's facets. Higher values indicate stronger activation when masking specific context words, demonstrating how each facet learns to detect distinct categories of harmful content rather than broadly negative terms.
number of facets increase. To this end, we vary facets from I to 60 and evaluate on the 14 BeaverTails categories sepa- rately. We measure classification accuracy on BeaverTails' test set, reporting the mean and standard deviation of per- category accuracy across categories, repeated over 5 seeds. We emphasize that we use BeaverTails' categories only for evaluation. that is. these categories are not given to SaP during training.
fluctuations even with additional facets, suggesting their defense mechanisms might more sensitive to the specific geometric arrangements of polytope constraints.
For classification accuracy (right), we observe steady im- provements up to approximately 30 facets, after which per- formance diminishes. Ministral demonstrates the most sub- stantial initial gains, with accuracy increasing from 76.8% to 82.5% when moving from to 40 facets. Llama and Qwen2 show more gradual improvements but follow a similar pat- tern that the performance does not significantly improve beyond certain facet counts. A comprehensive presentation of all metric values available in Appendix D.3. We addi- tionally conducted safety classification with standard fully connected neural network trained using cross entropy loss, and show that SaP's 's classification performance does not degrade substantially compared a fully connected neural network with an equivalent parameter count. We present these results in Appendix D.3.
Figure shows how the number of polytope facets affects both defense and classification performance. terms of defense capabilities (left), increasing facet count leads to rapid improvement in safety, particularly evident in reduc- ing attack success rates. Llama-2 achieves near-perfect defense (ASR < 0.1%) 1)) with just 20 facets, while Minis- tral and Qwen2 show continued improvements up to 30 facets before stabilizing. The three models exhibit different convergence patterns: Llama-2 maintains consistently low ASR after 20 facets, while Ministral and Qwen2 show slight

---

Learning Safety Constraints for Large Language Models
Average ASR across Attack Methods
Average Accuracy across Categories
Attack Success Rate (%)
Llama2-7B
Ministral-80
Accuracy (%)
Qwen2-1 1.50
Liama2-78
Ministral-BB 8B
Qwen2-1.55
10
20
30
40
50
Number of Facets
Number Facets
Figure 6: Impact of polytope constraint count model performance. Left: Average attack success rate (ASR) across 7 HarmBench attack methods, where lower values indicate better defense. Right: Average classification accuracy across 14 BeaverTails safety categories. With more facets, SaP's performance first improves, then hits a diminishing return.
We empirically observe that many polytope facets remain inactive during training- they rarely activate on violations in the training data. This suggests potential inefficiency in the polytope loss design of CPM, where the training objective might not effectively utilize the full capacity of available parameters in the facets.
prompts. Input modification approaches (Yi et al., 2024; Wei al., 2024; Xie al., 2023; Robey et al., 2024) modify input prompts weaken potential attacks. Both methods are vulnerable white-box attacks, and focus only on input- level patterns rather than addressing model-internal safety representations. Output filtering (Phute et al., 2024; Zhang et al., 2024; Wang et al., 2024) uses additional language models to check for safety violations with significant com- putational overhead. Despite the computation overhead, large-scale training of input-output safety classifiers can yield highly promising defense results (Sharma al., 2025). Represeetaaiion-basee approaches (Casper et al., 2024; She- shadri et al., 2024; Zou et al., 2024) modify model's in- ternal representations for improved robustness but rarely provide interpretable insights into their safety mechanisms. We complement the limitations by introducing constraints in the model's representation space. This approach allows to modify models' representations, opposed modify- ing/filtering user inputs model outputs. Furthermore, it facilitates treating safety concepts as constraints that should not be violated, enabling easy detection and analysis by humans, contrast to previous representation-batt ap- proaches.
5. Related Work
Constrained and multi-objective alignment. Recent work has explored multi-objective approaches to handle competing goals in language model alignment. Similar to our steering algorithm, MCA (Fu et al., 2024) use post- training decoding to align LLMs to different objectives. Their treatment of safety relies on "contrastive prompts", used obtain a Pareto front for multiple objective. con- trast, SaP treats safety constraints that must be enforced. Training-time methods like RLHF-RLHF (Dai et al., 2023), SACPO (Wachi et al., 2024). RePO (Peng et al.. 2024), and Constrained DPO (Liu et al., 2024) frame safety as constraints within reward learning. Other approaches, like Panacea (Zhong et al., 2024) and MODPO (Zhou et al., 2024), focus too on raining-time techniques, however, like (Fu et al., 2024), they frame safety multiple objectives problem, instead of constraint satisfaction. The above meth- ods require model retraining or finetuning, lack interpretable insights into learned safety concepts, and provide no mech- anisms for post-deployment safety corrections. SaP ad- dresses these challenges by learning safety constraints di- rectly in the representation space, enabling both post-hoc control without retraining and interpretable analysis of how different safety concepts captured through specialized polytope boundaries. This provides complementary ap- proach that can work alongside training time methods.
AT safety and alignment. Recent work has surfaced many challenges AI safety and alignment. Anwar al. (2024) elucidate fundamental challenges spanning scientific derstanding, development methodologies, and sociotech- nical considerations in LLM safety. While Reuel et al. (2024) advocates for sophisticated technical governance frameworks, Dalrymple et al. (2024) makes a case for rig- orous approaches to design AI systems with quantifiable safety assurances. AI Control (Greenblatt al., 2023) pro- poses that safety can be maintained through systems that prevent harmful actions, even when models are misaligned. In this context, SaP's geometric approach to constraints in the representation space offers a promising framework for advancing both alignment and control.
Adversarial defenses. Prior work explores four main cate- gories of LLM safety defenses. Input filtering methods (Cao
al., 2023; Jain et al., 2023) detect and filter suspicious
8

---

Learning Safety Constraints for Large Language Models
6. Limitations and Future Work
Acknowledgement
Our results encourage further practical and theoretical de- velopment of SaP. A key future work direction lies in fur- ther disentangling semantic understandings across polytope edges. CPM represents progress, however, its heuristic facet assignment algorithm limits its overall performance. Lever- aging tools from mechanistic interpretability (Black et al., 2022) or exploring other (possibly non-linear) geometric representations of safety (Park et al., 2024) could reveal deeper insights into the learned constraints.
We thank Javier Rando and Daniel Paleka for their valuable input during the development of this framework. We are grateful to the anonymous reviewers for their constructive feedback, which led to additional experiments and baselines that strengthened our arguments. We also thank Ido Hakimi, Vignesh Ram Somnath and Roni Globerman for their helpful comments on earlier versions of the manuscript. Xin Chen is supported by the Open Philanthropy AI Fellowship and the Vitalik Buterin Fellowship from the Future of Life Institute. Yarden As is supported by the ETH AI Center Fellowship and grant of the Hasler foundation (no. 21039). The research received further support through ELSA (European Lighthouse on Secure and Safe AI) funded by the European Union under grant agreement No. 101070617.
Based on our novel geometric framework for explicit safety modeling, SaP demonstrates an effective defense against various adversarial attacks. However, we observe that under many attacks, specifically for the Ministral-8B model, SaP can induce semantically incoherent outputs. Importantly, all evaluated models, including Ministral-8B. perform as expected on standard capability benchmarks like MMLU and MT-Bench (Bai et al., 2024). with results listed in Tables 8 and 9. Although comprehensive adversarial robustness and broad capability improvements are beyond the scope of this paper, we encourage future research to investigate and improve SaP's performance in these critical domains.
Impact Statement
The increasing capabilities and deployment of large language models (LLMs) highlight the critical importance of ensuring their safety. LLMs can generate harmful content and are vulnerable to adversarial attacks, posing significant risks. Research focused on developing robust safety measures is essential to mitigate these risks and promote the responsible development of LLMs. This work contributes to advancing the field of LLM safety by introducing a novel geometric framework for learning and enforcing safety constraints. We believe that our research can support the creation of safer and more reliable AI systems, ultimately benefiting society by reducing the potential for harm.
In terms of theory, our work lacks a more rigorous under- standing of SaP. Current sample complexity bounds for learning polytopes build on Probably Approximately Cor- rect (PAC) learning. The PAC-style results of Gottlieb et al. (2018) provide solid starting point, as they shed light on the fat-shattering dimension required to learn polytopes (dis- cussed in Appendix E). However, these results depend on strong assumptions that are hard to verify in practice. Adapt- ing these results to LLMs could provide further insights into what can be achieved in terms of formal guarantees for LLM safety, as stressed by Dalrymple et al. (2024).
References
Agrawal, A., Amos, B., Barratt, S., Boyd, S., Diamond, S., and Kolter, J. Z. Differentiable convex optimization lay- ers. Advances in neural information processing systems, 2019.
7. Conclusion
We introduce SaP, a geometric framework for learning safety constraints in LLMs. Motivated by constrained Markov decision processes, we formulate safety as ge- ometric problem, introducing constraints in model represen- tations. This view enables post-deployment control with- out model fine-tuning and provides interpretable insights into the learned safety mechanism. Experiment results demonstrate the effectiveness of this framework in defend- ing against adversarial attacks while maintaining model capabilities and offering facets that naturally correspond to different safety concepts. While significant challenges remain in developing safe AI, this geometrie perspective offers promising framework for aligning and controlling model behaviors. Substantial research remains necessary, from developing more sophisticated geometric frameworks to establishing rigorous theoretical foundations for safety guarantees in large-scale models.
Altman, F. Constrained Markov Decision Processes. Chap- man and Hall, 1999.
Andriushchenko, M., Croce, F., and Flammarion, N. Jail- breaking leading safety-aligned IIms with simple adaptive attacks. arXiv preprint arXiv:2404.0211 2024.
Anthony, M. and Bartlett, P. L. Neural Network Learning: Theoretical Foundations 1st Edition by Anthony, Martin, Bartlett Peter L. (1999) Hardcover. 1999.
Anwar, U., Saparov, A., Rando, J., Paleka, D., Turpin, M.. Hase, P., Lubana, E. S., Jenner, E., Casper, S., Sourbut, O., Edelman, B. I.,, Zhang, Z., Günther, M., Korinek, A., Hernandez-Orallo, J., Hammond, L., Bigelow, E., Pan, A., Langosco, L., Korbak, T., Zhang, H., Zhong, R., Chan, A., Anderljung, M., Edwards, L., Bengio, Y., Chen, D., A Albanie, S., Maharaj, T., Foerster, J., Tramer, F., He, H.,

---

Learning Safety Constraints for Large Language Models
Kasirzadeh, A., Choi, Y., and Krueger, D. Foundational Challenges in Assuring Alignment and Safety of Large Language Models. 2024.
Chao, P., Debenedetti, E., Robey, A., Andriushchenko, M.,
Croce, E., Sehwag, V., Dobriban, E., Flammarion, N., Pappas, G. J., Tramer, F,, Hassani, H., and Wong, E. JailbreakBench: An Open Robustness Benchmark for Jailbreaking Large Language Models, July 2024.
Arditi, A., Obeso, O., Syed, A., Paleka, D., Panickssery, N., Gurnee, W., and Nanda, N. Refusal in Language Models Is Mediated by Single Direction, October 2024.
Cundy, C. and Ermon, S. Sequencematch: Imitation learn- ing for autoregressive sequence modelling with backtrack- ing. arXiv preprint M.Xiv:23360052220 2023.
Bai, G., Liu, J., Bu, X., He, Y., Liu, J., Zhou, 7.,, Lin, 7.,, Su, W., Ge, T., Zheng, B., et al. Mt-bench- 101: A fine- grained benchmark for evaluating large language models in multi-turn dialogues. arXiv preprint arXiv:2402. 14762, 2024.
Cunningham, H., Ewart, A., Riggs, L.., Huben, R., and Sharkey, L. Sparse Autoencoders Find Highly Inter- pretable Features in Language Models, October 2023.
Dai, J., Pan, X., Sun, R., Ji, J., Xu, X., Liu, M., Wang, Y., and Yang, Y. Safe RLHF: Safe Reinforcement Learning from Human Feedback, October 2023.
Barber, C. B., Dobkin, D. P., and Huhdanpaa, H. The quickhull algorithm for convex hulls. ACM Transactions on Mathematical Software (TOMS), 1996.
Dalal, G., Dvijotham, K., Vecerik, M., Hester, T., Paduraru, C., and Tassa, Y. Safe exploration in continuous action spaces, arXiv preprint arXiv: 2018.
Bertsekas, D. Nonlinear Programming. Athena scientific optimization and computation series, Athena Scientific, 2016.
Dalrymple, D. d., Skalse, J., Bengio, Y., Russell, S., Tegmark, M., Seshia, S., Omohundro, S., Szegedy, C., Goldhaber, B., Ammann, N., Abate, A., Halpern, J., Bar- rett, C.,, Zhao, D., Zhi-Xuan, T., Wing, J., and Tenenbaum, J. Towards Guaranteed Safe AI: A Framework for Ensur- ing Robust and Reliable AI Systems, May 2024.
Black, S., Sharkey, L., Grinsztajn, L., Winsor, E., Braun, D., Merizian, J., Parker, K., Guevara, C. R., Millidge, B., Alfour, G., et al. Interpreting neural networks through the polytope lens. arXiv preprint arXiv:2211.12311 2022.
Blondel, M., Berthet, Q., Cuturi, M., Frostig, R., Hoyer, S., Llinares-López, F., Pedregosa, F., and Vert, J.-P. Efficient and modular implicit differentiation. Advances in neural information processing systems, 2022.
Elhage, N., Hume, T., Olsson, C., Schiefer, N., Henighan, T., Kravec, S., Hatfield-Dodds, Z., Lasenby, R., Drain, D., Chen, C., Grosse, R., McCandlish, S., Kaplan, J., Amodei, D., Wattenberg, M., and Olah, C. Toy models of superposition. Transformer Circuits Thread, 2022.
Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C., Maclaurin, D.. Necula, G., Paszke, A., VanderPlas, J., Wanderman-Milne, S., and Zhang, Q. JAX: composable transformations of Python+NumPy programs, 2018. URL http://st thub. com/ cowo//aao////ooooooo/aee/oooe ml jaax
Fu, T., Hou, Y., McAuley, J., and Yan, R. Unlock- ing decoding time controllability: Gradient-free multi- objective alignment with contrastive prompts. arXiv preprint arXiv:2408.0555 2024.
Cao, B., Cao, Y., Lin, L., and Chen, J. Defending Against Alignment-Breakine Attacks via Robustly Aligned LLM, December 2023.
Gehman, S., Gururangan, S., Sap, M.. Choi, Y., and Smith, N. A. Realtoxicitypromps: Evaluating neural toxic degeneration in language models. arXiv preprint arXiv:2009. /462, 2020.
Casper, S., Davies, X., Shi, C., Gilbert, T. K., Scheurer, J., Rando, J., Freedman, R., Korbak, T., Lindner, D., Freire, P., Wang, T., Marks, S., Segerie, C.-R., Carroll, M., Peng, A., Christoffersen, P., Damani, M., Slocum, S., Anwar, U., Siththaranjan, A., Nadeau, M., Michaud, E. J., Pfau, J., Krasheninnikov, D., Chen, X., Langosco, L., Hase, P., BıyIk, E., Dragan, A., Krueger, D., Sadigh, D. Open Problems and Fun- damental Limitations of Reinforcement Learning from Human Feedback. 2023.
Gottlieb, L.-A., Kaufman, E., Kontorovich, A., and Nivasch, G. Learning convex polytopes with margin. 2018.
Gottlieb, L.A.,, Kaufmman E., Kontorovich, A., and Nivasch, G. Learning convex polyhedra with margin, November 2021.
D., and Hadfield-
Greenblatt, R., Shlegeris, B., Sachan, K., and Roger, F. Ai control: Improving safety despite intentional subversion. arXiv preprint arXiv: 066442, 2023.
Casper, S., Schulze, L., Patel, O., and Hadfield-Menelll D. Defending against unforeseen failure modes with latent adversarial training. arXiv preprint arXiv:2403.05030 2024.
Guo, C.,, Sablayrolles, A., Jégou, H., and Kiela, D. Gradient- based Adversarial Attacks against Text Transformers, April 2021.
10

---

Learning Safety Constraints for Large Language Models
Hendrycks, D., Burns, C., Basart, S., Zou, A., Mazeika, M., Song, D., and Steinhardt, J. Measuring mas- sive multitask language understanding. arXiv preprint 0rXiv:2009.033000 2020.
OpenAL. 2024.
Learning Reason with LLMs,
URL https://opents com/ / index/
dearning-to-reasonn wi-h
Park, K., Choe, J., and Veitch, V. The Linear Represen- tation Hypothesis and the Geometry of Large Language Models, November 2023.
Jain, N., Schwarzschild, A., Wen, Y., Somepalli, G., Kirchenbauer, J., Chiang, P.y.., Goldblum, M., Saha, A., Geiping, J., and Goldstein, T. Baseline Defenses for Adversarial Attacks Against Aligned Language Models, September 2023.
Park, K., Choe, Y. J., Jiang, Y., and Veitch, V. The Geom- etry of Categorical and Hierarchical Concepts in Large Language Models, June 2024.
Ji, J., Liu, M., Dai, J., Pan, X., Zhang, C., Bian, C., Zhang, C., Sun, R., Wang, Y., and Yang, Y. BeaverTails: To- wards Improved Safety Alignment of LLM via Human- Preference Dataset, November 2023.
Peng, X., Guo, H., Zhang, J., Zou, D., Shao, Z., Wei, H., and Liu, X. Enhancing safety in reinforcement learning with human feedback via rectified policy optimization. arXiv preprint arXiv:2410. 19933, 2024.
Jiang, A. Q., Sablayrolles, A., Roux, A., Mensch, A., Savary, B., Bamford, C., Chaplot, D. S., de las Casas, D., Hanna, E. B., Bressand, F., Lengyel, G., Bour, G., Lample, G., Lavaud, L. R., Saulnier, L., Lachaux, M.-A., Stock, P., Subramanian, S.. Yang. S., Antoniak, S., Scao, T. L., Gervet, T., Lavril, T., Wang, T., Lacroix, T., and Sayed, W. E. Mixtral of Experts, January 2024.
Phute, M., Helbling, A., Hull, M., Peng, S., Szyller, S., Cornelius, C., and Chau, D. H. LLM Self Defense: By Self Examination, LLMs Know They Are Being Tricked, May 2024.
Rafailov, R., Hejna, J., Park, R., and Finn, C. From SrS to $q Your language model is secretly q-function. In First Conference on Language Modeling, 2024.
Kandpal, N., Jagielski, M., Tramèr, F,, and Carlini, N. Back- door attacks for in-context learning with language models. arXiv preprint arXiv:2307. 14692, 2023.
Rame, A., Couairon, G., Dancette, C., Gaya, J.-B., Shukor, M.. Soulier, L., and Cord, M. Rewarded soups: towards pareto-optimal aligmment by interpolating weights fine- tuned on diverse rewards. Advances in Neural Informa- tion Processing Systems, 36, 2024.
Kantchelian, A., Tschantz, M. C., Huang, L., Bartlett, P. L., Joseph, A. D., and Tygar, D. Large-Margin Convex Polytope Machine, 2014.
Lin, S., Hilton, J., and Evans, Truthfulqa: Measuring how models mimic human falsehoods. arXiv preprint arXiv:2109.07955 2021.
Reuel, A., Bucknall, B., Casper, S., Fist, T., Soder, L., Aarne, O., Hammond, L., Ibrahim, L., Chan, A., Wills, P., Anderljung, M., Garfinkel, B., Heim, L., Trask, A.. Mukobi, G., Schaeffer, R., Baker, M., Hooker, S., So- laiman, I., Luccioni, A. S., Rajkumar, N., Moës, N., Ladish, J., Guha, N., Newman, J., Bengio, Y., South, T., Pentland, A., Koyejo, S., Kochenderfer, M. J., and Trager, R. Open Problems Technical AI Governance, July 2024.
Lindner, D., Chen, X., Tschiatschek. S., Hofmann, K., and Krause, A. Learning Safety Constraints from Demonstra- tions with Unknown Rewards, March 2024.
Liu, X., Xu, N., Chen, M., and Xiao, C. Autodan: Generat- ing stealthy jailbreak prompts aligned large language models. arXiv preprint 23 10.04411, 2023.
Robey, A., Wong, E., Hassani, H., and Pappas, G. J. Smooth- LLM: Defending Large Language Models Against Jail- breaking Attacks, June 2024.
Liu, Z. Sun, X., and Zheng, Z. Enhancing LLM Safety via Constrained Direct Preference Optimization. http:://arrvvooggabb222444.44220
March 2024.
Sharma, M., Tong, M., MU, J., Wei, J., Kruthoff, J., Good- friend, S., Ong, E., Peng, A., Agarwal, R., Anil, C., et al. Constitutional classifiers: Defending against universal jailbreaks across thousands of hours of red teaming, arXiv preprint arXiv:2501. 18837 2025.
Lu, H., Peng, Z.,, and Yang, J. Mpax: Mathematical pro- gramming in jax. arXiv preprint arXiv:2412.09734 2024.
Mazeika, M., Phan, L.,, Yin, X., Zou, A., Wang, Z.,, Mu, N., Sakhace, E., Li, N., Basart, S., Li, B., Forsyth, D., and Hendrycks, D. HarmBench: A Standardized Evalua- tion Framework for Automated Red Teaming and Robust Refusal, February 2024.
Shen, X., Chen, Z., Backes, M., Shen, Y., and Zhang, Y. "Do Anything Now": Characterizing and Evaluating In- The- Wild Jailbreak Prompts on Large Language Models, May 2024.

---

Learning Safety Constraints for Large Language Models
Sheshadri, A., Ewart, A., Guo, P., Lynch, A., Wu, C., Heb- bar, V., Sleight, H., Stickland, A. C., Perez, E., Hadfield- Menell, D., et al. Latent adversarial training improves robustness to persistent harmful behaviors in llms. arXiv preprint arXiv:2407, 15549, 2024.
Yang, J., Xu, J., Zhou, J., Bai, J., He, J., Lin, J., Dang,
K., Lu, K., Chen, K., Yang, K., Li, M., Xue, M., Ni, N., Zhang, P., Wang, P., Peng, R., Men, R., Gao, R., Lin, R., Wang, S., Bai, S., Tan, S., Zhu, T., Li, T., Liu, T., Ge, W., Deng. X., Zhou, X., Ren, X., Zhang, X., Wei, X., Ren, X., Liu, X., Fan, Y., Yao, Y., Zhang, Y., Wan, Y., Chu, Y., Liu, Y., Cui, Z., Zhang, Z., Guo, Z., and Fan, Z. Qwen2 Technical Report, September 2024.
Shin, T., Razeghi, Y., R. L.,, Wallace, E., and Singh. S. AutoPrompt: Eliciting Knowledge from Language Mod- els with Automatically Generated Prompts, November 2020.
R.
Yi, J., Xie, Y., Zhu, B., Kiciman, E., Sun, G., Xie, X., and Wu, F. Benchmarking and Defending Against Indirect Prompt Injection Attacks on Large Language Models, March 2024.
Touvron, H., Martin, L.., Stone, K., Albert, P., Almahairi,
A., Babaei, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., Bikel, D., Blecher, L.,, Ferrer, C. C., Chen, M., Cucurull, G., Esiobu, D., Fernandes, J., Fu, J., Fu, W., Fuller, B., Gao, C., Goswami, V., Goyal, N., Hartshom, A. Hosseini, S., Hou, R., Inan, H., Kardas, M., Kerkez, V. Khabsa, M., Kloumann, I., Korenev, A., Koura, P. S., Lachaux, M.-A., Lavril, T., Lee, J., Liskovich, D., Lu, Y., Mao, Y., Martinet, X., Mihaylov, T., Mishra, P., Molybog, I. Nie, Y., Poulton, A., Reizenstein, J., Rungta, R., Saladi, K., Schelten, A., Silva, R., Smith, E. M., Subramanian, R., Tan, X. E., Tang. B., Taylor, R., Williams, A., Kuan, J. X., Xu, P., Yan, Z., Zarov, I., Zhang, Y., Fan, A., Kambadur, M., Narang, S., Rodriguez, A., Stojnic, R., Edunov, S., and Scialom, T. Llama 2: Open Foundation and Fine- Tuned Chat Models, July 2023.
Zeng, Y., Liu, G., Ma, W. Yang, N., Zhang, H., and Wang, J. Token-level direct preference optimization. arXiv preprint arXiv:2404.11999 2024.
Zhang, Z., Zhang, Q., and Foerster, J. PARDEN, Can You Repeat That? Defending against Jailbreaks via Repetition, May 2024.
Zhong, Y., Ma, C.,, Zhang, X., Yang, Z., Chen, H., Zhang, Q. Qi, S., and Yang, Y. Panacea: Pareto alignment via preference adaptation for 1lms. arXiv preprint arXiv:2402.02003, 2024.
Zhou, Z., Liu, J., Shao, J., Yue, X., Yang, C., Ouyang, W., and Qiao, Y. Beyond one- mne-preference-itte
alignment: Multi-objective direct preference optimization. In Find- ings of the Association for Computational Linguistics ACL 2024, 2024.
Wachi, A., Tran, T. Q.. Sato, R., Tanabe, T., and Akimoto, Y. Stepwise Alignment for Constrained Language Model Policy Optimization, October 2024.
Wallace, E., Feng, S., Kandpal, N., Gardner, M., and Singh, S. Universal Adversarial Triggers for Attacking and Ana- lyzing NLP, January 2021.
Zou, A., Wang, Z., Carlini, N., Nasr, M., Kolter, J. Z., and Fredrikson, M. Universal and Transferable Adversarial Attacks on Aligned Language Models, December 2023.
Wang, Y. Shi, Z., Bai, A., and Hsieh, C.-.. Defending LLMs against Jailbreaking Attacks via Backtranslation, June 2024.
Zou, A., Phan, L.., Wang, J., Duenas, D., Lin, M., An- driushchenko, M., Wang, R., Kolter, Z., Fredrikson, M., and Hendrycks, D. Improving Alignment and Robustness with Circuit Breakers, June 2024.
Wei, Z., Wang. Y., Li, A., Mo. Y., and Wang, Y. Jailbreak and Guard Aligned Language Models with Only Few In-Context Demonstrations, May 2024.
Wen, Y., Jain, N., Kirchenbauer, J., Goldblum, M., Geip- ing, J., and Goldstein, T. Hard Prompts Made Easy: Gradient-Bass Discrete Optimization for Prompt Tun- ing and Discovery. June 2023.
Xie, Y., Yi, J., Shao, J., Curl, J., Lyu, L., Chen, Q., Xie, X., and Wu. F. Defending ChatGPT against jailbreak attack via self-reminders. Nature Machine Intelligence, December 2023.
Yang, A., Yang, B., Hui, B., Zheng, B., Yu, B., Zhou, C., Li, C., Li, C., Liu, D., Huang, F., Dong, G., Wei, H., Lin, H., Tang, J., Wang, J., Yang, J., Tu, J., Zhang, J., Ma, J.,
12

---

Learning Safety Constraints for Large Language Models
A. Entropy-based Heuristic Facet Assigmment Strategy
Following the Convex Polytope Machine (Kantchelian et al., 2014) algorithm, we use an entropy- heuristic strategy to select constraint violations while preventing overreliance on individual facets. For each unsafe example where y = =-1, rather than simply choosing the facet with maximum violation max; ( (Φx Ex), the method considers the set of valid facets aaddd whose violations exceed threshold 5:
Kvalid
Ek >i>>>7 7} [k: Φ
(3)
The entropy of facet assignments is computed across the set of unsafe examples 24 as:
K
((C)
Σ
Pk loga Pk
(4)
where (aEM)k=a(a)] is the proportion of unsafe examples assigned to facet k.
Algorithm Entropy- based facet Assignment
Require: Feature matrix f, labels y, constraint matrix Φ, thresholds f
Require: Valid facet threshold т, entropy threshold Mangat max attempts T
1: U < {i: -1)
Set of unsafc examples
2: k; t arg max; - E.) for all i E u
Initial assignments
3: + ComputeEntropy (A: hew)
4: while H < hurget and attempts < T do
5: Sample randomly from U
6: Kvalid < {k: 6 Ex > 7}
7: Sample knew randomly from KCvalid
8: Hnew + Cotpputentroppy (ki, (eu) U {kiow)) )
9: if Hee H then
10:
K; f- knew
11:
H +- Hnew
end if
end while
output Facet assignments (kajiez)
The algorithm iteratively attempts to increase assignment entropy by randomly selecting alternative facets from Kvalid while maintaining strong constraint violations. The target entropy threshold Marget is heuristically set to 1 2 log, K, where is the total number of facets. 7 set to 0. This approach ensures diverse facet utilization across unsafe examples while preserving the effectiveness of the safety constraints.
B. Steering via Lagrangian Relaxation
Building on the learned polytope facets, SaP implements directional steering in the model's representation space to constrain generation while minimizing disruption to the model's capabilities. This approach enables fine-grained safety control through local activation editing, avoiding the need for model fine-tuning or prompt engineering.
One way to solve the constrained problem in Equation (2) is through a Lagrangian relaxation. To this end, in practice, we perform 100 gradient steps on the following loss function
min Fe (x) + Σ
OJ Eh) -
Eo + Aunsafe Σ Ec -G
where is the edited activation, (h(h) is its corresponding encoded feature, [:]++ = max(0, 0,,) -) denotes the positive part, and [-]- = min(0, denotes the negative part. The hyperparameters Asate and Aunsate control the penalties for negative and positive violations respectively, allowing for asymmetric treatment of constraint violations.
13

---

Learning Safety Constraints for Large Language Models
C. Harmbench Training Setup and Results
C.1. Harmbench Training Setup
In Harmbench (Mazeika et al., 2024), we save the per-token representation at layer 20 during the original model output generation (with attacked inputs). The total training dataset size is num.sent x Eokens.per.senten an output sentence is marked as unsafe, we label all token representations in the sentence as unsafe, and vice versa. For each model (Llama, Mistral, Qwen), we use its default system prompt template when inferencing both the original and the steered model.
For the polytope training phase, we use the Adam optimizer with learning rate of 10-2 and batch size of 128. The feature extractor projects hidden states to (6.384-dimension)) space followed by ReLU activation. The loss function uses an entropy weight 1.0 and feature 1.1 regularization weight A = 1.0. The margin parameter a. varies across model architectures: 60.0 for Llama-2 7B, 5.0 for Ministral 8B, and 30.0 for Qwen2 1.5B, reflecting different geometric requirements their respective representation spaces.
During the steering phase, we apply hidden state optimization at layer 20 with model-specific configurations. For Llama-2 7B, we set Aunsate = 4.0 and Aaae 10-4 for the optimization objective. Ministral 8B uses Aunsafe = 0.25 without safety violation penalty (Aaaee = 0). For Qwen2 1.5B, we set Aunsate = 10.0 and Xaae = 5000. All models are quantized to 16-bit precision (float16 for Llama-2 and b for Ministral and Qwen) to reduce memory requirements while maintaining numerical stability.
C.2. Harmbench Main Results
Figure 2 is created by aggregating the average of each defense method's performance over all attack algorithms. We present detailed Llama-2 7B results in Table and Table 3, Ministral-8B results in Table 4 and Table 5, and Qwen2-1.55 1.5B results in Table and Table 7. We perform experiments in the same setup with prompt-based baselines Self Reminder (Xie et al., 2023), Response Check (Wang et al., 2024), and SmoothLLM (Robey et al., 2024), implemented in the BaaTTranslation code base (Wang al., 2024).
We also implemented rejection sampling, which uses the polytope classifier. Instead of steering when token is detected unsafe, it immediately modifies the current token output to be "Sorry". It empirically works well for defense, but might result in over-rejecting reasonable reguests, as shown in the MMLU results (Table 8). All methods' results (mean + standard deviation) are obtained over seeds. Table 10 shows results over attack algorithms for each model, when trained with and without the concept encoder.
Table 2: Defense comparison for Llama2-7B across different attack methods. Lower values indicate better defense effectiveness (0 means perfect defense). Each method evaluated over seeds.
Attack Method
Original
Steering
Rejection
Response Check
Self Reminder
SmoothLLM
AutoDAN
1.77
0.00 + 0.00
0.00 0.00
0.51 # 0.00
0.00
2.33 + 0.33
AutoPrompt
16.5
0.00 # 0.00
0.05 0.11
8.25 # 0.00
1.25
10.30 1.58
DirectRequest
1.0
0.00 = 0.00
0.05 +.11
0.50
0.00
6.50 0.64
GBDA
0
0.00 # 0.00
0.00
0.04
0.00
4.21 # 0.39
GCG
30
0.40 = 0.29
0.22
1.000
2.50
12.10 0.68
HumanJailbreaks
0.5
+ 0.00
0.01
0.01
0.05
2.19
PEZ
0.00
0.00
0.00
0.00
4.03 0.48
UAT
6.5
0.00
0.00
2.00
0.50
8
AdaptiveAttack
60
2.00
0.00
0.00
2.00
2.04
14

---

Learning Safety Constraints for Large Language Models
Table 3: MLP defense comparison for Llama2-7B across different attack methods, evaluated over seeds. Lower values indicate better defense effectiveness.
Attack Method
Original
MLP
MLP+5
AutoDAN
1.77
0.25
0.00
AutoPrompt
16.5
1
12.85 0.14
DirectRequest
1.0
0.11
0.25 =
GBDA
0
0.00
0.00
GCG
30
27.10
2.290
0.5
0.266
0.02
PEZ
0.00
0.00
UAT
6.5
4.75
4.75 0.00
AdaptiveAttack
60
0.00 = 0.00
0.00 + 0.00
Table 4: Defense comparison for Mistral-8B across different attack methods. Lower values indicate better defense effectiveness (0 means perfect defense). Each method is evaluated over seeds.
Attack Method
Original
Steering
Rejection
Response Check
Self Reminder
SmoothLLM
AutoDAN
66.75
1.35 + 0.49
0.00 + 0.00
1.50 = 0.00
11.50 + 0.00
19.38 # 1.94
AutoPrompt
48.3
2.79 # 1.64
0.51 + 0.25
7.88 = 0.00
13.94 + 0.00
40.36 + 2.63
DirectRequest
42.8
1.20 + 0.54
0.33 + 0.13
5.25 # 0.00
2.50 0.00
16.30 2.18
GBDA
62.8
1.26
0.23
4.377
6.00
36.75 8.71
GCG
67.5
1.90
0.08 0.13
5.00
26.50
32.35 2.00
HumanJailbreaks
39.4
1.15
0.01
0.13
9.65
17.41 + 0.88
PEZ
33.6
1.44
0.00
41.38
# +
27.00 10.63
UAT
40.8
2.45 1.04
0.24
6.00
4.50 # 0.00
24.65 3.09
AdaptiveAttack
100
14.00 + 0.00
0.00 + 0.00
97.33
90.00 + 0.00
97.33 + 0.94
Table 5: MLP defense comparison for Ministral-8B across different attack methods, evaluated over 5 seeds. Lower values indicate better defense effectiveness.
Attack Method
Original
MLP
MLP+5
AutoDAN
66.75
16.90 5.14
17.80 + 1.34
AutoPrompt
48.3
74.30 2.04
73.58 # 1.58
DirectRequest
42.8
20.50
20.35 # 1.29
GBDA
62.8
66.31 22.82
80.91 # 3.12
GCG
67.5
6.63
6..10
HumanJailbreaks
39.4
2.00
25.62 + 1.75
PFZ
33.6
32.72 12.77
32.81 + 15.03
UAT
40.8
49.90 # 1.07
49.25 + 1.38
AdaptiveAttack
100
16.00
16.00
D. BeaverTails Training Setup and Results D.I. Mutual Information Calculation Details
For BeaverTails (Ji et al., 2023), the Mutual Information (MD) between each safety category and each facet's violations computed using sci fickiiti tilkkilllle learn's mutual m informatti on score function. We compare the facets training from two polytopes, both share the same set of hyperparameters, and the only difference is whether the Concept Encoder is included in the training. They are both trained with 50 facets, trained on the full BeaverTails dataset over 20 epochs, optimized with Adam optimizer with learning rate 0.01, batch size 128, and A 0.01. The polytope trained with Concept Encoder has an encoded feature dimension 16384, Ap 1.0, and margin 2.0.
15

---

Learning Safety Constraints for Large Language Models
Table 6: Defense comparison for Qwen-1.SB across different attack methods. Lower values indicate better defense effectiveness (0 means perfect defense). Each method evaluated over seeds.
Attack Method
Original
Steering
Rejection
Response Check
Self Reminder
SmoothLLM
AutoDAN
45.00
0.80 0.82
0.67 0.00
38.25 0.00
30.50 0.00
6.85 1.27
AutoPrompt
17.90
1.16 0.55
2.83 1.09
23.74 0.00
3.54 0.00
6.77 + 1.20
DirectRequest
13.25
0.15 = 0.22
0.60 = 0.29
5.00
2.00
5.60 + 1.04
GBDA
19.45
1.42 # 0.87
20.63 =
3.00
8.54 + 0.63
GCG
25.00
1.00 = 0.69
2.44
$5,00
2.08
11.43 + 1.16
HlumanJailbredt
4.14
0.13
0.23
5.40
4.00
4.42 + 0.50
PEZ. 10.90
0.44
1.76
4.85 0.00
0.70
4.63 0.48
UAT 12.50
1.15 0.52
2.70 1.14
14.75 0.00
1.75
7.30 1.56
AdaptiveAttack
100
100.00 0.00
90.00 0.00
100.00 0.00
100.00
100.00 # 0.00
Table MLP defense comparison for Qwen-115B across different attack methods, evaluated over seeds. Lower values indicate better defense effectiveness.
Attack Method
Original
MLP
MLP+5
AutoDAN
45.00
3,75 1.56
12.85 14.36
AutoPrompt
17.90
8.23 1.62
14.29 + 12.29
DirectRequest
13.25
3.10 # 0.63
4.00 = 2.67
GBDA
19.45
9.92 # 1.44
16.80 = 1 11.80
GCG
25.00
10.57 3.35
18.43 # 21.13
HumanJailbreaks
4.14
2.26
4.17 2 4.38
PEZ
10.9
2.81 0.50
5.07 # 2.21
UAT
12.5
6.75 1.08
AdaptiveAttack
100
100.00 0.00
100.00 0.00
Table 8: MMLU accuracy (%) comparison across different defense methods. Results show that most defense methods maintain original model performance, with some exceptions for rejection and self-reminder strategies.
Method
Llama2-7B
Ministral-8B Qwen2-1.58 1.5B
Original
45.8
63.4
52.7
Steering
45.7
63.3
52.7
Rejection
44.6
28.5
51.2
ICL
45.8
65.1
55.4
Response Check
45.8
65.1
55.4
Self Reminder
45.8
22.9
22.9
SmoothLLM
45.8
65.1
55.4
MLP
45.8
63.4
55.4
MLP+5
45.8
63.4
55.4
D.2. BeaverTails Training Hyperparameters
conduct polytope training on all BeaverTails classes. We first obtain (human question, model answer) pairs from the using Python f-string without any other
dataset, formatting the input as f" (twmamnqqusttiioo. prompts. We observe that this works better for classifying safety concepts than other basic types of input prompting. We then obtain the activation for each sentence at model's layer / at the sentence's last word.
The polytope model is optimized using Adam with a learning rate of 10 and batch size 128. For facet assignment, we use a valid facet threshold 0.0. The concept encoder uses feature dimension 16.384. The training runs for 20 epochs.
16

---

Learning Safety Constraints for Large Language Models
Table 9: Model performance across turns in MT-Bench. There are fluctuations in LLM-based judges' scores, but all models produce coherent and high-quality sentences, without signs of rejections.
Model
First Turn
Second Turn
Average
Llama2-7B
6.54
5.46
6.00
Llama2-7B SaP
6.98
6.21
6.59
Ministral-8B
7.35
6.98
7.26
Ministral-8B B + SaP
9.01
7.64
8.33
Qwen-1.5B
6.90
5.15
6.03
Qwen-1 + SaP
6.29
4.70
5.49
Table 10: Attack success rate comparison between models, with and without Concept Encoder. Results show mean standard deviation across different attack methods.
Llama-2 7B
Ministral-80
Qwen2 1.5B
Method
Without CE
With CE
Without CE
With CE
Without CE
With CE
AutoPrompt
10.7 0.0
0.00 + 0.00
7.0
2.79 1.64
0.45 0.33
1,16 0.55
DirectRequest
0.0 0.0
0.00 # 0.00
2.0
1.20
0.15 0.14
0.15 0.22
GBDA
0.0
0.00 0.00
8.0
1.26
0.71 # 0.28
1.42 0.87
GCG
0.40 2.0
6.0
1.90
0.57 0.54
1.00 + 0.69
HumanJ HumanJailbreaks ailbreaks
0.3
0.00 0.00
2.07
1.15
0.08 + 0.08
0.18 + 0.13
PEZ
0.0 + 0.0
0.00
4.0
1.14
0.32 0.07
0.45 + 0.44
UAT
4.0 + 0.0
0.00 =
5.0
2.45 + 1.04
0.45 0.33
1.15 0.52
KLD for Facet 43
KLD for Facet 51
KLD Facet
1.53
gang
6.001
gang 6.020
0.000
terrorist
0.000
kill 0.000
0.000
homb
homb
kill
illegal
illegal
illegal
0.006
0.600
barm
0.600
1.25
KLD for Facet 3
KLD for Facet 29
KLD for Facet 57
20.081
kill
596
homb
1,365
kill
5.964
illegal
24444
kill
illggal
terrorist
6.000
gang
6.000
bomb
0.600
terrorist
0.287
bomb
0.000
0.000
illegal
0.041
harm
0.000
harm
0.600
barm
10
15
20
0.1
0.2
0.3
0.4
0.5
0.6
0.00
0.25
0.50
0.75
1.00
1.25
Figure 7: KL divergence analysis revealing semantic specialization of SaP's facets BeaverTaill terrorism dataset. Facets can activate one single concept (top row plots), multiple concepts (bottom row plots)
We empirically find that margin parameter x and the polytope sparsity weight As significantly influence model performance. Different safety categories and model architectures require different hyperparameter values for optimal performance. Table 11 shows hyperparameters for Llama2-17 7B, and Table and 13 show hyperparameters for Ministral-8B and Qwen2-1.5B respectively.
17

---

Learning Safety Constraints for Large Language Models
Table 11: Capegory-specific hyperparameeerss for BeaverTails for Llama2-7B. The last row "All data" shows the hyperparam- eters used to train all BeaverTails data.
Category
Margin
Animal abuse
0.5
0.01
Child abuse
0.7
0.01
Politics
3.0
0.0001
Discrimination
2.0
0.0001
Drug abuse
1.0
0.01
Financial crime
2.0
0.01
Hate speech
10.0
0.0001
Misinformation
0.8
0.0001
Unethical behavior
20.0
0.0001
Privacy violation
1.0
0.01
Self-harm
1.0
0.0001
Adult content
1.0
0.01
Terrorism
1.0
0.01
Violence
20.0
0.0001
All data
20.0
0.01
Table 12: Category-specino hyperparameters for BeaverTails for Ministral-8B. The last row "All data" shows the hyperpa- rameters used to train all BeaverTails data.
Category
Margin
>
Animal abuse
1.0
0.0001
Child abuse
0.8
0.0001
Politics
10.0
0.0001
Discrimination
10.0
0.0001
Drug abuse
2.0
0,0001
Financial crime
2.0
0,0001
Hate speech
1.0
0.0001
Misinformation
10.0
0.0001
Unethical behavior
15.0
0.0001
Privacy violation
2.0
0.0001
Self-harm
2.0
0.01
Adult content
0.6
0.01
Terrorism
0.6
0.01
Violence
2.0
0.01
All data
20.0
0.0001
D.3. BaiverTails Results
We present test accuracy results across different configurations of facets of Llama2-7B. Ministral-8B. and Qwen2-1.50 in Tables 14, 15, and 16 respectively. For each category, we additionally run learning on a multi-layer perceptron (MLP) classifier using the same representation inputs. For fair comparison, the MLP is defined to have similar model size with the polytope trained for each safety category. The MLP architecture has two layers and a classification head:
The first layer is (e) + bi), with ((( being the LLM activation at layer / (i.e., the same representation input as the one used polytope training). This layer has the same number of parameters as the concept encoder, with Wi TR 16384xd; b, E R 16381
We then add a second layer ReLU(W221 + b2), with W2 E R x 16384 with 21 being the output of the first layer, by E Re matching the parameter size of the polytope o and Ç.
18

---

Learning Safety Constraints for Large Language Models
Table 13: Category-specific hyperparameters for BeaverTails for Qwen-1.5B. The last row "All data" shows the hyperparam-
eters used to train all BeaverTails data.
Category
Margin
Animal abuse
15.0
0.0001
Child abuse
2.0
0.01
Politics
20.0
0.0001
Discrimination
10.0
0.0001
Drug abuse
15.0
0,0001
Financial crime
10.0
0.0001
Hate speech
0.8
0.01
Misinformation
0.8
0.01
Unethical behavior
15.0
0.0001
Privacy violation
2.0
0.0001
Self-harm
2.0
0.0001
Adult content
3.0
0,0001
Terrorism
0.9
0,0001
Violence
10.0
0.0001
All data
10.0
0.01
To perform classification, MLP has an additional classification head defined by Softmax(W;;2) + by), where 22 is the second layer output, W3 E 22x by E R2a. outputs the model's s estimated probability whether an input is safe or unsafe (i.e., classes). Hence, for each category, the MLP has Plaeet
2 more parameters than the SaP polytope.
The MLP is trained using cross-entropy loss and Adam optimizer with learning rate 0.0001. We present the results in Table 17. More KL divergence analysis on SaP's facets are presented in Table 7.
Table 14: Test accuracy (%) comparison across different numbers of facets for each category in Llama2-7B. Results show mean + standard deviation. Bold numbers indicate the selected configuration in final model.
Category
facet
10 facets
facets
30 facets
40 facets
facets
60 facets
Animal abuse
60.0 18.0
92.0 + 3.0
93.0 + 1.4
92.0 3.0
91.0 + 4.0
93.0 + 1.0
89.0 + 2.0
Child abuse
85.0 +.0 8.0
91.8 + 2.2
4.0
88.0
8.00
2.0
89.0 + 3.0
Politics
82.0 3.0
8
84.1 2.0
80.0
8.0
80.0
Discrimination
86.0
8..0
8.00
8.00
92
8.0
8.0
Drug abuse
9..7
89.0
9.0
90.0
90.0
92.0 91.0
Financial crime
1.0
92.0
90.0
8.00
88.0
91.0
90.0
Hate speech
8.0
8.73
82.0
8..0
84.0#2.0
80.0
8..0
Misinformation
5.0
7..5
71.0
72.0
67.0
69.0
7.0
Unethical behavior
6.0
8.00
8.00
8.00
8..5
1.
8..0
Privacy violation
8.00
9..0
9..0
92.0+20
93.0
91.0#20
92.0
Self-harm
3.0
93.0
91.0
93.0
92.0+3.0
92.02.0
9.65
Adult content
9.0
93
91.0
90.0
89.0
89.0
86.0
Terrorism
9.0
91.0
89.1
92.0
88.0
90.0
85.0
Violence
87.0
8..0
8..0
8..0
9..9
86.0
8..0
19

---

Learning Safety Constraints for Large Language Models
Table 15: Test accuracy (%) comparison different numbers of facets for each category in Ministral-8B. Results show
mean - standard deviation. Bold numbers indicate the selected configuration in final model.
Category
facet
10 facets
20 facets
30 facets
40 facets
50 facets
60 facets
Animal abuse
50.0 = 0.0
80.0 = 20.0
93.6 + 1,4
88.0 + 6.0
80.0 = 17.0
61.0 + 11.0
84.0 + 4.0
Child abuse
86.0 1.0
89.9 # 1.3
85.0 + 2.0
86.0 4.0
1.0
8..0
85.0 + 2.0
Politics
86.9 0.4
5.00
62.0 = 13.0
8.0
6.0
8.0
65.0 + 17.0
Discrimination
82.0 2.0
76.0
61.0 13.0
7.0
7..0
0.3
69.0 13.0
Drug abuse
93.7 0.8
14.0
84.0 14.0
73.0
73.0
82.0 19.0
88.0
Financial crime
92.6 1.0
11.0
9.00
9.00
9..0
9.0
91.0 1.0
Hate speech
86.0
0.7
8.00
85.0
87.0
81.0 81.0
Misinformation
56.0 12.0
1.1
7..0
7.00
7..0
7.00
68.0
Unethical behavior
8.00
8.0
8.68
8.0
8.00
82.0
8.0
Privacy violation
84.0 10.0
9..0
93.0
91.0=
91.0
93.0
91.0
Self-harm
9..0
6.0
9.65
91.0
92.0
89.0
93.0
Adult content
9.0
93
91.0
93.0
93.0
92.0
91.0
Terrorism
8.0
8.0
8.11
88.0
8.0
88.0
8.00
Violence
8..0
1.0
8.0
9.5
8..0
8..0
89.0 1.0
Table 16: Test accuracy (%) comparison across different numbers of facets for each category in Qwen2-1.5B. Results show mean + standard deviation. Bold numbers indicate the selected configuration in our final model.
Category
facet
10 facets
20 facets
30 facets
40 facets
50 facets
60 facets
Animal abuse
93.5 # 1.1
89.0 + 6.0
88.0 # 5,0
84.0 + 7.0
8..0
86.0 + 4.0
83.0 # 10.0
Child abuse
83.0 8.0
86.2 2 2.2
83.0 2.0
8.00 8..0
82.0 4.0
Politics
81.0
8..0
83.3 7..0 6.0 8.0
77.0
Discrimination
7..0
5.0
7..0
8.4
76.0
81.0 80.0
Drug abuse
9..0
6.0
92.6
8.0
90.0
89.0
8.00
Financial crime
9.00
8..0
9.0
89.0
89.0
89.0
8.2
Hate speech
8.00
8.2
8.0
83.0
85.0+2.0
8..0
8..0
Misinformation
50.0
6.00
7..0
7..5
69.0
6..0
6.00
Unethical behavior
7.0
8.00
0.7
8.00
80.0
8..0
1.0
Privacy violation
90.0
88.0
9.00
8..0
9..1
91.0
2.0
Self-harm
77.0 14.0
86.0
1.64
88.0
88.0
88.0
8.00
Adult content
1.0
9.55
8.00
8..0
86.0
88.0
8.00
Terrorism
8.00
1.00
8..0
88.0
88.0
8.67
8..0
Violence
8.0
81.0
72.0 + 15.0
8..
84.0
86.0
84.0
20

---

Learning Safety Constraints for Large Language Models
Table 17: Classification accuracy (%) on the BeaverTails dataset. Results show mean standard deviation across different safety categories. MLP uses architecture with similar amount of Safety Polytope parameters, but has additional parameters in the classification layer.
Llama2-77 7B
Ministral-8B
Qwen2-1.5B
Category
facet
SaP
MLP
1 facet
SaP
MLP
facet
SaP
MLP
Animal abuse
60.0 + 18.0
93.0 1.0
96.0
50.0 0.0
93.6 1.4
95.5
93.5 1.1
93.5 1.1
94.8
Child abuse
8.0
91.8 2.2
92.8
86.0
89.9 1.3
91.7
83.0
86.2 2.2
87.6
Politics
82.0
84.1
84.1
8..
8.49
86.8
8.00
83.3
85.1
Discrimination
8.00
1.8
88.9
8..0
0.3
88.4
7.00
8..3
87.0
Drug abuse
93.7
9.87
93.7
9..7
90.0
9,66
93.6
Financial crime
92.6
9.06=100
94.0
92.6
9..6
94.7
90.0
9.00
93.5
Hate speech
8.00
8.7
88.5
8.00
8..3
88.3
82.0
8..0
86.9
Misinformation
55.0
7.15
71.7
16.0
7.15
77.6
50.0
7.45
74.5
Unethical behavior
83.0
0.6
88.8
8.0
8.6
88.3
81.0
8.7
86.8
Privacy violation
8..0
94.0
94.9
84.0 10.0
9.00
90.0
9.41
94.7
Self-harm
90.0
93.0
92.6
91.0
9..0
94.3
77.0
9.
89.7
Adult content
9..0
9.0
92.7
9.00
93.5
94.1
8.00
9.5
91.3
Terrorism
9.00
9..0
92.7
8.0
9.0
90.0
8.00
92.3
Violence
8..0
9.59
91.6
8.0
9.5
91.9
8..0
8.3
90.6
All data
81.2
8.4
83,9
80.2 0.5
828
83.9
0.8
80.3 1.0
83.1
21

---

Learning Safety Constraints for Large Language Models
E. Polytope Sample Complexity
We present the current theoretical results of learning a polytope. Note that these results make important assumptions: (i) reliance on certain restricted encoder nature; (ii) perfect data separability by the polytope; (iii) the feasibility of finding a solution.
The following theorem is derived from (Gottllieb al., 2021), which establishes the fat-shattering dimension of y-fat I-polyhedra.
Theorem E.1 (Theorem 2.2, Gottlieb et al. (2021)). The class --fat t-polyhedra in R with at most t hyperplanes, has fat-shattering dimension of order:
(tlogt] 722
D=0
where is the number of hyperplanes defining the polyhedron, and "y is the margin.
This result provides the shatttrag dimension D for our hypothesis class (the class of y-fat polytopes), which we will use derive the sample complexity.
The second theorem we rely is general result from Anthony & Bartlet (1999), which gives sample complexity bound based on the fat-shattering dimension.
Theorem E.2 (Theorem 19.1, Anthony & Bartlett (1999)). Suppose that F is a class of functions mapping from a domain X into the real interval [0, 1]. and suppose also that F has finite fat-shattering dimension. Let A be any approximate-SEM algorithm for F and define, for z E T L(z) = A (z, 6n/6), where co 16/ /mm Then T, is a learning algorithm for F, and its sample complexity satisfies
256 m(c, 6) S f8 In²
128
16
In
256
for all , > 0.
Using big-O notation, we simplify this to the following form:
m(c,8) 0 ( fatp (c)
( + log
We now state and prove the main result, which gives the sample complexity bound for learning f-fat polytopes.
Corollary E.3 (Sample complexity of y-fat polytopes). Let H be the class of y-fat t-polyhedra in Rd. The sample complexity m(c,5) required to learn polytope from H with accuracy and confidence 1 - is bounded by:
mle, = -0 (tlog
log" (2)
8 ))
(
+ log
Proof. By Theorem E. the fat-shattering dimension of the class of f-fat t-polyhedra is:
D D= (tlog (tllogt 72
Substituting this into the sample complexity bound from Theorem E.2, we get:
m(c, = 0 tlogt t log²
(t) + log
))
1
8
Thus, the sample complexity of learning a 7-fat polytope is as stated.
22


Reviews:
[{"review_id": "9qP7Pkh7qg", "reviewer": "Reviewer_BjPZ", "summary": "The paper presents a novel approach to increase safety and the adversarial robustness of LLM. Instead of fine-tuning the parameter of the model for safety alignment, the introduced approach SaP (Safety Polytope) is applied during inference by enforcing linear safety constraints using Convex Polytope Machines in the model's representation space or—inspired by sparse autoencoders—a projection (the concept encoder) of the model representation to a higher but sparse dimensional space. By steering the activation subsequently of detecting the representation of a unsafe concept SaP influence the next token prediction during the sampling process.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for their thoughtful analysis and constructive feedback. We appreciate your recognition of our paper's key contributions:\n\n1. The novel geometric approach to LLM safety through representation space constraints.\n\n2. The effectiveness of SaP in defending against adversarial attacks while maintaining model capabilities\n\n3. The interpretability insights provided by our analysis showing how different safety constraints specialize in detecting specific types of unsafe content.\n\nBelow, we address your specific questions and suggestions:\n\n**Regarding experimental methodology:**\n\n> Question: \"Could you clarify the number of seeded runs of the different experiments?\"\n\nWe apologize for the confusion. This is indeed a typo in Appendix C.2. All methods (including baselines) were evaluated over 5 seeds. As shown in Tables 1-4 in the appendix, we report means and standard deviations calculated over 5 seeds for all methods. We will correct this inconsistency in the revised manuscript.\n\n**Regarding computational overhead during inference:**\n\n> Question: \"While the authors state that 'our approach scales efficiently to large batches via existing tools for vectorized computation' the computational overhead during inference is unclear. Therefore some uncertainty of the approach's practicability remains.\"\n\nThank you for raising this important point. We have conducted additional experiments to quantify the inference cost of SaP. First, we measured the average per-token processing time with and without SaP across all three model architectures:\n\n| **Model**    | **Configuration** | **Avg Time per Token (s)** | **Overhead** |\n| ------------ | ----------------- | -------------------------- | ------------ |\n| Llama2-7B    | with SaP          | 0.0301                     | +29%         |\n| Llama2-7B    | without SaP       | 0.0234                     | -            |\n| Ministral-8B | with SaP          | 0.0381                     | +8%          |\n| Ministral-8B | without SaP       | 0.0353                     | -            |\n| Qwen-1.5B    | with SaP          | 0.0309                     | +35%         |\n| Qwen-1.5B    | without SaP       | 0.0228                     | -            |\n\nWhile SaP adds approximately 8-35% overhead at the per-token level, this efficiency advantage becomes even more significant when comparing end-to-end runtime for practical tasks. For the MMLU benchmark, all runs conducted on A100 40GB GPUs via clean slurm jobs with no other program interference. For these baselines, we use the implementation from the llm-jailbreaking-defense [benchmark](https://github.com/YihanWang617/llm-jailbreaking-defense).\n\n| **Method**                 | **Total Runtime** | **Relative to Baseline** |\n| -------------------------- | ----------------- | ------------------------ |\n| Llama2-7B (baseline)       | 24 min            | 1.0×                     |\n| Llama2-7B + SaP            | 38 min            | 1.6×                     |\n| Llama2-7B + ICL            | 11h               | 27.5×                    |\n| Llama2-7B + Response Check | 11h 17min         | 28.1×                    |\n| Llama2-7B + Self Reminder  | 10h 42min         | 26.7×                    |\n| Llama2-7B + SmoothLLM      | 11h 35min         | 28.8×                    |\n\nThese results demonstrate that SaP provides a dramatically more efficient safety mechanism compared to prompt-based alternatives, requiring only 1.6× the baseline runtime while other methods require 26-29× more time. This substantial efficiency advantage stems from SaP's direct manipulation of model representations rather than relying on multiple forward passes or additional inference steps that characterize most prompt-based approaches.\n\n**Regarding additional comparisons with related methods:**\n\n> Concern: \"[...] lacks a comparison to related inference time steering methods such as 1) steering vectors and 2) especially SAEs.\"\n\nWe appreciate this valuable suggestion. Our concept encoder indeed shares implementation similarities with Sparse Autoencoders (SAEs), and we view the polytope constraint mechanism as offering a complementary geometric perspective compared to steering vectors.\n\nWe agree these comparisons would strengthen our paper and plan to:\n\n1. Provide a more thorough analysis of how our concept encoder relates to and differs from SAEs and steering vectors\n2. Incorporate the suggested references and Wu et al. (2025) in our revised manuscript\n3. Expand our discussion of related inference-time approaches to better position our work within this literature\n\nWe believe these additions will address your concerns and enhance the paper's contribution to the field of LLM safety. Having addressed your concerns, we would appreciate your feedback during this discussion period. Let us know if there are any further questions that we can clarify, otherwise, we would appreciate it if you would consider increasing your score.", "claims_and_evidence": "The paper's primary claim is that SaP provides an inference-time safety mechanism for LLMs that increases safety while maintaining model accuracy.\n\nThis is supported by:\n\n1) Empirical evidence on the harmbench dataset while applying adversarial attacks to elicit unsafe behaviour and simultanously evaluating maintaining general LLM performance using MMLU. Considering both SaP outperforms the selected baselines.\n2) Further, the paper provides an interpretability analysis demonstrating that different safety constraints specialize in detecting specific types of unsafe content.\n\nHowever, the empirical evaluation lacks a comparison to related inference time steering methods such as 1) steering vectors (see [1] and [2] or see [3] for an overview of other related detection/steering approaches) and 2) especially SAE which the presented approach even draws inspirations from. \n\n**References**\n\n[1] Wang et al. Model Surgery: Modulating LLM's Behavior Via Simple Parameter Editing. (2024):  https://arxiv.org/abs/2407.08770v1 \n\n[2] Lee et al. A Mechanistic Understanding of Alignment Algorithms: A Case Study on DPO and Toxicity (ICML 2024): https://arxiv.org/abs/2401.01967\n\n[3] https://arxiv.org/pdf/2501.17148 While the benchmark was probably released after the ICML submission deadline, the paper provides an overview of (safety) detection/steering algorithms related to the presented method.", "methods_and_evaluation_criteria": "Yes, the choice of evaluation datasets makes sense for the type of assessment. The authors evaluate the introduced approach on the harmbench dataset while applying a range of adversarial attacks to elicit unsafe behavior. Additionally, the MMLU benchmark is used to demonstrate that the general performance of the model is maintained. Further, the introduced approach is compared against five rejection baselines to demonstrate the advantages of SaP. However, as described above, more related inference-time approaches, such as steering vectors or SAEs, should be considered.", "theoretical_claims": "The paper defines LLM safety as a CMDP problem, where constraints define safe and unsafe regions in the representation space. The claims are supported by references to prior work.", "experimental_designs_or_analyses": "In general, the experiments are well-structured, though additional comparisons with related safety methods would strengthen the evaluation.\n\nSpecifically, the experimental section evaluates SaP on three LLMs, namely Llama2-7B, Ministral-8B, and Qwen2-1.5B. \nAdditionally, ablations on the impact of the concept encoder and number of safety constraints are conducted.", "relation_to_broader_scientific_literature": "The paper is well-grounded in prior literature on LLM safety, and specifically, Section 5 describes the relation to prior literature quite well. However, even if inspired by SAEs, a comparison is missing. Further, the paper lacks a discussion on the relation to other recent inference-time steering methods, as mentioned above.", "essential_references_not_discussed": "See above.", "other_strengths_and_weaknesses": "**Other Strengths**\n- Limitations and future work are well discussed.\n\n**Other Weaknesses**\n- While the authors state that \"our approach scales efficiently to large batches via existing tools for vectorized computation“ the computational overhead during inference is unclear. Therefore some uncertainty of the approach’s practability remains. A additional analysis of the computational overhead during inference would strength the paper.", "other_comments_or_suggestions": "I suggest considering an additional comparison to the above mentioned safety methods.", "questions_for_authors": "- See comment in supplementary material section:\nCould you clarify the number of seeded runs of the different experiments and why you chose a different number of runs for the baselines?\n\n- Could you provide an overview of the additional computational overhead during inference?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "1dASXm1Osc", "reviewer": "Reviewer_FoeD", "summary": "The paper introduces SaP, a post-hoc safety mechanism that defines a convex polytope in an LLM’s feature space. Using a Concept Encoder to disentangle safety-related features, it learns linear constraints that steer unsafe outputs into a safe region without retraining the model. Experiments show that SaP dramatically reduces adversarial attack success rates while maintaining overall model performance.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for their thoughtful analysis and valuable feedback.\n\nWe would like to first clarify that SaP is not primarily a jailbreak defense paper. Our key contribution is reformulating safety as a geometric constraint learning problem in representation space, which provides both interpretable insights and an effective control mechanisms. Throughout the paper, we consistently frame our approach in terms of \"safety\" and \"constraints\" in their general sense. Please see our reply to Reviewer aHnQ for an expanded explanation on this.\n\n**Specific Responses to Concerns**\n> \"Overfit Problem\"\n\nWe thank you for this insightful comment. We agree that only training on GCG will cause SaP to overfit to its specific attack pattern. However, this is not directly a weakness of SaP, as it can be easily mitigated by training it with data collected from diverse sources. This is to be expected, as machine learning models, SaP being no exception, are generally only as good as their training data (no free lunch). \n\nTo address the concern about overfitting to specific attack patterns, we conducted additional experiments with attacks not seen during training. The polytopes used in the following experiments are the same ones we trained for reporting the main results in the paper, with no access to any new information.\n\nAutoDAN ASR (5 seeds):\n\n| **Method** | **Ministral-8B** | **LLaMA2-7B** | **Qwen2-1.5B** |\n| - | - | - | - |\n| Original | 66.75  | 1.77 | 45 |\n| Original + SaP | 1.35 ± 0.49 | 0.00 ± 0.00 | 0.80 ± 0.82 |\n\nAdaptive attack:\n\n| **Model Configuration**       | **Attack Success Rate (%)** |\n| - | - |\n| Llama2 Original | 100 |\n| Llama2 + SaP | 0 |\n| Ministral Original | 100 |\n| Ministral + SaP | 98 |\n| Ministral + SaP (increased λ) | 12 |\n| Qwen2 Original | 100 |\n| Qwen2 + SaP | 88 |\n\nThe adaptive attack experiments were conducted on AdvBench, following the original released code implementation, which contains requests not seen in HarmBench.\nTo the best of our knowledge, three papers to date report defense performance on adaptive attacks, and the most comparable result is from Yi et al., 2025, where they also train their defense methods on HarmBench and test on AdvBench for adaptive attacks. Their best reported performance across 9 defense methods on Llama3-8B is 61% ASR (Table 2 of Yi et al.), while our method achieves 0% ASR for Llama2-7B.\nFor Ministral, we found that a direct transfer of SaP hyperparameters trained on HarmBench doesn't defend well against adaptive attacks (98% ASR), but increasing $\\lambda$ (the unsafe penalty) can substantially reduce this to 12% ASR. For Qwen2, we discovered that adaptive attack learns to induce harmful responses in Chinese, which our English-trained SaP can't effectively counter. This highlights an interesting direction for future work on multi-lingual adversarial defense through geometric constraints.\n\n> \"Refusal Pattern After Training\"\n\nThroughout numerous adversarial inputs, the polytope learns to steer models toward safety in a nuanced way. We observe that it guides models to incorporate common refusal phrases like \"Sorry\", \"can't answer this question\", \"I'm just an AI\", etc. We are happy to include qualitative examples in the revised manuscript, if requested.\n\n> \"Unfair Baseline Comparison\"\n\nWe appreciate this concern. Our additional experiments with AutoDAN and adaptive attacks address this by testing on both attack types and adversarial requests not seen during training (generalization from HarmBench to AdvBench). These results strongly support the generalization capabilities of SaP beyond specific patterns used in training.\n\n> \"Simple MLP may also work\"\n\nFrom our experiments on BeaverTails, we do see that SaP performs similarly to an MLP with comparable parameter count. However, rejecting based on binary classification might result in over-rejection, which we experimented with and reported in Figure 2, Rejection Sampling (RJ). It reduces Ministral's MMLU performance from 63.4% to 28.5%. Furthermore, polytope constraints offer interpretability benefits that are not possible with a standard MLP, as discussed in our paper's Section 4.2.\n\n> \"MMLU does not prove that standard performance is maintained\"\n\nRegarding your suggestion to evaluate on benchmarks like OR-Bench, we see this as a valuable future direction. In our preliminary investigations with sensitive but benign prompts, we observe that SaP maintains high performance with minimal false positives. We would like to highlight again that the main contribution of this paper does not directly concern proposing a new SoTA defense method, but rather to propose a new framework for explicit modeling of safety in LLMs.\n\nWe thank you once again for your feedback and hope our responses adequately address your concerns. We'd be happy if you would consider revising our score, and we are happy to answer any other open questions. Thanks again for your feedback and active engagement in the review process.", "claims_and_evidence": "The claims are supported by a series of experiments on multiple LLMs, but I am not confident in some of them; please see the section for questions for the authors.", "methods_and_evaluation_criteria": "it makes sense.", "theoretical_claims": "I have checked the math in the main paper, and it looks reasonable to me.", "experimental_designs_or_analyses": "I have some concerns regarding the experimental design/ evaluation; please see the questions for the authors section.", "relation_to_broader_scientific_literature": "The paper builds on work in constrained decision processes and safe RL by framing safety as linear constraints in LLM representations. It leverages ideas from convex optimization and interpretability research to propose a novel post-hoc safety mechanism.", "essential_references_not_discussed": "all related works have been cited, and [1] from anthropic is also quite related to the topic shown in the paper (though it is released after the submission deadline, it could be added in the future version)\n\n[1] Sharma, M., Tong, M., Mu, J., Wei, J., Kruthoff, J., Goodfriend, S., ... & Perez, E. (2025). Constitutional classifiers: Defending against universal jailbreaks across thousands of hours of red teaming. arXiv preprint arXiv:2501.18837.", "other_strengths_and_weaknesses": "Strengths:\n1. introduce geometric safety constraints to detect harmful response.\n2. post-hoc mechanism preserves model performance while dramatically reducing adversarial attacks.\n3. provides interpretable safety facets, offering insights into distinct harmful content categories.\n\nWeaknesses:\n1. the LLMs used in the paper are somewhat outdated; Llama 3.1 8B would be preferable to Llama2-7B, but I admit it is just a minor problem.\n2. more discussion on the influence of the LLM's size on the proposed method is expected.\n3. the paper only considers gradient-based attacks (which optimize for unreadable adversarial strings) and does not consider other attack methods like AutoDan, which optimize for readable adversarial strings.", "other_comments_or_suggestions": "the writting looks good", "questions_for_authors": "Below is my main concern about this paper. If I have misunderstood anything, please correct me, and I would be happy to raise my score if the authors address these concerns well.\n\nFrom my perspective, the proposed method first trains the polytope based on the collected data (using only the last hidden representation of the final token). Once the polytope is obtained, during evaluation, it is applied at each token's position; if a token's representation is out of bounds, it is mapped back into the polytope. I then have several questions:\n\n(1) **Overfit Problem**: when training the polytope using gradient-based methods such as GCG, there is a concern that it might overfit to the specific attack pattern rather than learning robust safety constraints. GCG typically generates adversarial suffixes that are unreadable and have high perplexity. As a result, the hidden representations capture these characteristics through self-attention, and the polytope may rely on them for safe/unsafe classification instead of truly learning generalized safety constraints. This raises questions about defense transferability: if the polytope is trained on gradient-based attacks, will it effectively defend against other attack types such as human-crafted jailbreaks, AutoDan [1], or rule-based attacks [2], and vice versa?\n\n(2) **Refusal Pattern After Training**: Building on the previous point, I am surprised that simply mapping the representation back to the trained-internal polytope—as outlined in Algorithm 1—can prevent harmful responses. However, the paper does not provide examples of what the resulting refusal responses look like. Will these responses resemble the traditional refusal pattern, such as “I can’t…”? I am curious about the LLM's response style when the representation is mapped back, and it also raises the question of whether enforcing the safety constraint on only the initial tokens is sufficient to prevent attacks, that will improve the efficiency and reduce false positive rate.\n\n(3) **Unfair Baseline Comparison**: SaP is trained using the top 3 most effective attacks, which are also included in the evaluation. This means that the most effective attack patterns are already \"leaked\" during training for SaP but not for the other baselines. I recommend that the authors include held-out attacks—such as AutoDan [1] or rule-based attacks [2]—to ensure a fairer comparison.\n\n(4) **Simple MLP may also work**: Following from (3), if one were to directly train a simple MLP for binary classification on the hidden representations using the representation from the same top 3 attacks, I doubt it could achieve a much lower ASR than SaP while maintaining standard MMLU accuracy. The motivation comes from [4].\n\n(5) **MMLU does not prove that standard performance is maintained**: It is not surprising that MMLU performance does not drop because the prompts in MMLU are benign and do not include topics like fraud, hate, or violence as seen in HarmBench. In such cases, the polytope—or even a simple MLP—can easily learn and separate these patterns. In my view, it would be more effective to demonstrate that the false positive rate remains low when the LLM with the polytope is prompted with benign questions involving sensitive topics such as fraud, hate, or violence. This would show whether the polytope mistakenly rejects benign responses involving sensitive topics and prove that it is enforcing constraints based on the harmfulness of the response rather than merely the topic.\n\n[1] Liu, X., Xu, N., Chen, M., & Xiao, C. AutoDAN: Generating Stealthy Jailbreak Prompts on Aligned Large Language Models. In The Twelfth International Conference on Learning Representations.\n\n[2] Andriushchenko, M., Croce, F., & Flammarion, N. (2024). Jailbreaking leading safety-aligned llms with simple adaptive attacks. arXiv preprint arXiv:2404.02151.\n\n[3] Zou, A., Phan, L., Chen, S., Campbell, J., Guo, P., Ren, R., ... & Hendrycks, D. (2023). Representation engineering: A top-down approach to ai transparency. arXiv preprint arXiv:2310.01405.\n\n[4] Sharma, M., Tong, M., Mu, J., Wei, J., Kruthoff, J., Goodfriend, S., ... & Perez, E. (2025). Constitutional classifiers: Defending against universal jailbreaks across thousands of hours of red teaming. arXiv preprint arXiv:2501.18837.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ZyudDH1mbu", "reviewer": "Reviewer_aHnQ", "summary": "The authors propose to map unsafe model responses to save model regions in representation space without adjusting the weights of the respective model. Specifically, they represent safety constraints via polytopes and filter responses by assessing the similarity of latent features to the learned polytope. The method restricts models from generating safe outputs while model capabilities are maintained. \n\nReferences for the remaining review:\n\n[1] Chao et. al., \"Jailbreaking Black Box Large Language Models in Twenty Queries\", 2023\n\n[2] Andriushchenko et. al., “Jailbreaking Leading Safety-Aligned LLMs with Simple Adaptive Attacks”, 2024\n\n[3] Schwinn et. al., \"Soft Prompt Threats: Attacking Safety Alignment and Unlearning in Open-Source LLMs through the Embedding Space\", 2024\n\n[4] Carlini et. al., \"Adversarial Examples Are Not Easily Detected: Bypassing Ten Detection Methods\", 2017\n\n## update after rebuttal\n\nMost of my concerns were addressed. I still believe it's highly relevant (even in the scope of non-adversarial evaluations) to assess the refusal/safety trade-off of this defense. XS-Test, for example, has only ~100 prompts and would give insights into the overrefusal behavior. \n\nI will followed the discussion and increased my score from 2 to 4 and recommend accepting this work.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank you for your thoughtful review. We particularly appreciate your recognition of our paper's key strengths:\n- The \"novel method to steer model activations from harmful regions to safe regions\"\n- Our \"several ablation studies to investigate the individual components of [our] mechanism\"\n- The comprehensive evaluation across \"several adversarial attacks\" and \"multiple models\"\n\n**Clarification on Paper Focus and Contributions**\n\nWe would like to first clarify that SaP is not primarily a jailbreak defense paper. Our key contribution is reformulating safety as a geometric constraint learning problem in representation space, which provides both interpretable insights (through facet specialization) and an effective control mechanism (which includes but is not limited to defense). This geometric perspective offers a principled framework for conceptualizing and controlling model safety behaviors. Throughout the paper, we consistently frame our approach in terms of \"safety\" and \"constraints\" in their general sense:\n1. In Section 2, we discuss how \"safe and ethical language not only tries to maximize a reward function, but is also subject to some constraints. For instance, humans naturally avoid using language that would hurt someone's feelings, incite harm, or solicit unlawful actions.\"\n2. Our extensive experiments on BeaverTails (Sections 4.2-4.3) span 14 distinct safety categories, including animal abuse, discrimination, privacy violations, and misinformation—far beyond just jailbreak prevention.\n\nThat said, we agree that stronger defense performance can better support our perspective of geometric constraint learning. We therefore conducted the additional experiments you requested.\n\n**Regarding Stronger and Adaptive Attacks**\nAs you suggested, we have evaluated our defense against stronger and more recent attacks. In our response to Reviewer FoeD, we report results on AutoDAN and adaptive attacks. Please refer to it for detailed setting and discussions. Notably:\n1. For AutoDAN, our method reduces ASR from 66.75% to 1.35% on Ministral-8B, from 1.77% to 0% on Llama2-7B, and from 45% to 0.80% on Qwen2-1.5B.\n2. For adaptive attacks (on AdvBench), our method achieves 0% ASR for Llama2-7B by directly transferring our trained polytope and hyperparameter settings. With a slight increase in unsafe penalty, Ministral-8B can be steered from 100% ASR to 12%.  \n\n**Regarding Utility Evaluation**\n> Reviewer concern: \"Could the authors add more suitable utility evaluation benchmarks to investigate model capability? E.g., MT-Bench and over-refusal benchmarks\"\n\nWe agree that MMLU alone may not fully represent model capabilities. Following your suggestion, we conducted additional evaluations on MT-Bench:\n\n| **Model** | **First Turn Score** | **Second Turn Score** |\n| - | - | - |\n| Ministral-8B | 7.55 | 6.98 |\n| Ministral-8B + SaP | 9.01 | 7.71  |\n| Llama2-7B | 6.54 | 5.46 |\n| Llama2-7B + SaP | 6.98 | 6.21 |\n| Qwen-1.5B | 6.90 | 5.15 |\n| Qwen-1.5B + SaP | 6.85 | 5.06 |\n\nWhile the results suggest that models with SaP maintain performance and can generate coherent sentences without over-refusing on MT-Bench, we do not intend to claim that SaP improves performance. The slight differences observed might be due to the inherent noise in GPT-based evaluations and the limited sample size (80 evaluation examples). We manually inspected the outputs and did not find any nonsensical or rejection responses. The results are attached in the following URL: https://limewire.com/d/asA5e#Rlx4uqoGNp\n\nRegarding your suggestion to evaluate on benchmarks like OR-Bench, we see this as a valuable future direction. In our preliminary investigations with sensitive but benign prompts, we observe that SaP maintains high performance with minimal false positives. We would like to highlight again that the main contribution of this paper does not directly concern with proposing a new SoTA defense method, but rather to propose a new framework for explicit modeling of safety in LLMs.\n\n**Conclusion**\nWe thank you for your constructive feedback. The additional experiments have strengthened our confidence in SaP's learned safety knowledge, which supports our central contribution: reformulating language model safety as a geometric constraint learning problem in representation space. We believe this geometric perspective offers a principled framework for conceptualizing and controlling model safety behaviors beyond just jailbreak defense. \n\nWe hope our response addresses the reviewer’s remaining concerns, and given that we could address all the other concerns raised in the review, we would appreciate it if the reviewer would consider increasing our score. We are happy to answer any other open questions. Thanks again for your active engagement!", "claims_and_evidence": "The authors claim that their proposed algorithm can prevent harmful behavior in LLMs by steering latent activations associated to harmfulness to safe regions in the latent space. \n\nThey back up their claims through a large-scale empirical study using 7 different adversarial attacks (excluding recent Sota attacks), three different LLMs, and two common harmfulness benchmarks.\n\nMoreover, they conduct several ablation studies to verify the contribution of individual design choices within their algorithm", "methods_and_evaluation_criteria": "- The benchmarks datasets are widely used in the research field and appropriate \n- The used adversarial attacks are mostly weak (GBDA, PEZ , Human Jailbreak, Direct Request). Even GCG may be viewed as weak considering that multiple more efficient GCG variations have been proposed in the literature since the original paper.\n- The utility benchmark is not suitable. LLMs can achieve high scores on MMLU without being able to generate coherent sentences. \n- Measures for harmfulness evaluations are appropriate", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "The results of the conducted experiments are consistent with the results reported in the literature. I did not find anything unusual. However, as described in Methods And Evaluation Criteria some of the experiments are not suitable to back up the claims of the paper.", "relation_to_broader_scientific_literature": "The authors position their paper appropriately within the related work. However, some references/comparisons with more recent adversarial attacks are missing e.g., [1, 2]", "essential_references_not_discussed": "I am not aware of a paper that I would deem essential to discuss in the context of the presented work that is currently missing from the paper.", "other_strengths_and_weaknesses": "**Strengths**\n- The authors propose a novel method to steer model activations from harmful regions to safe regions\n- The authors conduct several ablation studies to investigate the individual components of their mechanism\n\n**Between Strength and Weakness**\n- The authors use several adversarial attacks and conduct experiments on multiple models. However, used attacks are not sufficient\n\n**Weakness**\n- The authors only use MMLU to analyze model capabilities\n- The authors do not try to adaptively break their defense. Detecting adversarial examples (e.g., in latent space) has shown to be very difficult [4]. Adaptive evaluations are necessary to evaluate new defenses.", "other_comments_or_suggestions": "Text in Figure 4 is not readable. Can the authors increase the font size of the axis labels? Numbers are less relevant. Colorbar could be helpful as well.", "questions_for_authors": "- Could the authors conduct evaluations with stronger and more recent attacks?\n- Could the authors conduct a sanity check with very strong continuous attacks if they are able to bypass their defense? A negative result would point towards errors in their evaluation [3]\n- Could the authors add more suitable utility evaluation benchmarks to investigate model capability? E.g., MT-Bench and over-refusal benchmarks\n\n# After the rebuttal\n\nMost of my concerns were addressed. I still believe its highly relevant (even in the scope of non-adversarial evaluations), to assess the refusal/safety trade-off of this defense. XS-Test, for example, has only ~100 prompts and would give insights into the overrefusal behavior. \n\nI will follow the discussion and am currently leaning toward increasing my score to three if no other major concerns are raised by the other reviewers.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "mEVDFzoAac", "reviewer": "Reviewer_ZBtn", "summary": "The study proposes a geometric approach called SaP (Safety Polytope) for large language models (LLMs) to mitigate safety risks. SaP learns and enforces linear safety constraints directly in the model's representation space, identifying safe and unsafe regions. Experiments show it reduces adversarial attack success rates and provides interpretable insights into its safety mechanisms.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful review and accurate summary of our paper. We appreciate your recognition that our \"claims are supported by the evidence\" and that \"SaP indeed addresses the problem of learning safety constraints automatically for LLMs.\"\n\n**Regarding Scalability**\n\nWe appreciate your question about the scalability of our approach. We acknowledge that polytope learning presents inherent scalability challenges, which is precisely why we adopted the Convex Polytope Machine (CPM) algorithm instead of traditional methods like QuickHull. This design choice was deliberate to address the computational complexity issues in high-dimensional representation spaces.\n\nOur empirical results demonstrate CPM's effectiveness at scale:\n\nOn the BeaverTails dataset (330K examples), our method achieves performance comparable to MLPs with similar parameter counts (Appendix D.3, Table 12).\nFor the HarmBench defense task with ~890K data points, SaP achieves an average safety classification accuracy of 91.3% and maintains attack success rates below 5%, further validating its effectiveness at scale.\n\nBased on these results, we expect CPM to perform robustly (on par with MLP) on datasets of similar or even larger scale. In Section 6, we discuss future directions for improving beyond CPM, including references to promising work like Hashimoto et al. (2023) on neural polytope learning.\n\nFor practical applications, the inference-time steering mechanism remains efficient regardless of training data size, as it only requires a forward pass through the concept encoder followed by verification against the learned polytope facets.\n\n**Regarding Missing References**\n\nThank you for suggesting the valuable references (Li et al., 2023 and Song et al., 2024) on MDP modeling for LLMs. We will incorporate these references in our revised manuscript to strengthen the discussion on using MDPs as a theoretical foundation for language model behavior. These papers provide important context for our approach to modeling language generation as a constrained MDP.\n\nWe greatly appreciate your feedback and we are happy to address any further questions or comments.\n\nHaving addressed your concerns, we would appreciate your feedback during this discussion period. Let us know if there are any further questions that we can clarify, otherwise, we would appreciate it if you would consider increasing your score.", "claims_and_evidence": "The claims are supported by the evidence. For example, SaP operates post-hoc in the representation space without modifying model weights. The approach relies on modifying the internal activations of LLMs rather than retraining or fine-tuning.", "methods_and_evaluation_criteria": "The proposed method, SaP, indeed address the problem of learning safety constraints automatically for LLMs.", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "The paper evaluates SaP’s ability to reduce attack success rates against seven adversarial attack methods, including gradient-based (e.g., GCG, GBDA) and human-crafted jailbreaks.", "relation_to_broader_scientific_literature": "The paper address the gap of leverage token vectors to build a safety constraint for LLM.", "essential_references_not_discussed": "The MDP modelling is a key concept of the paper. While there is some exploration work on modelling LLMs as MDP is missing:\n- Li, Kenneth, Oam Patel, Fernanda Viégas, Hanspeter Pfister, and Martin Wattenberg. \"Inference-time intervention: Eliciting truthful answers from a language model.\" Advances in Neural Information Processing Systems 36 (2023): 41451-41530.\n- Song, Da, Xuan Xie, Jiayang Song, Derui Zhu, Yuheng Huang, Felix Juefei-Xu, and Lei Ma. \"Luna: A model-based universal analysis framework for large language models.\" IEEE Transactions on Software Engineering (2024).", "other_strengths_and_weaknesses": "The scalability of the work is doubtable. When the dataset gets very large, I am not sure whether the learned safety constraints are still effective.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "How scalable is the proposed method?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: Strong Accept because (1) the paper addresses an important area, namely safety for LLMs, (2) develops a novel geometric approach to enforce safety constraints with interpretable insights, (3) jailbreak defence experiments are convincing and (4) unanimously received accept or weak accept from all reviewers.
Decision: Accept (spotlight poster)