Source-linked AI summary
DocBERT: BERT for Document Classification
Ashutosh Adhikari, Achyudh Ram, Raphael Tang, Jimmy Lin
TL;DR
BERT had not yet been fine-tuned for document classification, a task involving less syntactic structure and often multiple labels. The paper fine-tunes BERT and distills its knowledge into a lightweight BiLSTM, achieving competitive results with substantially lower computational cost.
Problem
BERT had not yet been fine-tuned for document classification, where syntactic structure is less important and documents often have multiple labels.
Method
The paper fine-tunes BERT for document classification and applies knowledge distillation from BERTlarge to a smaller BiLSTM model.
Results
The distilled LSTMreg achieves BERTbase parity on a majority of datasets with over 30× parameter compression and at least 40× faster inference.
Takeaways & Limitations
Fine-tuned BERT establishes stronger document-classification baselines, while distillation provides competitive accuracy at substantially lower computational cost.
Takeaways & Limitations
The paper leaves distillation effects across a wider range of neural-network architectures and transformer-specific compression techniques for future exploration.
Abstract
from arXiv · showhide
We present, to our knowledge, the first application of BERT to document classification. A few characteristics of the task might lead one to think that BERT is not the most appropriate model: syntactic structures matter less for content categories, documents can often be longer than typical BERT input, and documents often have multiple labels. Nevertheless, we show that a straightforward classification model using BERT is able to achieve the state of the art across four popular datasets. To address the computational expense associated with BERT inference, we distill knowledge from BERT-large to small bidirectional LSTMs, reaching BERT-base parity on multiple datasets using 30x fewer parameters. The primary contribution of our paper is improved baselines that can provide the foundation for future work.
1 Introduction
The paper investigates BERT for document classification despite task characteristics that may make it less suitable, then addresses its computational cost through knowledge distillation.
- Motivation: Document classification differs from typical BERT tasks because syntactic structure may matter less and documents often have multiple labels.The authors also note that strong logistic regression and support vector machine baselines support the reduced importance of syntax claim.
- Motivation: BERT had not previously been fine-tuned for document classification, motivating evaluation on this task.
- Contributions: Fine-tuning BERT establishes state-of-the-art results on four popular document-classification datasets.
- Contributions: BERT’s quality improvement carries heavy computational expense because it has hundreds of millions of parameters and the prior baseline infers forty times faster.
- Contributions: Knowledge distillation transfers information from BERT-large to a much smaller BiLSTM, achieving results comparable to BERT-base with 30× fewer parameters.
- Contributions: The paper contributes improved document-classification baselines and a simpler neural model with competitive accuracy at lower computational cost.
2 Background and Related Work
The related work contrasts neural and pre-trained language-model approaches for document classification and motivates knowledge distillation as a model-agnostic compression technique.
- Document classification models: Neural architectures such as XMLCNN and HAN have been developed to address document classification, including its multi-label nature and long-document features.
- Pre-trained models: Pre-trained models including BERT and XLNet achieve state-of-the-art results across several NLP tasks but require substantial computational resources.
- Knowledge distillation: Knowledge distillation transfers information from a larger teacher to a smaller student using the teacher’s soft class-probability targets.
- Knowledge distillation: Because knowledge distillation is model agnostic, it can transfer knowledge between different types of architectures.
3 Our Approach
The approach fine-tunes BERT with a document-classification head, then distills the fine-tuned BERT-large model into a smaller LSTMreg using labeled and soft-target objectives.
- BERT fine-tuning: BERTbase and BERTlarge use a fully connected classifier over the final [CLS] hidden state and are optimized end-to-end.The classifier has parameters W ∈IRK×H, where H is the hidden-state dimension and K is the number of classes.
- BERT fine-tuning: The models use cross-entropy for single-label datasets and binary cross-entropy for multi-label datasets.
- Knowledge distillation: Knowledge is distilled from fine-tuned BERTlarge into the smaller LSTMreg using training examples and minor augmentations as the transfer set.
- Knowledge distillation: The distillation objective combines target-label classification loss with soft-target distillation loss.
- Knowledge distillation: The distillation loss minimizes KL divergence between student and teacher class probabilities.The paper denotes the student probabilities by p and the teacher probabilities by q.
4 Experimental Setup
Experiments compare fine-tuned BERT and distilled LSTMreg with established neural and linear baselines across four datasets, using dataset-specific training settings.
- Models and baselines: The evaluation compares BERT models with HAN, KimCNN, XMLCNN, SGM, LSTMreg, logistic regression, and support vector machine baselines.The linear baselines use tf–idf document vectors.
- Datasets: The four datasets are Reuters-21578, AAPD, IMDB reviews, and Yelp 2014 reviews, spanning multi-label and single-label classification.
- Datasets: Dataset splits use standard or previously provided partitions for Reuters and AAPD, while IMDB and Yelp use 80% training and 10% each for validation and test.
- Training and hyperparameters: BERT training optimizes epochs, batch size, learning rate, and maximum sequence length, with a 512-token maximum sequence length selected across datasets.The authors report dataset-specific epoch counts because model quality is sensitive to the number of epochs.
- Training and hyperparameters: Distillation trains LSTMreg with batch sizes of 128 for multi-label tasks and 64 for single-label tasks.
- Training and hyperparameters: Transfer sets are formed by augmenting training data with POS-guided word swapping and random masking, except Yelp2014 uses no augmentation due to computational restrictions.
5 Results and Discussion
BERT-large achieves the strongest results across all four datasets, while distilled LSTMreg matches BERT-base on most datasets with substantially lower computational requirements.
- Table 2 reports validation and test results, with reimplemented values averaged over five runs and original values reported as point estimates.
- BERT-large achieves state-of-the-art results on all four datasets, followed by BERT-base.
- KD-LSTMreg achieves average parity with BERT-base on Reuters, AAPD, and IMDB, but narrows the gap less on Yelp.
6 Conclusion and Future Work
The paper improves document-classification baselines by fine-tuning BERT and distilling its knowledge into a lightweight BiLSTM, achieving competitive quality with lower model cost. Future work includes studying distillation across architectures and transformer-specific compression.
- Conclusion: Fine-tuned BERT improves document-classification baselines, while knowledge distillation improves a lightweight single-layer BiLSTM model.
- Conclusion: Over 30× compression and at least 40× faster inference accompany BERT-base parity on a majority of datasets for distilled LSTMreg.
- Future Work: Future work should study distillation across neural-network architectures and model-compression techniques designed for transformer models.
A Appendix
On Reuters, using four epochs produces an F1 score worse than logistic regression.
- On Reuters, four training epochs yield an F1 score worse than logistic regression.
A.1 Hyperparameter Analysis
Figure 2 examines how maximum sequence length and fine-tuning epochs affect validation performance across document-classification datasets. Truncation is especially damaging on IMDB, while smaller datasets require more epochs to converge.
- MSL analysis: Lowering maximum sequence length causes only a minor F1 loss on Reuters, possibly because its documents are shorter.
- MSL analysis: Lowering maximum sequence length causes a drastic accuracy decline on IMDB, suggesting the entire document is necessary.
- MSL analysis: 256-token truncation lowers accuracy below the much smaller LSTMreg baseline, showing that truncation can be highly detrimental.
- MSL analysis: Any amount of truncation is detrimental in document classification, although the degradation level differs across datasets.
- Epoch analysis: Four epochs produce a significant drop in BERT model quality on AAPD and Reuters, indicating that smaller datasets require more epochs to converge.