SimpleStrat: Diversifying Language Model Generation with Stratification
Justin Wong, Yury Orlovskiy, Alexander Shypula, Michael Luo, Sanjit A. Seshia, Joseph E. Gonzalez
Abstract
Generating diverse responses from large language models (LLMs) is crucial for applications such as planning/search and synthetic data generation, where diversity provides distinct answers across generations. Prior approaches rely on increasing temperature to increase diversity. However, contrary to popular belief, we show not only does this approach produce lower quality individual generations as temperature increases, but it depends on model's next-token probabilities being similar to the true distribution of answers. We propose SimpleStrat, an alternative approach that uses the language model itself to partition the space into strata. At inference, a random stratum is selected and a sample drawn from within the strata. To measure diversity, we introduce CoverageQA, a dataset of underspecified questions with multiple equally plausible answers, and assess diversity by measuring KL Divergence between the output distribution and uniform distribution over valid ground truth answers. As computing probability per response/solution for proprietary models is infeasible, we measure recall on ground truth solutions. Our evaluation show using SimpleStrat achieves higher recall by 0.05 compared to GPT-4o and 0.36 average reduction in KL Divergence compared to Llama 3.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext 2e9d0446-781a-479e-b965-4fe9bbc2808bCited by top-tier papers3
- TreeSynth: Synthesizing Diverse Data from Scratch via Tree-Guided Subspace PartitioningSheng Wang, Pengan Chen, Jingqi Zhou, Qintong Li et al.NeurIPS 2025 · 9 citations
- Parametric Social Identity Injection and Diversification in Public Opinion SimulationHexi Wang, Yujia Zhou, Bangde Du, Qingyao Ai et al.KDD 2026 · 2 citations
- VOYAGER: A Training Free Approach for Generating Diverse Datasets using LLMsAvinash Amballa, Yashas Malur Saidutta, Chi-Heng Lin, Vivek Kulkarni et al.ACL 2026
Builds on12
- Reflexion: language agents with verbal reinforcement learningNoah Shinn, Federico Cassano, Ashwin Gopinath, Karthik Narasimhan et al.NeurIPS 2023 · 5,828 citations
- Tree of Thoughts: Deliberate Problem Solving with Large Language ModelsShunyu Yao, Dian Yu, Jeffrey Zhao, Izhak Shafran et al.NeurIPS 2023 · 5,068 citations
- Effective Data Augmentation With Diffusion ModelsBrandon Trabucco, Kyle Doherty, Max Gurinas, Ruslan SalakhutdinovICLR 2024 · 380 citations
- Rainbow Teaming: Open-Ended Generation of Diverse Adversarial PromptsMikayel Samvelyan, Sharath Chandra Raparthy, Andrei Lupu, Eric Hambro et al.NeurIPS 2024 · 231 citations
- Pretraining task diversity and the emergence of non-Bayesian in-context learning for regressionAllan Raventós, Mansheej Paul, Feng Chen, Surya GanguliNeurIPS 2023 · 174 citations
Related papers
- Improving Diversity in Language Models: When Temperature Fails, Change the LossAlexandre Verine, Florian Le Bronnec, Kunhao Zheng, Alexandre Allauzen et al.ICML 2025
- Group-Aware Reinforcement Learning for Output Diversity in Large Language ModelsOron Anschel, Alon Shoshan, Adam Botach, Shunit Haviv Hakimi et al.EMNLP 2025 · 1 citation
- Random Selection Reveals Implicit Knowledge Consensus in Code GenerationRen-Biao Liu, Li Xin-Ye, Hui Sun, Yali Du et al.ICML 2026
- Hot or Cold? Adaptive Temperature Sampling for Code Generation with Large Language ModelsYuqi Zhu, Jia Li, Ge Li, Yunfei Zhao et al.AAAI 2024 · 68 citations
- Improving Diversity of Demographic Representation in Large Language Models via Collective-Critiques and Self-VotingPreethi Lahoti, Nicholas Blumm, Xiao Ma, Raghavendra Kotikalapudi et al.EMNLP 2023 · 14 citations
