Technical Summary: Geometry-Lite: Interpretable Safety Probing via Layer-Wise Margin Geometry

Problem Definition

As large language models (LLMs) are deployed at scale, prompt-level harmfulness detection has become a critical safety component. While hidden-state probes using linear classifiers on single-layer snapshots (typically the final prompt token) are standard baselines, they rely on the assumption that a single layer suffices to summarize safety-relevant signals. This assumption warrants reexamination for two reasons:

1. Borderline Cases: Aggregate performance on standard benchmarks often obscures behavior on ambiguous prompts, which are critical for safety systems operating at low false-positive rates.
2. Out-of-Distribution (OOD) Generalization: Models performing well on benchmark splits may rely on shortcut features that fail to transfer. Recent work indicates that activation-based safety probes suffer substantial degradation under cross-dataset evaluation.

The core research question is not merely whether multiple layers help, but what form of multi-layer evidence remains useful in decision-critical regimes (low false-positive rates) and under benchmark shifts. Specifically, the authors investigate whether the safety signal is distributed across layers, driven by layer-to-layer motion (displacement), or defined by the persistent position of a margin profile relative to safety boundaries.

Methodology: Geometry-Lite

The authors introduce Geometry-Lite, a compact prompt-level probe designed to decompose multi-layer safety signals along labeled geometric axes. The method operates in three stages:

1. Layer-Wise Margin Reduction

For each transformer layer $\ell$, the $D$-dimensional hidden state of the final prompt token, $h^{(\ell)}(x)$, is reduced to a scalar signed margin using three complementary geometric readouts:

• Centroid Distance: The difference in Euclidean distance to the class-conditional means ($\mu^+$ for safe, $\mu^-$ for unsafe).
  $$m^{cent}_\ell(h) = \|h - \mu^-_\ell\|_2 - \|h - \mu^+_\ell\|_2$$
• k-NN Local-Neighborhood: The difference in mean cosine distance to the $k$ nearest neighbors in the safe and unsafe training pools. This is the only non-linear readout.
  $$m^{knn}_\ell(h) = \bar{c}^-_k(h) - \bar{c}^+_k(h)$$
• Supervised Linear Boundary: An L2-regularized logistic regression boundary fitted per layer on standardized features, yielding a signed score.
  $$m^{lin}_\ell(h) = w^\top_\ell z + b_\ell$$

This process converts the hidden state trajectory $H(x) \in \mathbb{R}^{L \times D}$ into three scalar margin trajectories $m^G(x) \in \mathbb{R}^L$ (where $G \in \{cent, knn, lin\}$).

2. Profile Summarization via Named Axes

Each margin profile is summarized along three named axes to separate distinct forms of layer-wise signal:

• Margin Level: Captures the persistent boundary position (e.g., endpoint margins, extremal values, negative-side occupancy, and area).
• Layer-to-Layer Change: Isolates local movement via finite differences ($\Delta m_\ell = m_{\ell+1} - m_\ell$), separating motion from absolute position.
• Structural Shape: Captures coarse profile properties (e.g., first boundary crossing, variability of change, total path length, and directional consistency).

These axes generate a 39-dimensional feature vector $\phi(x)$ (13 features per geometry).

3. Classification

The 39-dimensional vector is fed into an L2-regularized logistic regression classifier. The parameter count is independent of the backbone model size and hidden dimension.

Baselines and Comparisons

The study compares Geometry-Lite against:

• Single-layer probes: Best-performing layer selected via validation.
• MultiLayer-DIM: Aggregates projections onto the difference-in-means direction across layers.
• MultiLayer-Linear: Aggregates raw supervised linear scores across layers.
• TaT-Disp-LSTM: A trajectory-based baseline modeling layer-wise displacements with an LSTM.

Key Contributions

1. Geometry-Lite Probe: A compact, interpretable probe that converts layer-wise hidden states into signed margin profiles under three geometric readouts and summarizes them along named axes (boundary position, change, and shape).
2. Decomposition of Multi-Layer Evidence: The framework disentangles three explanations of multi-layer signals: aggregation across informative layers, boundary-relative margin position, and explicit layer-to-layer motion.
3. Empirical Findings:
   • Signal Concentration: Safety evidence is primarily expressed through persistent boundary-position geometry (margin level and occupancy). Layer-to-layer motion (drift) contributes minimally to aggregate AUROC.
   • Drift Utility: While drift adds little to pooled AUROC, it provides sparse, recall-oriented corrections under shifted low-FPR thresholds.
   • Readout Stability Trade-off: Optimized linear boundaries are sharper in-distribution, whereas class-conditional mean geometry (centroid/DIM) retains separation more reliably under benchmark shifts (LOBO protocol).

Experimental Results

The evaluation spans nine instruction-tuned backbones (Llama, Gemma, Qwen families; 1.2B–70B) and seven safety benchmarks (XSTest, WildJailbreak, JBB-Behaviors, DoNotAnswer, BeaverTails, PKU-SafeRLHF, ToxicChat).

• In-Distribution Performance: Geometry-Lite (AUROC 0.955) and MultiLayer-Linear (0.954) outperform the best single-layer probe (0.941) and significantly outperform MultiLayer-DIM (0.911) and TaT-Disp-LSTM (0.945). Geometry-Lite matches raw score stacking while offering interpretability.
• Decision-Critical Regimes: At low false-positive rates (3% and 5% FPR), the multi-layer advantage is more pronounced. Geometry-Lite improves True Positive Rate (TPR) by ~0.06–0.08 over the best single-layer probe at 5% FPR.
• Generalization (LOBO): Under Leave-One-Benchmark-Out evaluation, the performance gap between supervised linear methods and class-mean methods (MultiLayer-DIM) narrows. MultiLayer-DIM (0.889) matches MultiLayer-Linear (0.889) and exceeds the best single-layer probe (0.842) in many cells.
• Feature Ablation:
  • Magnitude features alone match the full 39-feature model within 0.001–0.002 AUROC.
  • Structural features and Neg-drift features alone are significantly weaker.
  • Readout Ablation: In-distribution, the supervised linear boundary is strongest. Under hard benchmark shifts, the centroid readout becomes the most stable single branch, outperforming linear boundaries on the hardest held-out subsets.

Significance and Claims

The paper claims that prompt-level safety evidence is not primarily a layer-to-layer motion signal, but a persistent layer-wise margin geometry. The useful components of this signal are the position of the margin relative to safety boundaries and its persistence across layers.

The study highlights a critical trade-off in safety probing:

• Optimized boundaries (linear classifiers) capture strong separation in-distribution but degrade under distribution shift.
• Class-mean geometry (centroids/DIM) is less sharp in-distribution but more robust to benchmark shifts.

Consequently, the authors argue that multi-layer safety probes should be evaluated not only by in-distribution AUROC but also by which geometric biases transfer under shift. Geometry-Lite serves as a tool to make these components and biases visible, particularly in decision-critical regimes where false positives are costly.

Limitations: The study is limited to hand-designed margin summaries, single-turn safety prompts, and binary safety labels. Future work could explore learned profile functionals, multi-turn monitoring, and applications to other prompt-level tasks.