Source-linked AI summary
MIL-BERT: Classification of Arbitrarily Large Text with Performance and Explanatory Guarantees
John Cadigan, Dayne Freitag, Eric Yeh
TL;DR
Long-document classification is constrained by transformer length and memory limits, especially when useful evidence may occur anywhere in a document. MIL-BERT selects informative excerpts through a double-pass multiple-instance-learning procedure. It achieves state-of-the-art results on several tasks and competitive long-document performance while exposing reviewable decision-driving passages.
Problem
Long-text classification methods face fixed-length and quadratic transformer costs, while truncation can omit evidence needed for accurate prediction.
Method
MIL-BERT segments documents into excerpts, embeds them with RoBERTa, selects key excerpts in a gradient-free first pass, and trains on the selected subset in a second pass.
Results
MIL-BERT achieves state-of-the-art results for political-bias detection, trigger-warning detection, and several author-profiling tasks, with competitive long-document results and passage-level generalization.
Takeaways & Limitations
The approach classifies arbitrarily large text while identifying excerpts that can be reviewed as explanations and can support constituent-instance classification.
Takeaways & Limitations
Instance-labeled data for the article-bias and triggering-passage MIL experiments are limited, and whether their bags appear in training was not identified.
Abstract
from arXiv · showhide
Many text classification decisions are viable based on constituent excerpts alone. Taking inspiration from the field of multiple instance learning, we present an algorithm for training a neural network to classify text by selecting such excerpts. We show that our approach is also scalable with demonstrated learning against samples with nearly 1M tokens. We evaluate our methods on 7 datasets with emphasis on long-textual collections that far exceed the encoding limit of our base model. We present state-of-the-art results with this algorithm on 3 datasets: identification of political bias in news outlets, trigger warnings in long stories, and demographic characteristics of authors in tweet collections. Furthermore, the model trained on weakly-labeled collections of text (bags) generalizes to accurately classify constituent, smaller instances. Besides a new state-of-the-art for these problems, this approach is one of the few neural methods to excel in these datasets.
1 Introduction
MIL-BERT addresses long-text classification by searching for locally informative excerpts rather than processing entire documents. It frames such tasks as multiple instance learning and emphasizes competitive accuracy alongside reviewable explanations.
- Transformer classifiers face fixed-length limits and quadratic computation, making truncation a common compromise for long documents.
- Signal locality means that one or a few passages can support a particular prediction, enabling multiple instance learning over long texts.
- MIL-BERT uses fixed selectors to identify passages most relevant to a prediction before classification.
- Across long-text classification problems, MIL-BERT is competitive with whole-document alternatives and surpasses previously reported accuracies in several cases.
- The selected passages support explainability because they are key, easily reviewable excerpts that drive the classification decision.
2 Method
The method represents arbitrarily large documents as bags of excerpt embeddings and uses a double-pass selector to train with memory proportional to selected points rather than all points.
- MIL classifies a bag as positive when at least one constituent instance is positive, while training supplies only bag labels.
- Documents are segmented with a rolling window, and each excerpt is embedded using the start-token representation from a RoBERTa model.
- The double-pass pooling trick selects key excerpts without gradients, then computes gradients only on the selected subset.
- The method handles n points with k memory at the cost of O(n + k) compute.
- Selectors use affine transformations and extreme-value selection, with shared feature-map and feature-tree variants for combining selected excerpts into predictions.
- Gumbel-softmax makes individual or top-k excerpt selection differentiable, allowing more powerful classifier layers to operate over instances.
3 Datasets
The experiments cover classic MIL, implicit-label demographic profiling, and long-document classification across datasets containing long documents or text collections.
- The study evaluates MIL with identifiable instance labels, MIL with implicit labels, and long-text documents.
- Trigger-warning experiments train on story-level labels and evaluate both story bags and labeled passages containing fine-grained trigger warnings.
- Demographic profiling predicts gender, occupation, birthyear, and fame from celebrity authors’ tweet collections with noisy, implicit indicators.
- The long-document benchmark includes Hyperpartisan, 20Newsgroups, EURLEX-57K, and Book Text datasets with specified train, validation, and test partitions.
4 Hyperparameters
Hyperparameters vary excerpt windows, strides, selector types, and classifier configurations across datasets, with RoBERTa large used for embedding.
- Experiments tune excerpt window size, stride, selector choice, and the corresponding classification layer.
- The default target is a 256-token window with stride 1, using the shared classifier because it generally performs better and requires less tuning.
- Shared feature-map selectors can require more memory as the number of classes grows, motivating shorter excerpts for extreme classification.
- CLEF23 uses shared selection with top-k 3, while PAN19-BIRTHYEAR uses a feature-tree selector after shared selection underperformed.
- RoBERTa large is used to embed the excerpts.
- The experiments span datasets with longer documents and text collections, all carrying bag labels and some also carrying instance labels.
5 Results
MIL-BERT achieves strong classification performance while scaling to very long documents with controlled memory use. Its selected excerpts also support visual and qualitative analysis of how instances contribute to predictions.
- Classification performance: State-of-the-art results are achieved for news-source political bias and trigger-warning classification, with respectable passage-level trigger-warning accuracy.
- Classification performance: New state-of-the-art results are reported for gender and occupation in author profiling, but not for age or fame.
- Memory and compute: O(n + k) compute and O(k) memory characterize MIL-BERT, while training on a single 1M-token instance can approach 100 seconds.
- Memory and compute: Compared with naive max-pooling MIL-BERT, the double-pass approach uses less memory and can process far more tokens than long-context encoders such as Longformer-large and DeBERTa base.
- Chosen-instance analysis: In the CLEF23 visualization, selected excerpts lie on each bag’s convex hull and therefore contribute the largest feature-score magnitudes to the classification boundary.
- Explainability: Selector scores identify impactful instances and provide qualitative and quantitative insight into the decision function, including correlations with author birthyear of 0.42-0.62.
6 Discussion
MIL-BERT is presented as an efficient, explainable approach for long-text classification, with strongest gains on tasks exhibiting signal locality and limitations elsewhere.
- Results: MIL-BERT achieves new state-of-the-art results for trigger-warning detection, political-bias detection, and selected author-profiling tasks.The authors also report competitive results on the long-document benchmark, where no single method is dominant.
- Explainability: The method identifies and interprets selector scores for individual excerpts and in aggregate, supporting reviewable explanations.The study includes visualizations of selector behavior for age profiling and excerpt-level attribution for political-bias classification.
- Evaluation: The evaluation includes trigger warnings, political bias, demographic profiling, and long-document classification against prior neural and traditional approaches.The discussed datasets were difficult for conventional neural methods because of memory limits and truncation-related information loss.
- Scope: MIL-BERT is most suitable for multiple-instance-learning tasks and otherwise remains merely competitive on long-text classification.The authors associate weaker birthyear and fame performance with the absence of signal locality in those tasks.
- Future work: Future work targets improved selector and classifier layers, more efficient selection, and applications beyond text classification.The authors mention retrieval, matching, entity typing, and distant information extraction as possible directions, while noting risks in mining protected characteristics.
7 Related Work
MIL-BERT combines multiple-instance learning with selector-based explainability to process long texts through indicative excerpts rather than wholesale document encoding.
- MIL-BERT: MIL-BERT represents long documents as arbitrarily large bags of excerpts when classification exhibits signal locality.The approach is motivated by scaling and explainability requirements for long-document classification.
- Multiple instance learning: Multiple instance learning classifies a bag from constituent instances and can learn instance classifiers from bag-level labels.The canonical formulation treats a bag as positive when at least one constituent instance is positive.
- Multiple instance learning: Prior NLP work includes event detection, neural decision trees, and protein-interaction prediction, but some methods use all instances or sample limited contexts.These approaches differ in how they aggregate instances and manage memory.
- Explainability: Selector-predictor explainability methods include input perturbation, attention analysis, gradient analysis, REINFORCE, and Gumbel-softmax selection.These methods address explanation or differentiable discrete selection in different ways.
- Long-document classification: Long-document research has expanded transformer attention or identified key excerpts with systems such as Longformer, Infinite Memory, CogLTX, and BERT+TextRank.CogLTX uses separate judge and reasoner BERT models, while BERT+TextRank uses unsupervised content selection.
8 Conclusion
The paper presents MIL-BERT for classifying arbitrarily large text collections by exchanging additional computation for reduced memory use through excerpt selection and double-pass pooling.
- Conclusion: MIL-BERT classifies arbitrarily large text collections by exchanging compute for memory savings through double-pass pooling and idempotent selectors.The approach is intended for long-text and multiple-instance-learning problems.
- Conclusion: The method produces new state-of-the-art results for trigger-warning detection and political-bias detection, with limited authorship-profiling improvements and competitive long-document results.The authors propose applying the techniques to more multiple-instance-learning, retrieval, and related problems.
- Future work: The paper identifies future improvements to selectors and classifiers for applications involving more data and related tasks such as retrieval.The authors also express interest in applying the techniques beyond text.
9 Limitations
The authors identify limitations in instance-label availability and memory profiling, particularly because PyTorch memory reservation obscures actual allocation behavior.
- Instance-labeled data for article-bias and triggering-passage experiments are limited.The authors also did not identify whether bags containing these instances appeared during training.
- PyTorch memory reservation impeded in-depth profiling of algorithm memory usage.Despite clearing cached CUDA memory and running models in separate processes, reported values may be closer to GPU high-water marks.
- Memory and compute requirements increase with more classes, but measured memory does not scale as expected between 3-class and 8-class MIL-BERT.The authors attribute this discrepancy to limited visibility into memory allocation and fragmentation.
10 Ethics Statement
The algorithm can support classification of otherwise infeasible arbitrarily large text and provides interpretability through passage selection, while raising risks around protected-characteristic inference.
- Learning from arbitrarily large data creates opportunities for classification problems that would otherwise be infeasible.
- Selecting passages makes the algorithm interpretable, but improved author-profiling performance could enable inference of protected characteristics from suitable training data.
- The approach also presents capabilities for identifying trigger warnings.
A Appendix
The appendix records the computing environments used for training and for evaluating memory requirements and single-bag training time.
- The models were trained on a variety of machines.
- The longest-text dataset was trained using 6 A6000 GPUs for 3 days.
- Memory requirements and single-bag training time were evaluated using a single RTX 3090.