ESLM: Risk-Averse Selective Language Modeling with Hierarchical Batch Selection

Melis Ilayda Bal, Volkan Cevher, Michael Muehlebach

Keywords: selective language modeling, risk-averse pretraining, online batch selection, large language model pretraining

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), an online, risk-aware batch selection algorithm that improves training efficiency and distributional robustness. ESLM operates in two phases: $(i)$ instance-level selection via a shallow early-exit model pass that computes proxy per-instance statistics (e.g., loss or entropy) and retains data points using value-at-risk thresholding; and $(ii)$ loss shaping with token-level selection via risk-aware thresholding on per-token scores. This data-centric mechanism reshapes the training objective, 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, linking selective pretraining 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 perplexity and downstream performance compared to baselines. Our approach also scales across model sizes, pretraining corpora, and integrates naturally with knowledge distillation.

Primary Area: foundation or frontier models, including LLMs

Submission Number: 11185