BERT-base fine-tuned on AG News (Topic Classification)
Model description
This model is a partial fine-tuning of bert-base-uncased for topic
classification on the AG News dataset (4 classes: World, Sports, Business,
Sci/Tech). Only the top 2 encoder layers and the classification head were
trained; the rest of the encoder (embeddings + lower layers) was kept frozen.
Training data
- Dataset: fancyzhx/ag_news
- Subsample used: 10,000 training examples, 2,000 test examples
- Max sequence length: 128 tokens
Training procedure
- Base model:
bert-base-uncased - Adaptation method: Partial fine-tuning (top 2 layers + head unfrozen)
- Trainable parameters: 14,769,412 / 109,485,316 (~13.5%)
- Two learning-rate groups: head = 1e-3, encoder top layers = 2e-5
- Epochs: 2, batch size: 16
- Optimizer: AdamW
Evaluation results
| Metric | Value |
|---|---|
| Accuracy | 0.9025 |
| F1 macro | 0.9035 |
Compared against a feature-based baseline (frozen BERT-base + logistic regression), which achieved accuracy 0.8735 / F1 macro 0.8742. The fine-tuned model was chosen as the delivered model based on this comparison.
Intended use
Topic classification of short English news headlines/snippets into one of four categories: World, Sports, Business, Sci/Tech.
Limitations
- Trained on a 10k-example subsample of AG News, not the full training set.
- English only; performance on other languages or domains is not evaluated.
- Business and Sci/Tech categories show the highest confusion rate due to semantic overlap between tech-business news.
References
- Devlin et al., 2018. BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding
- Tunstall, von Werra & Wolf. Natural Language Processing with Transformers (O'Reilly), chapters 1-3.
- Dataset: fancyzhx/ag_news