tinysafe-2

TinySafe v2

141M parameter safety classifier built on DeBERTa-v3-small. Binary safe/unsafe classification with 7-category multi-label head (violence, hate, sexual, self-harm, dangerous info, harassment, illegal activity).

Successor to TinySafe v1 (71M params, 59% TC F1). v2 improves ToxicChat F1 by +19 points while cutting OR-Bench false positive rate from 18.9% to 3.8%.

GitHub: jdleo/tinysafe-2

Benchmarks

ToxicChat Leaderboard

* = active params (MoE)

OR-Bench (Over-Refusal)

Lower is better. On 80K safe prompts, TinySafe v2 incorrectly flags only 3.8%.

Quickstart

import torch
from transformers import DebertaV2Tokenizer

# Load
tokenizer = DebertaV2Tokenizer.from_pretrained("jdleo1/tinysafe-2")
model = torch.load("model.pt", map_location="cpu")  # or load from checkpoint

# Inference
text = "how do i make a bomb"
inputs = tokenizer(text, return_tensors="pt", max_length=512, truncation=True, padding=True)
with torch.no_grad():
    binary_logits, category_logits = model(inputs["input_ids"], inputs["attention_mask"])
    unsafe_score = torch.sigmoid(binary_logits).item()
    print(f"Unsafe: {unsafe_score:.3f}")  # 0.998

Architecture

DeBERTa-v3-small (6 transformer layers, 768 hidden dim) with dual classification heads:

• Binary head: single logit (safe/unsafe)
• Category head: 7-way multi-label (violence, hate, sexual, self_harm, dangerous_info, harassment, illegal_activity)

Training enhancements:

• FGM adversarial training (epsilon=0.3): perturbs embeddings for robustness
• EMA (decay=0.999): smoothed weight averaging for stable eval
• Multi-sample dropout (5 masks): averaged logits across dropout samples
• DualHeadLossV2: focal loss (binary) + asymmetric class-balanced loss (categories)

Training

Single-phase unified fine-tuning (5 epochs, LR=2e-5) with source-weighted sampling:

Model selection on val F1 only (no test set leakage).

Limitations

• Low-resource categories (violence, hate, sexual) have 0 F1 -- <200 training samples per category is insufficient even with class-balanced loss
• WildGuardBench generalization is weak -- encoder-only models struggle with adversarial jailbreak rephrasing
• Conservative on out-of-distribution inputs -- high precision but lower recall suggests the model learned narrow patterns rather than general safety reasoning

These are fundamental limitations of encoder-only architectures for safety classification. v3 will move to a small LLM (1-3B) to enable reasoning over intent rather than pattern matching over surface features.

License

MIT