bandarra.me

Would Fine-Tuning Have Beaten Jev?

In my last post, I compared Jev with two custom models built from labeled comments: a small neural network (MLP) and k-nearest neighbors (k-NN). Both custom models start from an embedding of each comment, and with cloud embeddings they beat Jev once they had about a hundred labels.

Another common way to build a text classifier is to fine-tune a pretrained language model like BERT, and I was curious how that would perform against Jev. It works differently from the custom models in the last post. In both of them, an embedding model turns each comment into a list of numbers once, and that list never changes. The MLP can only learn which of those numbers matter, not how to compute better ones. Fine-tuning adjusts every one of the language model's weights on the labeled data instead.

So I fine-tuned BERT on the same 8,000 comments and scored it on the same 2,000. The short version: on the common labels, it's the best model in the comparison, and nothing goes to the cloud. But it's slower, much bigger, and it misses almost every comment on the two rarest labels.

Fine-tuning vs. frozen embeddings. In the last post, the embedding model was a fixed first step: it read the comment and produced a list of numbers, and the MLP only ever saw those numbers. Fine-tuning removes that boundary. A small classification layer goes on top of the language model, and training updates the whole stack, so the model can learn to pay attention to the words and phrases that matter for toxicity, like tone, slurs or a "you" aimed at another person, instead of whatever the embedding model was originally trained to capture.

The setup

Everything that was in the last post stays the same: the first 10,000 comments from the Jigsaw Toxic Comment Classification Challenge, the same 8,000 for training and 2,000 for evaluation, the same six labels and the same metrics. The other approaches' numbers are unchanged.

The new model is bert-base-uncased, a 110-million-parameter model from 2018, with one output for each of the six labels. I trained it once with a standard recipe: 3 passes over the 8,000 comments, a learning rate of 2e-5, batches of 16, and each comment cut to its first 128 tokens. I didn't tune any of these on the evaluation set, or at all. Training took 9 minutes on the GPU of an Apple M1 Pro laptop.

The core of the training loop is short, using the Hugging Face transformers library:

model = AutoModelForSequenceClassification.from_pretrained(
    "bert-base-uncased",
    num_labels=6,
    problem_type="multi_label_classification",
)
optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5, weight_decay=0.01)
loss_fn = nn.BCEWithLogitsLoss()

for batch, labels in batches:
    loss = loss_fn(model(**batch).logits, labels)
    loss.backward()
    optimizer.step()
    scheduler.step()
    optimizer.zero_grad()

The difference from the last post is in the optimizer line. With frozen embeddings, the optimizer only gets the MLP's weights. Here it gets model.parameters(), all 110 million of them. The full script, with batching and evaluation, is src/eval_bert.py in the repository.

Ranking comments

Ranking quality on the toxic label for fine-tuned BERT, the MLP and k-NN on cloud embeddings, and Jev, measured with ROC-AUC and average precision Figure 1: How well each approach ranks toxic comments above clean ones on the toxic label. Each line is a 95% bootstrap interval. Fine-tuned BERT has the best score on both.

On toxic, BERT reaches 0.98 ROC-AUC and 0.88 average precision, the best of the four on both. The gap to the MLP on cloud embeddings, 0.88 against 0.85, is small but holds up: in a paired bootstrap, BERT comes out ahead in 99% of resamples. Against Jev, it's ahead in every resample.

ROC-AUC hides where the models differ, because all four curves crowd into the top-left corner. Zooming into that corner, where only a small share of clean comments get flagged, makes the differences visible:

ROC curve on the toxic label, zoomed to false positive rates up to 20%, for fine-tuned BERT, the MLP and k-NN on cloud embeddings, and Jev Figure 2: The top-left corner of the ROC curve on toxic. Each curve shows the share of toxic comments caught (up) against the share of clean comments wrongly flagged (right). Higher is better.

If you accept flagging 2% of clean comments, about 36 of the 1,815, BERT catches 82% of the toxic ones. The two cloud models catch about 76%, and Jev catches 60%. Jev only catches up with the MLP when about 10% of clean comments get flagged. That's why Jev's ROC-AUC (0.97) is higher than the MLP's (0.96) while its average precision is clearly lower: Jev does well in the part of the curve that matters least in practice. The k-NN curve moves in big steps because, with 7 neighbors, it can only give 8 different scores.

It isn't the best ranker I've built on this dataset, though. In the last post, k-NN on cloud embeddings with 51 neighbors reached 0.90 average precision. BERT and that k-NN are a statistical tie.

Making decisions

The scores are also well calibrated. BERT's default 0.5 threshold gives 0.81 F1 on toxic: it catches 146 of the 185 toxic comments and flags 30 clean ones. With a threshold picked fairly, the same way as in the last post, the best threshold lands at 0.44, right next to the default.

Approachtoxic F1 with a fair thresholdAverage F1, all six labels
BERT, fine-tuned0.800.49
k-NN, cloud embeddings0.800.57
MLP, cloud embeddings0.770.54
Jev0.700.51

On toxic, BERT ties k-NN on cloud embeddings for the best F1. Better ranking doesn't always become better decisions, which is what I saw with 51-neighbor k-NN too. Averaged over all six labels, in the right-hand column, BERT drops below Jev, because of the two rarest labels.

