Classifier Model Project
Introduction
The model is a DistilBERT encoder with 4-way head, fine-tuned to label a news blurb as World, Sports, Business, Sci/Tech. This is a classification model; it is not a chatbot or a generator.
Data
The training data used is the fancyzhx/ag_news dataset officially consisting of 120,000 training samples and 7,600 test samples. The training sample was cut into 108,000 sample training split and a 12,000 validation split (using seed 42 for deterministic splitting). The official test split was never seen during training or checkpoint evaluation.
Training
Training started from the distilbert/distilbert-base-uncased model and tokenizer. We tokenized each headline by truncating it to 128 tokens (or padding when necessary). We ran one epoch of the training split from the base model using a starting learning rate of 2e-5. Then we reloaded the model from this checkpoint and trained again for two more epochs of the training split with a fresh starting learning rate of 2e-5 as the LR had decayed to almost zero during the previous run. During this run, we saved the model with the highest accuracy on the validation split epoch. The training used MPS device backend for PyTorch as it was performed on Apple Silicon.
You can find the training Python code on Github here.
Accuracy
- Final Test split accuracy: 0.9441
- Validation split accuracy at selection: 0.9462
The minor gap of 0.0021 isn't significant enough to signal the model training was not successful.
Per-class Accuracy
The test split accuracy across the 4 classes is as follows:
- World 0.9511 (samples = 1900)
- Sports 0.9847 (samples = 1900)
- Business 0.9089 (samples = 1900)
- Sci/Tech 0.9316 (samples = 1900)
The model is quite accurate as classifying World and Sports headlines. The slight degradation in accuracy for Business and Sci/Tech is due to the model incorrectly classifying headlines which contain a mix of business and science/technology (like business articles about a technology company, or a technology article citing a business leader). For example:
Headline: Intel to delay product aimed for high-definition TVs SAN FRANCISCO -- In the latest of a series of product delays, Intel Corp. has postponed the launch of a video display chip it had previously planned to introduce by year end, putting off a showdown with Texas Instruments Inc. in the fast-growing market for high-definition television displays.
Truth: Business
Prediction: Sci/Tech
Reason: It seems like a lot of text does indeed have overlap between technology firms and technology centres like San Francisco, but the article is about business. The model is having a hard time with these distinctions when it deals with such companies.
- Downloads last month
- 21
Model tree for GVRV/ag_news
Base model
distilbert/distilbert-base-uncased