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

2026-05-22 · Source: cs.LG updates on arXiv.org · Field: Technology & Digital — Artificial Intelligence & Machine Learning, Cybersecurity & Data Privacy · Depth: Expert, extended

Summary

Geometry-Lite is a novel, compact prompt-level probe designed for interpretable safety analysis in large language models, addressing how safety evidence manifests across layers and its robustness to benchmark shifts. It operates by mapping each layer's final prompt-token representation to signed margins using three geometric readouts: centroid distance, k-nearest-neighbor local structure, and a supervised linear boundary. These margin profiles are then summarized by boundary position, layer-to-layer change, and coarse shape. Evaluated across nine instruction-tuned backbones (1.2B–70B) and seven safety benchmarks, Geometry-Lite achieved a pooled AUROC of 0.955, outperforming single-layer probes by 0.013–0.014 and matching raw multi-layer score stacking. The study reveals that safety evidence primarily resides in persistent boundary-position geometry, with layer-to-layer motion offering only sparse, recall-oriented corrections at low false-positive rates. Crucially, while optimized linear boundaries are sharp in-distribution, class-conditional mean geometry proves more stable under benchmark shift.

Key takeaway

For AI Scientists and ML Engineers building LLM safety guardrails, you should adopt multi-layer safety probes like Geometry-Lite to significantly improve detection performance, especially in critical low-false-positive regimes and under benchmark shifts. When designing or selecting probes, prioritize class-conditional mean geometry readouts for enhanced stability and generalization to unseen safety distributions, even if optimized linear boundaries appear sharper in-distribution. Always evaluate your probes across diverse benchmarks and low-FPR thresholds to ensure real-world robustness.

Key insights

Geometry-Lite decomposes multi-layer safety signals into interpretable geometric margin profiles.

Principles

• Safety evidence is primarily persistent boundary-position geometry.
• Optimized linear boundaries are sharper in-distribution but less stable under shift.
• Class-mean geometry offers greater stability under benchmark shift.

Method

Reduce D-dimensional hidden states to scalar margins using centroid, k-NN, and linear boundary readouts. Summarize these profiles by level, change, and shape. Classify with L2-regularized logistic regression.

In practice

• Evaluate multi-layer probes using both in-distribution AUROC and shift stability.
• Prioritize class-mean readouts for robustness against benchmark shifts.
• Focus on boundary-relative margin position for aggregate safety detection.

Topics

• Large Language Models
• LLM Safety
• Hidden State Probes
• Interpretability
• Margin Geometry
• Benchmark Shift

Best for: Research Scientist, AI Scientist, Machine Learning Engineer, AI Security Engineer

Related on AIssential

• Guardrails for LLMs: Measuring AI ‘Hallucination’ and Verbosity — Controlling LLM verbosity via readability metrics can reduce hallucination risks and improve response clarity.
• Harmful Intent as a Geometrically Recoverable Feature of LLM Residual Streams — Harmful intent is geometrically recoverable in LLM residual streams, independent of alignment or model scale.
• Trajectory Geometry of Transformer Representations Across Layers — Trajectory geometry offers a principled, probe-free lens to understand Transformer representation evolution across layers.
• Linear Probes Detect Task Format, Not Reasoning Mode in Language Model Hidden States — Linear probes in LLMs often detect task format confounds, not distinct reasoning mode representations, challenging common interpretability claims.
• Risk-Aware LLM Agents for Geospatial Data Retrieval: Design and Preliminary Adversarial Evaluation — The framework uses LLM agents to convert natural language into structured API calls for geospatial data retrieval, integrating safety guardrails.
• LLMs Count with Geometry: Claude’s 6D Helix for Line Breaks — LLMs count characters and manage line breaks using geometric representations in high-dimensional spaces.

Open in AIssential →

Editorial summary, takeaway, and curation by AIssential. Original article published by cs.LG updates on arXiv.org.