Building a PII Classifier in 2 Hours for $3: A Teacher-Student Distillation Story

How I distilled Roblox’s 0.6B parameter PII classifier into a lightweight Llama 3.2 3B model using synthetic data and Fireworks AI

Table of Contents

The Problem
The Solution: Distill the Knowledge
Architecture Deep Dive
The Key Insight: Tokens as Labels
The Math: Why This Works
Applying This to Our PII Classifier
Why Not Use a Separate Classification Head?
Lessons Learned
Results
Try It Yourself
What’s Next?
Acknowledgments

The Problem

Moderating online platforms for PII (Personally Identifiable Information) is critical for user safety, especially in gaming environments where children interact. Roblox recently open-sourced their PII classifier, a powerful XLM-RoBERTa-Large model trained on billions of chat messages.

But there’s a catch: The Roblox model is 0.6B parameters and requires GPU inference. What if you want something:

Faster for real-time moderation?
Cheaper to deploy at scale?
Available as a simple API with confidence scores?

Enter teacher-student distillation and the “tokens as labels” approach.

The Solution: Distill the Knowledge

I built a pipeline that:

Generates synthetic chat messages using DeepSeek’s API (~$0.50)
Labels them with the Roblox teacher running on RunPod (~$1.50)
Fine-tunes a small student model on Fireworks AI (~$1.00)

Total cost: ~$3. Total time: ~2.5 hours (mostly waiting for data generation).

The result? 87.4% accuracy with a model that fits in a 46MB LoRA adapter.

Architecture Deep Dive

┌──────────────────────────────────────────────────────────────────────────────┐
│                     THE TEACHER-STUDENT DISTILLATION PIPELINE                 │
└──────────────────────────────────────────────────────────────────────────────┘

     ┌─────────────┐                                                           
     │  DeepSeek   │  "Generate gaming chat messages that ask for personal    
     │    API      │   information like phone numbers, social media, etc."    
     └──────┬──────┘                                                           
            │                                                                  
            ▼                                                                  
   ┌────────────────────────────────────────────────────────────────────────┐  
   │                        5,000 SYNTHETIC MESSAGES                         │  
   │                                                                         │  
   │  "none":    "gg wp", "nice shot!", "rematch?"                          │  
   │  "asking":  "whats ur snap?", "drop ur discord", "u got insta?"        │  
   │  "giving":  "my snap is jake2024", "add me on dc: gamer#1234"          │  
   └────────────────────────────────────────────────────────────────────────┘  
            │                                                                  
            ▼                                                                  
     ┌─────────────────────────────────────────────────────────────────────┐   
     │                        ROBLOX TEACHER MODEL                          │   
     │                   (XLM-RoBERTa-Large, 0.6B params)                   │   
     │                                                                      │   
     │  Input: "whats ur instagram?"                                        │   
     │  Output: {                                                           │   
     │    "asking_sharing_score": 0.82,                                     │   
     │    "giving_pii_score": 0.15,                                         │   
     │    "sharing_contact_score": 0.71,                                    │   
     │    ...                                                               │   
     │  }                                                                   │   
     │                                                                      │   
     │  → Apply thresholds → Label: "asking"                                │   
     └─────────────────────────────────────────────────────────────────────┘   
            │                                                                  
            ▼                                                                  
   ┌────────────────────────────────────────────────────────────────────────┐  
   │                          LABELED TRAINING DATA                          │  
   │                                                                         │  
   │  {"messages": [                                                         │  
   │    {"role": "system", "content": "Classify PII..."},                   │  
   │    {"role": "user", "content": "Message: "whats ur instagram?""},    │  
   │    {"role": "assistant", "content": "asking"}   ← SINGLE TOKEN LABEL   │  
   │  ]}                                                                     │  
   └────────────────────────────────────────────────────────────────────────┘  
            │                                                                  
            ▼                                                                  
     ┌─────────────────────────────────────────────────────────────────────┐   
     │                         FIREWORKS SFT                                │   
     │                                                                      │   
     │  Base Model: Llama 3.2 3B Instruct                                   │   
     │  Method: LoRA (rank=16, alpha=32)                                    │   
     │  Epochs: 3                                                           │   
     │  Training Time: ~5 minutes                                           │   
     │  Cost: ~$1.00                                                        │   
     └─────────────────────────────────────────────────────────────────────┘   
            │                                                                  
            ▼                                                                  
     ┌─────────────────────────────────────────────────────────────────────┐   
     │                         STUDENT MODEL                                │   
     │                                                                      │   
     │  LoRA Adapter: 46MB                                                  │   
     │  Accuracy: 87.4%                                                     │   
     │  Inference: Single API call → "asking" (with logprobs)               │   
     │                                                                      │   
     │  → Deployed to HuggingFace: sugiv/sugiv-pii-classifier               │   
     └─────────────────────────────────────────────────────────────────────┘

