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.