One encoder, seven heads: what we learned training a unified security classifier with masked losses [P]

Setup and Architecture

We spent the last months consolidating seven separate sequence classifiers into one multi-head model - our apex model, so to speak. Since the weights are now public, I wanted to share what worked and what surprised us.

The setup: a shared mmBERT-small encoder with seven task heads:

• Binary injection (BCE)
• Document class (7‑way)
• Tool type (14‑way)
• Tool operation (6‑way)
• Tool data-flow tags (3× BCE, multi‑label)
• Intent routing (5‑way)
• Threat type (7‑way)

Training with Masked Losses

The part that needed care: our training rows only carry labels for a subset of tasks, so absent tasks are masked out of the loss entirely. We ended up writing a self‑test that asserts absent‑task gradients are exactly zero, which caught two subtle bugs, and I'd recommend it to anyone doing similar masking.

Data and Evaluation

About 5k synthetic/real multi‑task rows help the heads co‑train; the test sets stay 100% real data. Held‑out results per head:

• Injection F1: 0.962
• Documents: 0.980
• Tool type: 0.957
• Tool operation: 0.945
• Tool tags: 0.958
• Routing: 0.916
• Threat: 0.952

Quantization

Both the unified model and the dedicated single‑task variants ship quantized -edge builds (ONNX INT8 + INT4 embeddings, from 96 MB) with measured parity benchmarks in the repos. The worst head loses 0.012 against FP32.

Comparison: Unified vs. Dedicated

Was it worth it vs. seven dedicated models? We released both variants, so you can judge for yourself. The dedicated models score marginally higher on most tasks, but the unified one does one encoder pass instead of up to seven.

Weak Spot: Routing

Our weak spot: routing, at 0.916. The intent classes overlap semantically ("write code that analyzes my data" - is that code or analytics?), and I suspect the ambiguity is genuinely in the data. If you have ideas beyond relabeling, let me know :)

Weights and per‑head metrics: https://huggingface.co/patronus-studio