katanastas/llm-guardrails-probing

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)