The Key Insight: Tokens as Labels

The magic comes from Fireworks’ blog post on turning LLMs into classifiers.

Traditional classification: Train a separate classification head, get softmax probabilities.

Tokens as labels: Train the LLM to output a single token ("none", "asking", or "giving"). Then use logprobs to get calibrated probabilities!

response = client.chat.completions.create(
    model="my-classifier",
    messages=[...],
    max_tokens=5,
    logprobs=True,       # ← This is the magic
    top_logprobs=5
)

# Extract probability from log-probability
import math
token = response.choices[0].logprobs.content[0]
confidence = math.exp(token.logprob)  # e.g., 0.95 for "asking"

This approach is:

Simpler: No custom architecture needed
Calibrated: Logprobs naturally give you confidence scores
Flexible: Works with any LLM that supports logprobs

The Math: Why This Works (And Why No Renormalization Is Needed)

This section explains the mathematical foundation behind using LLM token probabilities as classifier outputs, based on Fireworks’ excellent appendix.

PMF vs PDF: Discrete vs Continuous

Key insight: A PMF (Probability Mass Function) is for discrete outcomes; a PDF (Probability Density Function) is for continuous ones. Token probabilities are a PMF over a discrete vocabulary, and cross-entropy + gradient descent reshapes that PMF so almost all mass lies on your class tokens.

Think of three stages:

PMF vs PDF - understanding the difference
PMF over all tokens - before fine-tuning
PMF concentrated on class tokens - after training

Stage 1: PMF vs PDF (Discrete vs Continuous)

PMF (discrete) – probability bars at separate points:

Probability
   ^
   |            *
   |      *     |     *
   |      |     |     |
   +----------------------------->  x
        1     2     3       (discrete values)
        p(1)  p(2)  p(3)

(sum of bar heights = 1)

PDF (continuous) – smooth curve over an interval:

Probability density
   ^
   |        /   |       /     |      /       |_____/____________________>  y
       a        b        (continuous range)

P(a ≤ Y ≤ b) = area under the curve between a and b

PMF: probabilities at isolated points, sum to 1
PDF: density over a continuum, integral (area) = 1

LLM token probabilities are a PMF: For input x, the model outputs probabilities p_t = P(next token = t | x) for each token t, and Σ p_t = 1 over all tokens in vocabulary V.

Stage 2: Before Fine-Tuning - PMF Over Full Vocab

Let the vocabulary be tokens V (~32K-128K tokens), and class tokens be a small subset C ⊂ V.

Before fine-tuning, the PMF spreads mass across many tokens:

Tokens:  [ ... non-class ... | CLASS_1 | CLASS_2 | CLASS_3 | ... non-class ... ]

Probability
   ^
   |     *       *         *      *        *
   |   * |     * |       * |    * |      * |
   +-----------------------------------------------> tokens
       t1   t2   CLASS_1  t3  CLASS_2  CLASS_3 ...
         ^          ^          ^
         |          |          |
       non-class  class      class

Total class mass = p(CLASS_1) + p(CLASS_2) + p(CLASS_3) < 1
(there is non-trivial mass on many non-class tokens)

This is still a valid PMF: all bars ≥ 0, sum to 1. But class tokens do not yet dominate.

Stage 3: After Cross-Entropy Training - Mass Concentrated on Class Tokens

Training objective: for classification prompts, make the correct class token very probable and push non-class tokens’ probabilities toward 0.

Over training:

Logits for correct class tokens increase
Logits for non-class tokens decrease
The PMF gets reshaped so almost all mass is on class tokens

After sufficient fine-tuning:

Tokens:  [ ... non-class ... | CLASS_1 | CLASS_2 | CLASS_3 | ... non-class ... ]

Probability
   ^
   |                              *
   |                            * |   *
   |                          *   | * |
   |                        *     | | |
   +-----------------------------------------------> tokens
       t1   t2   CLASS_1  t3  CLASS_2  CLASS_3 ...
       ~0   ~0    ~0.6    ~0   ~0.3     ~0.1

Non-class bars ≈ 0
Class bars sum ≈ 1.0

So for classification prompts:

Σ p_c ≈ 1 for c ∈ C (class tokens)
Σ p_j ≈ 0 for j ∉ C (non-class tokens)

This is exactly a PMF over the label set C!

Now you can read the LM head’s probabilities for class tokens directly as calibrated class probabilities:

PII_labels = ["none", "asking", "giving"]

