You can deploy this model and start earning money today!

katanastas/
llm-guardrails-probing

Name: llm-guardrails-probing
Availability: InStock

Views

No views yet

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)

katanastas/llm-guardrails-probing