Geometry-Guided Constraint Learning for LLM Safety Classification

[Submitted on 9 Jun 2026]

View PDF HTML (experimental)

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

Submission history

From: Koo Imai [view email]
[v1] Tue, 9 Jun 2026 08:47:00 UTC (112 KB)