p("none")   ≈ 0.05
p("asking") ≈ 0.30
p("giving") ≈ 0.65

# Sum ≈ 1.0, non-class tokens ≈ 0

No extra renormalization step or separate classifier head is needed: the cross-entropy + gradient descent process has already concentrated the PMF’s mass onto your discrete class tokens.

The Math: Cross-Entropy + Gradient Descent

Let’s define:

V = full vocabulary (discrete outcomes, ~32K-128K tokens)
C ⊂ V = class tokens (our labels: “none”, “asking”, “giving”)
For input x, model logits z ∈ ℝ^|V| and PMF via softmax:

p_k = exp(z_k) / Σ exp(z_r)  for all r ∈ V

This defines a probability distribution (PMF) over tokens.

The Cross-Entropy Loss

For one training example with correct label token c ∈ C, the cross-entropy loss is:

L = -log(p_c)

The gradient of this loss with respect to logits is:

∂L/∂z_c = p_c - 1    (negative when p_c < 1, pushes z_c UP)
∂L/∂z_j = p_j        (positive for j ≠ c, pushes z_j DOWN)

Gradient Descent Updates

With learning rate η > 0, gradient descent updates logits:

z_c^(t+1) = z_c^t + η(1 - p_c)    ← correct token gets pushed UP
z_j^(t+1) = z_j^t - η·p_j         ← all others get pushed DOWN

The Logit Gap Analysis

Define the logit gap between a non-class token j ∉ C and correct class token c:

Δ_j^t = z_j^t - z_c^t

Then after one update:

Δ_j^(t+1) = z_j^(t+1) - z_c^(t+1)
          = (z_j^t - η·p_j) - (z_c^t + η(1 - p_c))
          = Δ_j^t - η(1 + p_j - p_c)

Since probabilities sum to 1:

1 - p_c = Σ p_k for k ≠ c  ≥  p_j

So:

1 + p_j - p_c ≥ 1 + p_j - (1 - p_j) = 2·p_j

Therefore:

Δ_j^(t+1) ≤ Δ_j^t - 2η·p_j

Interpretation: If p_j > 0 (the model gives token j any probability), the gap Δ_j decreases by at least 2η·p_j each step.

From Logit Gap to Probability Ratio

The gap relates to probabilities via:

p_j / p_c = exp(z_j) / exp(z_c) = exp(z_j - z_c) = exp(Δ_j)

So:

p_j = p_c · exp(Δ_j) ≤ exp(Δ_j)

As training proceeds:

If p_j remains non-zero, updates keep subtracting at least 2η·p_j from Δ_j
This forces Δ_j → -∞
Then exp(Δ_j) → 0
Thus p_j → 0

The Key Conclusion

For any non-class token j ∉ C:

Cross-entropy + gradient descent pushes its probability p_j toward zero
Because the PMF must still sum to 1, the mass “removed” from non-class tokens flows onto the class tokens C
So Σ p_c → 1 for c ∈ C

This is why raw class-token probabilities are “effectively calibrated without renormalization” for classification prompts. The training dynamics naturally concentrate all probability mass on your label tokens!

Applying This to Our PII Classifier

Step 1: Define Discrete Labels as Tokens

We chose a discrete label set:

C = {"none", "asking", "giving"}

Each label is a single vocabulary token. Our classification prompt ends so the next token should be one of these:

You are a PII classifier. Reply with exactly one word:
"none", "asking", or "giving".

Message: "<chat text>"

At generation time, the model’s next-token PMF p_t = P(next token = t | x) gives us a discrete distribution over the vocabulary.

Step 2: Fine-tune on PII Labels with Cross-Entropy

We:

Take Llama 3.2 3B Instruct as the backbone
Add LoRA adapters to attention/MLP blocks (rank=16, alpha=32)
Freeze base weights
Train with standard cross-entropy on class tokens

For each training example:

Correct token is the label token c ∈ {"none", "asking", "giving"}
Loss L = -log(p_c)
Backprop + gradient descent update only LoRA parameters

The math above applies: logits for correct class tokens get pushed up relative to non-class tokens, driving non-class probabilities toward 0.

Step 3: Use LM-Head Probabilities as Calibrated Confidences

At inference:

response = client.chat.completions.create(
    model="my-pii-classifier",
    messages=[...],
    max_tokens=5,
    logprobs=True,      # Get probability distribution
    top_logprobs=5
)

# The logprob IS the calibrated confidence
import math
logprob = response.choices[0].logprobs.content[0].logprob
confidence = math.exp(logprob)  # e.g., 0.92 for "asking"

