DistilBERT Fine-Tuned on IMDB (Whole Word Masking)
This repository contains a domain-adapted version of distilbert-base-uncased fine-tuned on the IMDB movie
reviews dataset using Whole Word Masking (WWM) for Masked Language Modeling (MLM).
Unlike standard training, this model was fine-tuned using Whole Word Masking (WWM) collator strategy. By masking entire words rather than individual subwords
during training, the model is deprived of partial word hints (e.g., seeing ##tography to guess cinema). This creates a significantly harder predictive objective
designed to force the model to learn deeper sentence-level context, higher-level semantic contexts
and multi-subword film concepts (e.g., "cinematography", "masterpiece").
📊 Performance & Benchmark
Performance was evaluated using Perplexity on a large evaluation benchmark of 10,000 sequence chunks derived from the official IMDB test split.
Because each data collator (see note below) evaluates the model under its own masking strategy, baseline pre-trained perplexity varies depending on the evaluation objective.
| Experiment Run | Masking Strategy | Eval Mask Prob | Pre-Trained Perplexity | Fine-Tuned Perplexity | Relative Reduction |
|---|---|---|---|---|---|
| Standard Baseline Run | Subword Masking | 15% | ~23.00 | ~10.00 | ~2.3x lower |
| WWM Run (This Model) | Whole Word Masking | 20% | ~45.00 | ~15.00 | ~3.0x lower 🟢 |
Note: While the raw final perplexity for WWM (15.00) appears higher than the subword model (10.00), WWM is an inherently harder predictive task. Relative to its starting point, Whole Word Masking achieved a larger relative drop in uncertainty (~3.0x vs. ~2.3x). For a discussion and comparison with DistilBERT-finetuned-imdb-standard, which uses standard subword masking, see the Comparison with Standard Subword Masking section below.
🛠️ Training Details
- Base Model:
distilbert-base-uncased - Training Dataset Size: 50,000 sequence chunks (128 max token length)
- Evaluation Dataset Size: 10,000 sequence chunks (Official IMDB test split)
- Hardware: NVIDIA T4 GPU (Google Colab)
- Collator Strategy: Custom Whole Word Masking (
whole_word_masking_data_collator, 20% probability)
Training hyperparameters
The following hyperparameters were used during training:
- learning_rate: 2e-05
- weight_decay: 0.01
- train_batch_size: 64
- eval_batch_size: 64
- seed: 42
- optimizer: Use OptimizerNames.ADAMW_TORCH_FUSED with betas=(0.9,0.999) and epsilon=1e-08 and optimizer_args=No additional optimizer arguments
- lr_scheduler_type: linear
- num_epochs: 3.0
- mixed_precision_training: Native AMP (
fp16=True)
Training results
| Training Loss | Epoch | Step | Validation Loss | Model Preparation Time |
|---|---|---|---|---|
| 2.9545 | 1.0 | 782 | 2.8006 | 0.0016 |
| 2.8715 | 2.0 | 1564 | 2.7388 | 0.0016 |
| 2.8216 | 3.0 | 2346 | 2.7194 | 0.0016 |
Framework versions
- Transformers 5.13.1
- Pytorch 2.11.0+cu128
- Datasets 4.0.0
- Tokenizers 0.22.2
🚀 How to Use
You can use this model directly with the Hugging Face fill-mask pipeline:
from transformers import pipeline
# Load fine-tuned WWM model
fill_mask = pipeline("fill-mask", model="AlexStamp/distilbert-base-uncased-finetuned-imdb-wwm")
# Test domain-specific fill-mask prompt
prompt = "The director created a visual [MASK] with stunning cinematography."
predictions = fill_mask(prompt)
for pred in predictions:
print(f"Token: {pred['token_str']:<15} Score: {pred['score']:.4f}")
Comparison with Standard Subword Masking
This experiment was performed to investigate whether Whole-Word Masking (WWM) could provide an alternative objective for domain adaptation compared with standard subword-level masking.
With WWM, perplexity decreased from 45.0 to 15.0, indicating substantial improvement in the model's ability to predict masked words in the movie-review domain.
For comparison, the standard subword-masking experiment achieved a perplexity reduction from 23.0 to 10.0 (DistilBERT-finetuned-imdb-standard).
These values should not be interpreted as a strict head-to-head comparison, since the two experiments use different masking strategies and masking probabilities, and the evaluation masking is dynamically generated. Nevertheless, both experiments provide evidence that continued MLM training successfully adapts DistilBERT to the language of movie reviews.
Key takeaway
Whole-Word Masking produced strong domain adaptation, reducing perplexity by approximately two-thirds. The experiment also highlights how the choice of masking strategy influences the difficulty and evaluation of the MLM objective.