# Weighted Sampling for Masked Language Modeling

## TL;DR

This paper introduces weighted sampling strategies based on token frequency and training loss to improve masked language modeling, especially for rare tokens, leading to better sentence embeddings and transfer learning performance.

## Contribution

The paper proposes two novel weighted sampling methods for MLM that enhance token representation learning and transferability of language models.

## Key findings

- Weighted sampling improves sentence embeddings on STS benchmark.
- Combining WSBERT with calibration and prompt learning yields further gains.
- Weighted sampling enhances transfer learning on GLUE benchmark.

## Abstract

Masked Language Modeling (MLM) is widely used to pretrain language models. The standard random masking strategy in MLM causes the pre-trained language models (PLMs) to be biased toward high-frequency tokens. Representation learning of rare tokens is poor and PLMs have limited performance on downstream tasks. To alleviate this frequency bias issue, we propose two simple and effective Weighted Sampling strategies for masking tokens based on the token frequency and training loss. We apply these two strategies to BERT and obtain Weighted-Sampled BERT (WSBERT). Experiments on the Semantic Textual Similarity benchmark (STS) show that WSBERT significantly improves sentence embeddings over BERT. Combining WSBERT with calibration methods and prompt learning further improves sentence embeddings. We also investigate fine-tuning WSBERT on the GLUE benchmark and show that Weighted Sampling also improves the transfer learning capability of the backbone PLM. We further analyze and provide insights into how WSBERT improves token embeddings.