---
tags:
- probing
- interpretability
- safety
- content-moderation
- logistic-regression
base_model:
- google/gemma-2-2b
- google/gemma-3-4b-pt
- Qwen/Qwen3-4B-Base
- Qwen/Qwen3-4B
- nvidia/Nemotron-Mini-4B-Instruct
---

# 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](https://huggingface.co/datasets/nvidia/Aegis-AI-Content-Safety-Dataset-2.0)
- **aegis-test-ru** — [katanastas/aegis-safety-ru](https://huggingface.co/datasets/katanastas/aegis-safety-ru)
- **mixed / mixed-ru** — balanced ~1K dataset; RU version: [katanastas/mixed-safety-ru](https://huggingface.co/datasets/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

![guard vs probe](images/16_guard_vs_probe.png)

### Confusion matrices — guard vs probe per model

![confusion matrices](images/20_guard_vs_probe_confusion.png)

### Per-class F1 — guard vs probe

![per-class F1](images/21_guard_vs_probe_perclass.png)

## 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)
```