DocBERT: BERT for Document Classification

TL;DR

This paper demonstrates that fine-tuning BERT on document classification yields state-of-the-art results across four datasets, but at considerable computational cost. To address efficiency, the authors distill knowledge from BERT_large into a lightweight LSTM_reg, achieving parity with BERT_base on several datasets while reducing parameters by ~30× and speeding up inference by ~40×. The KD-LSTM_reg approach uses a joint objective combining standard classification loss with a distillation term, guided by transfer-set augmentation, and shows promising results as a practical, scalable alternative to full BERT inference. Overall, the work strengthens document-classification baselines with strong, efficient alternatives and lays groundwork for broader transformer compression research.

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.

Figures (2)

• Figure 1: Effectiveness of KD-LSTM$_{reg}$ vs. BERT$_{base}$ and BERT$_{large}$
• Figure 2: Results on the validation set from varying the MSL and the number of epochs.