Qwen3.5-4B Activation Safety Classifier

Activation probe for prompt safety classification.

This is not a normal text classifier. It runs Qwen3.5-4B, extracts hidden states, then classifies the activations.

Labels:

safe   = 0
unsafe = 1

Checkpoint

checkpoint/qwen35_4b_activation_safety_probe_mlp.pt

This checkpoint is for Qwen3.5-4B activations.

Do not use it with another model unless you retrain the probe on that model's activations.

Training data

Dataset:

nvidia/Nemotron-3.5-Content-Safety-Dataset

Columns used:

text_col  = prompt
label_col = input_label

Activation setup

Expected setup:

chat template: enabled
thinking: disabled
pooling: mean + last token
classifier: MLP

Layer choices and extraction metadata are stored inside the .pt checkpoint.

Install

pip install torch transformers datasets accelerate

Predict

git clone https://github.com/Banaxi-Tech/activation-safety

python training_code/activation_safety_probe.py predict \
  --probe checkpoint/qwen35_4b_activation_safety_probe_mlp.pt \
  --model_path /path/to/Qwen3.5-4B \
  --text "How do I bake a cake?"

Output format:

{
  "label": "safe",
  "prob_unsafe": 0.01,
  "threshold": 0.5
}

Train

Extract activations:

python training_code/activation_safety_probe.py extract \
  --model_path /path/to/Qwen3.5-4B \
  --dataset nvidia/Nemotron-3.5-Content-Safety-Dataset \
  --split train \
  --text_col prompt \
  --label_col input_label \
  --out features/qwen35_4b_safety_activations.pt \
  --layers auto \
  --pool mean,last \
  --max_length 512 \
  --batch_size 8 \
  --dtype bf16

Train MLP probe:

python training_code/activation_safety_probe.py train \
  --features features/qwen35_4b_safety_activations.pt \
  --out checkpoint/qwen35_4b_activation_safety_probe_mlp.pt \
  --probe_type mlp \
  --hidden_dim 2048 \
  --epochs 10 \
  --batch_size 256 \
  --lr 1e-4

Other backbones

The code can train probes for other models.

The checkpoints are not universal.

For another model:

python training_code/activation_safety_probe.py extract \
  --model_path /path/to/OtherModel \
  --out features/other_model_activations.pt \
  --layers auto \
  --pool mean,last

python training_code/activation_safety_probe.py train \
  --features features/other_model_activations.pt \
  --out checkpoint/other_model_probe.pt \
  --probe_type mlp

Use each probe only with the same backbone and extraction settings it was trained with.