The rare labels

Average precision on all six labels for fine-tuned BERT, the MLP and k-NN on cloud embeddings, and Jev Figure 3: Average precision on each label, with the number of examples in the evaluation set in parentheses. BERT leads on the four larger labels and scores close to zero on the two smallest.

On the four labels with at least 80 training examples, BERT is first or tied for first on average precision:

Label (training examples)BERTMLP, cloudk-NN, cloudJev
toxic (786)0.880.850.830.79
obscene (438)0.870.870.810.85
insult (399)0.830.750.770.63
severe_toxic (81)0.540.390.230.32

BERT's biggest lead is on insult. In the last post, Jev trailed there because the raters used the label more narrowly than my question did. BERT learned that narrower definition better than any other model, though with a fair threshold, k-NN still edges it on F1, 0.73 against 0.71.

Then there are threat and identity_hate, with 27 and 71 training examples. BERT doesn't learn to find either of them. Across all 2,000 evaluation comments, its highest threat score is 0.23, so at any sensible threshold it never flags a threat. The six real threats score 0.09 or lower, and the median one ranks around 120th. Its average precision on identity_hate is 0.12, against 0.54 for the MLP on cloud embeddings.

The likely explanation is that all six labels share one model, trained together for three short passes. With 27 threats among 8,000 comments, they barely move the training loss, and the model spends its effort on the common labels. The cloud-embedding models don't have this problem to the same degree: the embedding already puts threats close together, so even a handful of examples is enough to find them. Weighting the rare labels more heavily, training longer, or training a separate model per label would probably help, but I haven't tested that.

How many labels does it need?

The last post's most useful chart showed how quality changes with the number of labels. I repeated it for BERT, training on 50 to 1,000 comments, three times per size, with the same random subsets the MLP and k-NN used. Each run got 1,500 training steps, the same as the full 8,000-comment run, so a small training set is seen many more times. That's the same rule the MLP sweep follows, and it matters. My first sweep gave the small runs only 200 steps, and 4 of its 15 models didn't catch a single toxic comment at the default 0.5 threshold.

Average precision on the toxic label as the number of labeled training examples grows, for fine-tuned BERT, the MLP on cloud embeddings, and Jev Figure 4: Average precision on toxic as the training set grows. Light dots are individual BERT runs, and the MLP line is the average over the same three subsets. The dotted segment joins the sweep to the single full-data run, with no runs in between.

Labeled commentsBERTMLP, cloudJev
500.630.830.79
1000.680.850.79
2500.770.830.79
5000.820.840.79
1,0000.850.850.79
8,0000.880.870.79

Average precision on toxic, averaged over the same three training subsets. The 8,000 row for BERT is the single full run. Jev uses no labels, so its score doesn't change. The MLP's 8,000 row comes from the sweep, which trains it a little differently from the single run in Figure 1, so it reads 0.87 here instead of 0.85.

The default threshold doesn't hold up on small training sets. At 50 labels, BERT's average F1 at 0.5 is 0.31, against 0.69 for the MLP on cloud embeddings. Small runs rank comments better than their F1 suggests, but their scores sit in the wrong place, so they'd need a threshold picked on held-out labels, the way Jev's threshold was picked in the last post.

Averaged over all six labels, BERT trained on 1,000 comments reaches 0.53 average precision. That's about as good as with all 8,000 and better than the MLP on cloud embeddings with the same 1,000. The 1,000-comment runs went through their data 24 times, against 3 times for the full run. That's a hint that the full run's failure on threat and identity_hate comes from too little training, not too little data, but I haven't tested it directly.

Speed, size and cost

ModelMedian time per commentSizeWhere the comment goes
BERT, laptop CPU52 ms438 MBNowhere
BERT, Apple M1 Pro GPU13 ms438 MBNowhere
Cloud embeddings + MLP254 ms2 MBGoogle
Jev, all six questions209 msnoneTypeSafe

I timed BERT on the same AMD Ryzen 7 5800H CPU as the local pipelines in the last post, one comment at a time with four threads, and on the GPU of an M1 Pro. Even on the CPU, it's four times faster than a round trip to either cloud service. It costs nothing per comment.

The costs are elsewhere:

What this doesn't tell you

Conclusion

The last post ended with "the first thing to get right is the embedding model". Fine-tuning adds a second answer: let the model learn its own representation. A fine-tuned BERT running on my own hardware ranked toxic comments better than anything built on cloud embeddings, except a k-NN it tied, and its default threshold worked without tuning. Nothing leaves the machine.

But it needs a GPU to train, it's a 438 MB model, and on the rarest labels it did worse than Jev, which never saw a single example. For the common labels, with thousands of examples and a GPU, fine-tuning is the strongest option I've tried. For my blog, with a hundred comments and one question, it doesn't change my choice: Jev is still the right trade.

What to try next

The code, cached predictions and analysis scripts are in the repository, including the sweep script, so you can try either one yourself. If you do, I'd love to hear how it goes!

Note: This post was created with the assistance of AI. While a human carefully reviewed and edited the content, it's important to remember that AI tools may introduce errors or biases. If you have any concerns or questions, please feel free to reach out.

Comments

No comments yet. Be the first to comment!

Leave a comment

Comments are moderated and may take some time to appear.

You may also want to read