Trainable Weighted Pooling Method for Text Classification with BERT

Abstract

Text classification is one of the central challenges in natural language processing, encompassing techniques for cate-gorizing large amounts of text data into meaningful categories. This field plays an important role in many applications, such as information retrieval, sentiment analysis, and recommendation systems. In recent years, the remarkable development of deep learning technology has led to the proposal of large language models, which have achieved high performance in various tasks. BERT is one of the large language models widely recognized for its potential in text classification. Although BERT can effectively learn context-dependent word representations, an appropriate pooling strategy is necessary to obtain a representation of the entire document. In this study, we propose a pooling method called CLS-average pooling (CAP) that combines the commonly used the [CLS] embedding and the average pooling method in BERT for text classification. We obtain the sentence representations by taking the weighted sum of the embeddings obtained from the [CLS] embedding and the average pooling. At this time, we treat the weights used in CAP as trainable parameters to automatically acquire appropriate weights for text classification. We demonstrated that the proposed method is more effective than conventional pooling methods in text classification tasks by applying it to a dataset for text classification.