Because Σ p_c ≈ 1 for class tokens (non-class probabilities near 0), these p_c form a PMF over labels that is already normalized and empirically well-calibrated.

You can then:

Threshold: Flag high-risk when P("giving" | x) > 0.8
Rank: Choose label with max probability
Monitor: Track calibration and adjust thresholds

Why Not Use a Separate Classification Head?

The traditional approach:

Backbone (frozen) → Pooled representation → Linear head → Softmax → Labels

This works, but:

Requires new architecture: Define and wire up a classification head
Different training: Head-only fine-tuning vs full model
No logprobs API: Most inference APIs don’t expose intermediate representations

The LM-head-as-classifier approach:

Reuses existing infrastructure: Same model, same API
Native logprobs: Confidence scores come for free
Mathematically grounded: Cross-entropy naturally concentrates mass on class tokens

This is exactly what Fireworks designed their SFT pipeline around, and why we chose it for our PII classifier.

Lessons Learned

1. Not All LLMs Are Created Equal

I initially tried Qwen3-4B as the student model. It has an “extended thinking” mode that outputs <think>...</think> traces before answering:

<think>Let me analyze this message for PII content. The user is asking 
for a Snapchat handle, which is a social media platform where personal 
contact information is exchanged...</think>

asking

This resulted in 0% accuracy because my classification code was looking for a single-word response!

Solution: Switch to Llama 3.2 3B Instruct, which follows instructions precisely and outputs just the label.

2. Firectl > REST API

Fireworks’ REST API for dataset uploads gave cryptic errors:

invalid character '-' in numeric literal

The CLI tool firectl worked flawlessly:

firectl create dataset pii-train data/train.jsonl
firectl create sftj --base-model accounts/fireworks/models/llama-v3p2-3b-instruct ...

3. Teacher Thresholds Matter

The Roblox model outputs raw scores, not labels. I experimented with thresholds:

if asking_score >= 0.2:           # Sensitive to asking
    label = "asking"
elif giving_score >= 0.3:         # Slightly less sensitive
    label = "giving"
elif combined_max >= 0.2691:      # Catch edge cases
    label = "giving"
else:
    label = "none"

Getting these right is crucial for training data quality.

4. Synthetic Data Works!

DeepSeek generated realistic gaming chat messages. The key was providing good few-shot examples:

prompt = """Generate 100 gaming chat messages. Examples:
- "gg" (none - just gaming)
- "whats ur snap?" (asking - requesting social media)
- "my insta is @gamer123" (giving - sharing contact info)

Make them realistic for gaming contexts (Roblox, Discord, etc.)"""

Results

Accuracy: 87.4%

On 500 held-out test examples:

Class	Precision	Recall	F1-Score
none	0.94	0.93	0.93
asking	0.75	0.66	0.70
giving	0.74	0.88	0.80

Confusion Matrix

               Predicted
            none  asking  giving
Actual none  316     16      8     (93% correct)
    asking   15     57     15     (66% correct - hardest class)
    giving    6      3     64     (88% correct)

Cost Breakdown

Step	Cost	Time
DeepSeek (5K messages)	$0.50	15 min
RunPod A40 (labeling)	$1.50	2 hrs
Fireworks SFT	$1.00	5 min
Total	$3.00	~2.5 hrs

Try It Yourself

Option 1: Use the Pre-trained Model

from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

base_model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.2-3B-Instruct",
    torch_dtype="auto",
    device_map="auto"
)
model = PeftModel.from_pretrained(base_model, "sugiv/sugiv-pii-classifier")
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-3B-Instruct")

messages = [
    {"role": "system", "content": 'Classify if this chat message involves PII. Reply with one word: "none", "asking", or "giving".'},
    {"role": "user", "content": 'Message: "whats your instagram?"'}
]
inputs = tokenizer.apply_chat_template(messages, return_tensors="pt").to(model.device)
outputs = model.generate(inputs, max_new_tokens=5, temperature=0)
print(tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True))
# → "asking"

Option 2: Train Your Own

Check out the full pipeline on GitLab.

What’s Next?

More data: Scale to 50K+ examples for better minority class performance
Multilingual: The Roblox teacher supports 10+ languages
Calibration analysis: Measure reliability diagrams for production use
Smaller models: Try Llama 3.2 1B or phi-2 for even faster inference

Acknowledgments

Roblox for open-sourcing their PII classifier
Fireworks AI for the SFT infrastructure and the excellent blog post
Meta AI for Llama 3.2
DeepSeek for affordable API access

Model available at: huggingface.co/sugiv/sugiv-pii-classifier

Code available at: gitlab.com/sugix/sugiv-pii-classifier

PII Classifier Distillation Related Posts