Research map / Large language models / Language, safety and interpretability

SimpleStrat: Diversifying Language Model Generation with Stratification

NeurIPS 2025 · Large language models › Language, safety and interpretability · approach: Uncertainty quantification for LLMs · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes SimpleStrat, a training-free sampling method that has the LLM auto-stratify the solution space into balanced partitions, estimates stratum probabilities, and samples a stratum before generating, improving answer diversity over temperature scaling; introduces CoverageQA.

Problem

LLM resampling for diverse responses is limited by temperature scaling, which degrades quality and fails to fix mode collapse in next-token probabilities (e.g., over-preferring "California" for "Name a US state").

Approach

Use the LLM itself to propose balanced binary partitions of the solution space (auto-stratification), estimate stratum probabilities via forecasting-style prompts (heuristic estimation), and sample a stratum to augment the prompt before generating (probabilistic prompting), orthogonal to temperature.

Key results

• On CoverageQA, SimpleStrat with GPT-4o improves recall by about 0.05 on CoverageQA-Wikipedia across temperatures (up to 0.2 on Curated).
• Across four Llama 3/3.1 models, SimpleStrat reduces average KL divergence from uniform by 0.36 on CoverageQA-Wikipedia (1.14 on Curated).
• SimpleStrat gives gains even at temperature zero, since diversity comes from prompt sampling rather than temperature.

Measured effects reported in the paper

Surprising findings

• Larger Llama models showed on average worse diversity (higher KL divergence) than smaller ones, possibly due to memorization.
• Raising temperature to 1.5 still left GPT-4o sampling Lake Huron once and never Lake Ontario in 100 samples.

Limitations

• Effectiveness depends on the model proposing meaningful partitions and accurate joint probability estimates.
• Possible race and gender biases may propagate into auto-stratification and estimation.
• CoverageQA contains only short-answer questions; multi-step reasoning settings not evaluated.

Experimental setup

Datasets: WikiData

Benchmarks: CoverageQA-Curated CoverageQA-Wikipedia

Metrics: recall over 100 samples KL divergence from uniform precision/F1

Baselines: temperature scaling GPT-4o Claude 3.5 Sonnet Llama 3 baseline sampling

Builds on: temperature scaling stratified sampling diverse beam search

Closest papers by meaning

• Improving Diversity in Language Models: When Temperature Fails, Change the Loss (ICML 2025 · similarity 0.88)
• Free Lunch for Pass@$k$? Low Cost Diverse Sampling for Diffusion Language Models (NeurIPS 2026 · similarity 0.88)
• Towards Scalable Data Diversification for Language Model Pretraining via Leverage Score Sampling (NeurIPS 2026 · similarity 0.88)
• Diversity as a Reward: Fine-Tuning LLMs on a Mixture of Domain-Undetermined Data (NeurIPS 2025 · similarity 0.88)
• Less Diverse, Less Safe: The Indirect But Pervasive Risk of Test-Time Scaling in Large Language Models (ICML 2026 · similarity 0.88)
• Prismatic Synthesis: Gradient-based Data Diversification Boosts Generalization in LLM Reasoning (NeurIPS 2025 · similarity 0.88)