Geometry-Guided Constraint Learning for LLM Safety Classification

2026/06/09 by Fumiaki Uehara, Koo Imai, Masato Tsutsumi +3
#cs.AI

paper · pdf

Abstract

Safety as Polytope (SaP) learns linear half-space constraints in LLM hidden space but requires per-category tuning of the constraint count K. We show that sparse autoencoder (SAE) feature extraction resolves this: K=2 becomes optimal for 12/14 categories on Qwen3.5-9B, achieving 96-99% accuracy per category on our BeaverTails classification benchmark, largely eliminating the need for exhaustive sweeps (K=4-25 with random initialization). This convergence to two planes is consistent with the Linear Representation Hypothesis, providing suggestive evidence that safety boundaries in this setting admit a low-dimensional linear description in the SAE feature space. Building on this geometric perspective, we introduce a cone constraint whose learnable aperture adapts to each category's cluster concentration, stabilized by a three-phase training

Citations

• Affective Coherence Monitoring for Transformer-Based Language Models
• Training language models to follow instructions with human feedback

Related

• SLM: End-to-end Feature Selection via Sparse Learnable Masks
• Structure-aware error bounds for linear classification with the zero-one loss
• Constrained Machine Learning Through Hyperspherical Representation
• Safe Feature Elimination in Sparse Supervised Learning
• Sparse but Wrong: Incorrect L0 Leads to Incorrect Features in Sparse Autoencoders