Weighted Ensemble method improves rare failure estimates in language models
Weighted Ensemble Sampling for Rare-Event Estimation in Language Models
Computation and Language
Summary
Large language models sometimes make rare but serious mistakes that are hard to predict because these mistakes happen very infrequently. The authors introduce a new method called Weighted Ensemble that helps estimate how often these rare failures occur more efficiently without needing to train any extra models. Instead, their method uses signals inside the original language model to focus on parts of the text that are more likely to cause errors. Tests show that this method gives better estimates for rare failures on a variety of language models.
What this means in practice
- •For ai safety engineers: Estimate the probability of rare harmful model outputs more efficiently during AI system deployment.
- •For cloud ai service operators: Monitor and manage rare failure risks in deployed language models without additional model training overhead.
Authors
Zixuan Liu, Gustavo Didier, Fangzheng Wu, Brian Summa, Zizhan Zheng
Abstract
Rare but severe failures of large language models can become practically significant at deployment scale. Estimating their probabilities with direct Monte Carlo sampling is highly inefficient, while proposal-based methods typically require auxiliary learning of an additional proposal distribution. We introduce \textbf{Weighted Ensemble (WE)}, a training-free framework for rare-event estimation in language models. WE samples exclusively from the original language model and improves upon direct Monte Carlo by periodically reallocating the sampling budget toward prefixes that appear more likely to reach the target rare event. To guide this reallocation without learning a separate proposal distribution, we exploit internal LLM representations as prefix-level signals and empirically validate their effectiveness. We provide a theoretical analysis of the proposed algorithm. Experiments on safety-related tasks spanning both small-scale and modern large-scale LLMs demonstrate that WE improves the efficiency and reliability of rare-event probability estimation.