Masked Language Modeling › Key Parameters

When training MLM models, several parameters significantly impact performance. Understanding what each parameter controls and why the canonical values were chosen helps you make principled decisions when adapting MLM training to new domains or resource constraints.

The core training parameters are:

• mask_prob (default: 0.15): Fraction of tokens to mask per sequence. Higher values provide more training signal but destroy more context. The 15% rate from BERT remains standard, though some work explores 40% or higher with adjusted strategies. SpanBERT found that masking contiguous spans with a geometric length distribution performs better than uniform token masking, even at the same 15% rate.
• d_model: Hidden dimension of the transformer. BERT-base uses 768, BERT-large uses 1024. Larger values increase model capacity but require proportionally more compute and memory. The hidden dimension controls how much information each token's representation can encode. Too small, and the model cannot represent complex contextual relationships. Too large, and the model becomes expensive to deploy.
• n_heads: Number of attention heads. Should divide d_model evenly. BERT-base uses 12 heads (64 dimensions each), BERT-large uses 16 heads. Each attention head can specialize in a different type of relationship: some heads learn syntactic dependencies, others learn semantic associations, and still others track coreference chains. Having more heads gives the model more capacity to represent diverse linguistic relationships in parallel.
• n_layers: Number of transformer layers. BERT-base uses 12, BERT-large uses 24. Deeper models capture more complex patterns but are slower to train and infer. Probing studies have shown that lower layers tend to encode syntactic information, while higher layers encode more abstract semantic content. This stratification is one reason why removing the top few layers (or using intermediate layer representations) sometimes improves performance on certain tasks.
• max_len: Maximum sequence length the model can process. BERT uses 512 tokens. Longer contexts require quadratically more memory for attention but capture more context. Most real-world NLP tasks fit within 512 tokens, but document-level tasks like long-form question answering and summarization benefit from longer contexts.
• learning_rate: Typically 1e-4 to 5e-4 for MLM pretraining. BERT used 1e-4 with a linear warmup over the first 1% of training steps. Higher rates speed training but risk instability, especially for large models. RoBERTa found that larger batch sizes combined with higher learning rates (up to 4e-4) accelerated convergence without sacrificing quality.
• batch_size: Larger batches provide more stable gradients. BERT used effective batch sizes of 256 sequences. MLM benefits from large batches since only 15% of tokens contribute to each gradient. RoBERTa trained with batch sizes of 8,000 sequences, finding that large batch training at high learning rates significantly reduced training time while achieving better final performance.