Token Masking Improves Transformer-Based Text Classification
Xianglong Xu, John Bowen, Rojin Taheri
TL;DR
The paper addresses overfitting in transformer-based text classification by introducing token masking regularization that randomly replaces tokens with a [MASK] token during training. This data-level perturbation induces gradient averaging across masked configurations, acting as an implicit ensemble and encouraging models to rely on deeper contextual relations rather than surface co-occurrences. Across LinCE benchmarks for language identification and sentiment analysis using multiple models, the method achieves consistent gains, with a robust default masking rate around 0.1 and capacity- and task-dependent behavior. The approach is simple yet effective, offering practical cross-lingual transfer benefits and a foundation for adaptive masking strategies in future work.
Abstract
While transformer-based models achieve strong performance on text classification, we explore whether masking input tokens can further enhance their effectiveness. We propose token masking regularization, a simple yet theoretically motivated method that randomly replaces input tokens with a special [MASK] token at probability p. This introduces stochastic perturbations during training, leading to implicit gradient averaging that encourages the model to capture deeper inter-token dependencies. Experiments on language identification and sentiment analysis -- across diverse models (mBERT, Qwen2.5-0.5B, TinyLlama-1.1B) -- show consistent improvements over standard regularization techniques. We identify task-specific optimal masking rates, with p = 0.1 as a strong general default. We attribute the gains to two key effects: (1) input perturbation reduces overfitting, and (2) gradient-level smoothing acts as implicit ensembling.