Downloads · 30 days
0
llm-guardrails-probing
katanastas/llm-guardrails-probing
llm-guardrails-probing is a machine learning model from katanastas. Use it for the machine learning task on the model card, and read the license before you ship it in a product.
Logistic regression probes trained on hidden states of LLMs to detect unsafe content. Part of a research project investigating whether internal representations of LLMs can serve as lightweight guardrails.
Access
Public
Updated May 19, 2026
Repo size
46.6 MB
Likes
0
Public
Repo makeup
Click a slice to open those files.
.parquet54.2 MB · 90%
At a glance
- Access
- Public
- Created
- May 18, 2026
- Updated
- May 19, 2026
- SHA
- 97283328
Base models
Model card
From the Hugging Face model README
LLM Guardrails Probing — Artifacts
Logistic regression probes trained on hidden states of LLMs to detect unsafe content. Part of a research project investigating whether internal representations of LLMs can serve as lightweight guardrails.
Models covered
| Slug | Base model |
|---|---|
gemma-base-2B | google/gemma-2-2b |
gemma-3-4B | google/gemma-3-4b-pt (instruct) |
gemma-3-4B-base | google/gemma-3-4b-pt (base) |
qwen3-4B-base / qwen3-base-4B | Qwen/Qwen3-4B-Base |
qwen3guard-gen-4B | Qwen/Qwen3-4B (guard fine-tune) |
nemotron-4B | nvidia/Nemotron-Mini-4B-Instruct |
nemotron-4B-think | Nemotron with thinking |
yufeng-xguard-8B | YuFeng xGuard 8B |
Probe types
- Binary probe (
probes_binary/) — safe vs unsafe (logistic regression per layer, 5-fold CV) - Quad probe (
probes_quad/) — TP / TN / FP / FN relative to the guard model's own predictions
Datasets
- aegis-train / aegis-test — nvidia/Aegis-AI-Content-Safety-Dataset-2.0
- aegis-test-ru — katanastas/aegis-safety-ru
- mixed / mixed-ru — balanced ~1K dataset; RU version: katanastas/mixed-safety-ru
Binary probe results — AEGIS test set
Probes trained on AEGIS-train (5-fold CV), evaluated on held-out AEGIS-test. Best layer selected by CV PR-AUC.
Base models (no safety fine-tuning)
| Model | Best layer | CV PR-AUC | PR-AUC (EN) | ROC-AUC (EN) | F1 (EN) | F1-macro (EN) | PR-AUC (RU) | F1-macro (RU) |
|---|---|---|---|---|---|---|---|---|
| gemma-base-2B | 14 | 0.950 | 0.932 | 0.893 | 0.850 | 0.806 | 0.886 | 0.460 |
| gemma-3-4B (instruct) | 16 | 0.946 | 0.928 | 0.882 | 0.842 | 0.794 | 0.898 | 0.702 |
| gemma-3-4B-base | 34 | 0.937 | 0.935 | 0.893 | 0.855 | 0.810 | 0.715 | 0.398 |
| qwen3-4B-base | 36 | 0.938 | 0.920 | 0.877 | 0.845 | 0.797 | 0.746 | 0.403 |
| qwen3-base-4B | 26 | 0.945 | 0.930 | 0.888 | 0.844 | 0.796 | 0.851 | 0.441 |
F1-macro (RU) is low for base models because probes trained on English hidden states do not generalise well to Russian inputs — the instruct-tuned gemma-3-4B is the notable exception (0.702).
Results — guard models vs probes
Guard model vs probe — F1 by dataset
Confusion matrices — guard vs probe per model
Per-class F1 — guard vs probe
Repo structure
{model}/{experiment}/{probe_dir}/
models/binary.pkl # trained LogisticRegression (best layer)
models/quad.pkl # quad probe (guard models only)
probe_results.csv # per-layer CV metrics
summary.txt # best layer summary
weight_analysis.pkl # feature importance
{model}/{experiment}/{embeddings_dir}/
meta.parquet # text + labels for each sample
silhouette_scores.npy # silhouette score per layer (TP/TN/FP/FN, PCA)
- Created
- May 18, 2026
- Updated
- May 19, 2026