ESLM: Risk-Averse Selective Language Modeling for Efficient Pretraining

TL;DR

ESLM introduces a risk-aware, token-level selection method for language model pretraining that reduces computational costs while maintaining or improving performance, by focusing on high-risk, informative tokens during training.

Contribution

This paper presents ESLM, a novel risk-averse token selection algorithm for efficient language model pretraining, with theoretical foundations and practical improvements demonstrated on GPT-2.

Findings

Reduces training FLOPs significantly.

Maintains or improves perplexity and downstream tasks.

Scales across model sizes and datasets.

Abstract

Large language model pretraining is compute-intensive, yet many tokens contribute marginally to learning, resulting in inefficiency. We introduce Efficient Selective Language Modeling (ESLM), a risk-aware algorithm that improves training efficiency and distributional robustness by performing online token-level batch selection. ESLM leverages per-token statistics (e.g., entropy or loss) and applies value-at-risk thresholding to retain only the most informative tokens per batch. This data-centric mechanism reshapes the training loss, prioritizing high-risk tokens and eliminating redundant gradient computation. We frame ESLM as a bilevel game: the model competes with a masking adversary that selects worst-case token subsets under a constrained thresholding rule. In the loss-based setting, ESLM recovers conditional value-at-risk loss minimization, providing a principled connection to distributionally robust optimization. We extend our approach to Ada-ESLM, which adaptively tunes the selection confidence during training. Experiments on GPT-2 pretraining show that ESLM significantly reduces training FLOPs while maintaining or improving both perplexity and downstream performance compared to baselines. Our approach also scales across model sizes, pretraining corpora, and integrates naturally with knowledge distillation.