[Paper Review] DocBERT: BERT for Document Classification

Ashutosh Adhikari, Achyudh Ram+2 more|arXiv (Cornell University)|Apr 17, 2019

Text and Document Classification TechnologiesComputer Science23 references215 citations

TL;DR

Fine-tuning BERT achieves state-of-the-art results on four document classification datasets; a distilled KD-LSTM reg can match BERT base with ~30x fewer parameters and ~40x faster inference.

ABSTRACT

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.

Motivation & Objective

• Demonstrate that fine-tuned BERT achieves state-of-the-art results on standard document classification datasets.
• Investigate the feasibility of using BERT for long, multi-label documents common in 1- to 4-label settings.
• address computational costs of BERT by distilling knowledge into a smaller model (KD-LSTM reg) for faster inference.

Proposed method

• Fine-tune BERT base and BERT large for document classification by adding a final classification layer over the [CLS] token.
• Use cross-entropy (single-label) or binary cross-entropy (multi-label) losses for optimization.
• Apply knowledge distillation from fine-tuned BERT large to a lightweight single-layer BiLSTM (LSTM reg) using KL-divergence on a transfer set.
• Combine classification loss with distillation loss via a weighted sum to train the student (KD-LSTM reg).
• Create a transfer set with POS-guided word swapping and random masking to improve distillation.
• Evaluate on Reuters, AAPD, IMDB, and Yelp 2014 with standard splits and reported baselines.

Experimental results

Research questions

• RQ1 Can fine-tuned BERT achieve new state-of-the-art results on standard document classification datasets?
• RQ2 Is a lightweight BiLSTM able to approximate BERT base performance via knowledge distillation?
• RQ3 What are the trade-offs in accuracy versus model size and inference time for BERT versus the distilled student?
• RQ4 How do different datasets (single-label vs multi-label) affect training dynamics and performance of BERT fine-tuning?

Key findings

• BERT large achieves state-of-the-art results on all four datasets.
• BERT base also achieves strong results, closely following BERT large.
• KD-LSTM reg attains parity with BERT base on Reuters, AAPD, and IMDB, and offers substantial speedups (≥40x faster inference).
• KD-LSTM reg reaches ~1–3% of BERT base parameter count while maintaining competitive accuracy across datasets.
• Inference latency is dramatically reduced for the distilled model compared to BERT base (about 40x on the tested hardware).
• The distilled model demonstrates that a simpler architecture can recover much of BERT’s performance with far fewer parameters.