PaperlayerView exact revision on arXiv

Source-linked AI summary

PRISM: Promoting Sampling Diversity of Large Language Models with Implicit Stochastic Modulation

Qixin Xu, Zihan Zhu, Xiaojiang Liu, Bohan Lyu

arXiv:2610.32101v1cs.LG

TL;DR

Post-trained LLMs concentrate on familiar semantic modes, limiting discovery through repeated sampling, while existing approaches leave a gap in model-level control of multiple valid outputs. PRISM combines submodular curation of diverse-yet-correct solutions with stochastic special-token embedding modulation. It shifts the creative-writing quality-diversity frontier outward, improves reasoning results, and provides a stronger foundation for RL.

  • Problem

    Post-trained LLMs concentrate on familiar semantic modes, limiting discovery through repeated sampling, while existing methods rarely provide model-level control for producing different valid outputs from the same prompt.

    Source: §1 · pp. 1–2

  • Method

    PRISM pairs prompts with diverse-yet-correct solutions through submodular curation and uses Gaussian noise on special-token embeddings as an implicit signal for selecting among them.

    Source: §1 · p. 2

  • Results

    PRISM shifts the writing quality-diversity frontier outward, improves Pass@N on reasoning, and provides a stronger RL foundation that prevents high-pass collapse during standard post-training.

    Source: §1 · p. 2, §6 · p. 9

  • Takeaways & Limitations

    The results indicate that diversity treated as a learnable training-time signal can serve as a unified mechanism across generation, reasoning, and RL.

    Source: §6 · p. 9

    Evaluation relies on LLM judges and one embedding model, and experiments cover only 7B–8B models.

    Source: Limitations · p. 9

Abstract

from arXiv · show

Large language models (LLMs) exhibit mode collapse, where their output distribution concentrates on familiar semantic modes. When diversity disappears, repeated sampling becomes semantic repetition rather than discovery, making rare but stronger solutions unlikely to surface and leaving self-evolution and continual learning difficult to sustain. Most remedies intervene only at inference time, shaping sam- ples without making diversity a model capability. We propose PRISM, a training-time framework that internalizes diversity within the model. PRISM combines (i) submodular curation that pairs each prompt with diverse-yet-correct multi-teacher solutions, and (ii) stochastic embedding modulation, which turns special-token perturbations into implicit signals for sampling distinct outputs. A theoretical analysis establishes the stability of this mechanism. On WritingPrompts and BookMIA, PRISM moves the quality-diversity frontier outward, achieving the best generation quality in all model-benchmark settings while maintaining the strongest or near-strongest semantic diversity across two open-weight model families. It also improves sampling efficiency on math and code benchmarks. Beyond inference, PRISM yields a stronger foundation for reinforcement learning, sustaining diverse rollouts and mitigating the diversity collapse commonly observed during RL training.

1 Introduction

PRISM addresses mode collapse by teaching models to select among diverse, correct outputs through submodular curation and stochastic embedding modulation. It improves creative-writing quality and diversity, math/code sampling efficiency, and RL rollout diversity.

1 Introduction

  • Post-trained LLMs concentrate on familiar semantic modes, so repeated sampling revisits trajectories and makes rare, stronger outputs harder to discover; existing remedies leave a gap in model-level control.Source: §1 · pp. 1–2
  • PRISM pairs per-prompt submodular curation of diverse-yet-correct solutions with stochastic special-token embedding perturbations that act as signals for selecting valid trajectories.The paper also analyzes the stability of this noise-injection mechanism.Source: §1 · p. 2
  • PRISM shifts the creative-writing quality–diversity frontier outward, achieving the best quality and strongest or near-strongest semantic diversity across model-benchmark settings.The evaluation covers WritingPrompts and BookMIA across two open-weight model families.Source: §1 · p. 2
  • During RL, PRISM preserves diverse rollouts and avoids the sampling collapse observed with standard SFT initialization in math and code reasoning.Source: §1 · p. 2
  • On eight math and code benchmarks, PRISM discovers more correct solutions than SFT and prior diversity-enhancing baselines under the same inference budgets.Source: §1 · p. 3

2 Methodology

PRISM first curates diverse, valid reasoning trajectories, then trains models to use noise on special-token embeddings as a signal for selecting among them. Its analysis links this training objective to embedding regularization and noise-conditioned generation of valid outputs.

2.1 Diverse Reasoning Trajectory Sampling via Submodular Optimization

  • For each prompt, PRISM selects a fixed-size subset of candidate trajectories to maximize coverage of discrete reasoning features.Features are constructed per domain: clustered SBERT embeddings for creative writing, solution-token n-grams for math, and n-grams plus AST subtrees for code.Source: §2.1 · p. 3
  • Human annotations found that submodular selection produced more queries with at least three distinct solutions than random sampling.The greedy selection algorithm guarantees a 1−1/e approximation to optimal feature coverage with n samples.Source: §2.1 · p. 3, Table 1 · p. 3

2.2 Noisy Special-token Embeddings as Implicit Diversity Signals

  • PRISM adds Gaussian noise to special-token embeddings during training and inference while leaving other token embeddings unchanged.The perturbation is intended to stimulate diverse outputs while keeping the input’s core semantic information intact.Source: §2.2 · p. 3
  • Because each input is paired with multiple curated answers under fresh noise draws, the model can learn noise as an implicit selector among valid answers.Varying the noise at inference then yields varying outputs.Source: §2.2 · p. 4

2.3 Training Target Analysis

  • The analysis treats Gaussian embedding noise as regularization that penalizes high-curvature regions of the loss landscape.It also decomposes per-instance loss into answer correctness and noise-conditioned specificity, with the latter supporting selection of a particular valid trajectory as noise varies.Source: §2.3 · p. 4

3 Experiments

PRISM is evaluated on creative writing, math and code reasoning, and reinforcement-learning post-training, testing both sampling outcomes and RL initialization. Across these settings, the reported results show improved quality-diversity or sampling-efficiency outcomes and a stronger RL foundation.

3 Experiments

  • The evaluation spans creative writing, verifiable math and code reasoning, and RL post-training; creative-writing tests use two benchmarks and two open-weight model families.WritingPrompts and BookMIA use 100 test prompts and 30 responses per prompt; reasoning evaluation covers eight benchmarks.Source: §3 · p. 4, §3.1 · p. 4, §3.2 · p. 6

3.1 Creative Writing

  • PRISM achieves the highest quality in all four creative-writing model-benchmark cells and shifts the quality-diversity frontier outward; Table 2 compares diversity and quality across both benchmarks.On Qwen, GEM exceeds PRISM by at most 1.3 diversity points while giving up 1.6 to 6.4 quality points; on Llama, PRISM is higher than GEM on both axes in both benchmarks.Source: §3.1 · p. 5, Table 2 · p. 5
  • Ablations show that curated data without noise loses 6 to 9 quality points, while random training data with noise loses 3.6 to 4.6 diversity points.Inference-time noise raises diversity but costs 3 to 6 quality points on an SFT checkpoint; under the noise-trained checkpoint it raises both, and removing inference noise forfeits quality and diversity.Source: §3.1 · pp. 5–6

3.2 Math and Code Reasoning

  • PRISM attains the highest Pass@N on seven of eight reasoning benchmarks, while NEFTune is higher on MBPP+ and AIME 2025 is tied with SFT.Against SFT, PRISM gains +13.33 on AIME 2024 and +10.0 on AMC 23; Figure 2 shows gains across temperatures.Source: Figure 2 · p. 5, §3.2 · p. 6, Table 3 · p. 6

3.3 Reinforcement Learning

  • PRISM has higher pre- and post-RL pass@64 than SFT in every cell; after RL, SFT loses pass@64 in all four code settings, while PRISM improves in three and drops marginally in one.Table 4 reports pass@64 before and after RL, with AIME averaged over AIME 2024 and AIME 2025.Source: §3.3 · p. 7, Table 4 · p. 7

4 Analyses

The analyses compare PRISM with reasoning-diversity baselines and examine whether algorithmic diversity accompanies its sampling gains. They report stronger high-budget code results and a shift toward more distinct algorithms per query.

4.1 Other Diversity Enhancing Baselines

  • PRISM achieves the highest Pass@64 on three code benchmarks, outperforming GEM and the extensively trained Llama-3.1-8B-Instruct model, while GEM performs well on HumanEval.Table 5 reports Pass@64 performance on code benchmarks and compares sampling efficiency with GEM and DBS.Source: §4.1 · p. 7, Table 5 · p. 7
  • DBS lowers all Pass@1 and Pass@N scores versus its respective original model, whereas PRISM leads at higher sampling budgets on three benchmarks.The analysis compares DBS on base and instruct models and includes GEM as an SFT baseline.Source: §4.1 · p. 7

4.2 Sourcing the Effectiveness of PRISM

  • PRISM shifts the per-query distribution toward more unique algorithms than SFT, and queries with more distinct algorithms gain substantially on Pass@N, with the gap widening from Pass@64 to Pass@128.Generated solutions were categorized into 39 algorithm classes and automatically labeled for the analysis.Source: §4.2 · pp. 7–8, Figure 3 · p. 8

5 Related Work

Related work describes diversity loss after post-training, diversity interventions at inference time, and training-time approaches that address data coverage or model control. It also notes that RL objectives can raise Pass@1 while narrowing rollouts.

Mode Collapse in LLMs

  • Post-training can improve helpfulness or correctness while making outputs more repetitive, with repetition appearing in open-ended generation and verifiable reasoning.Prior analyses trace much of the loss to post-training data and learned weights.Source: §5 · p. 8

Inference-Time and Search-Based Diversity

  • Inference-time decoding, prompting, guidance, and structured search can improve exploration, but their diversity comes from external interventions such as prompts, archives, or evaluators.Reasoning-oriented search includes program search, mathematical discovery, reward-guided reasoning, behavioral clustering, and parallel thought generation.Source: §5 · p. 8

Training-Time Approaches for Diversity

  • Training-time diversity work often separates data construction from model control, while one-problem-one-solution supervision can miss valid paths and extra traces may still collapse into one dominant mode.Embedding-noise approaches are described as robustness regularizers, while data-centric methods target coverage.Source: §5 · pp. 8–9

Diversity in RL Post-Training

  • Outcome-only and reverse-KL-style RL objectives can raise Pass@1 while narrowing rollouts and hurting Pass@N and test-time scaling.Recent approaches respond with rare-outcome rewards, mass-covering divergences, diversity rewards, or count-based exploration.Source: §5 · p. 9

6 Conclusion

The conclusion presents PRISM as a training-time approach that makes output diversity controllable through data curation and embedding modulation. It reports benefits across creative writing, reasoning, and RL training, and connects controllable diversity to continual self-improvement.

6 Conclusion

  • PRISM uses submodular data curation and stochastic embedding modulation to make output diversity a controllable model capability.The conclusion reports an outward shift in the creative-writing quality-diversity frontier and improved Pass@N on verifiable reasoning.Source: §6 · p. 9
  • Preserving exploratory rollouts keeps sampling from degenerating into repetition and helps models discover stronger candidates over time.The conclusion links controllable diversity to the foundation for continual self-improvement.Source: §6 · p. 9

Limitations

PRISM’s curation relies on domain-specific discrete features and greedy coverage to select diverse trajectories. The section also identifies limits in candidate-pool costs, task-specific noise selection, analysis, and evaluation scope.

Limitations

  • PRISM’s candidate pool requires multiple solutions; creative-writing pools were generated by sampling eight teacher models, a cost that recurs without off-the-shelf multi-solution data.Noise scale is task-dependent and requires a selection pass for each task family.Source: Limitations · p. 9

A.1 Feature Engineering from Reasoning Trajectories

  • Trajectory features are domain-specific: natural-language solutions use unique n-grams, while code trajectories use Abstract Syntax Tree features.These discrete features represent local patterns in text or structural and logical idioms in code.Source: §A.1 · p. 14

A.2 The Submodular Optimization Problem

  • Feature coverage is monotone and submodular: adding a trajectory contributes its features not already covered, with diminishing marginal gains as the selected set grows.The section states monotonicity and submodularity as propositions.Source: §A.2 · p. 14

A.3 The Greedy Selection Algorithm

  • The greedy selector repeatedly adds the remaining trajectory with the largest number of features not yet covered.It starts with the trajectory having the most unique features and updates the covered-feature union after each selection.Source: §A.3 · p. 15

A.4 Theoretical Performance Guarantee and Its Significance

  • For non-negative monotone submodular coverage under a size constraint, greedy selection guarantees f(A_n) ≥ (1 − 1/e)f(A_OPT).The guarantee is attributed to Nemhauser et al. (1978).Source: §A.4 · p. 15

B Proof of the Regularization Effect of Noise Injection

The proof estimates how independent Gaussian perturbations to selected embedding vectors change expected loss. A second-order expansion leaves a Hessian-trace correction at order σ^2, with higher-order terms of order σ^3.

B Proof of the Regularization Effect of Noise Injection

  • The analysis perturbs selected embedding rows with independent zero-mean Gaussian noise and expands the loss around the unperturbed embedding matrix.The expansion uses gradient and Hessian blocks evaluated at the unperturbed matrix.Source: B Proof of the Regularization Effect… · p. 15
  • Taking expectation removes the linear contribution because the noise has zero mean; independent noise vectors also make the off-diagonal quadratic terms vanish.The quadratic contribution is handled by separating diagonal and off-diagonal terms.Source: B Proof of the Regularization Effect… · p. 15
  • The resulting expected loss adds a second-order correction proportional to σ^2 and the summed Hessian-block traces, with remainder O(σ^3).The remainder order assumes sufficient smoothness of the loss.Source: B Proof of the Regularization Effect… · p. 15

C Detailed Analysis of the Loss Objective Decomposition

The training objective is negative log-likelihood over curated correct answers under freshly sampled embedding noise. Its decomposition separates assigning probability to the correct-answer set from selecting a particular answer within that set.

C Detailed Analysis of the Loss Objective Decomposition

  • Training minimizes NLL over prompts, curated targets, and freshly sampled noise, conditioning on the prompt’s noisy embedding representation.The targets are drawn from the curated answer set A(x).Source: C Detailed Analysis of the Loss… · p. 15
  • The semantic-correctness component encourages the model to maximize total probability mass assigned to all correct answers.The passage describes this as learning the task’s semantic constraints so generations are valid solutions.Source: C Detailed Analysis of the Loss… · p. 16
  • The noise-conditioned specificity component scores the selected target conditionally within the correct-answer set; independent noise draws favor output-specific routing during finite-step optimization despite a noise-invariant optimum also attaining the expected-loss optimum.The same prompt is paired with different correct targets under independently sampled noise, and varying noise at inference yields varying outputs.Source: C Detailed Analysis of the Loss… · p. 16

D.1 Benchmarks

The appendix specifies evaluation datasets, model and baseline comparisons, artifact terms, hyperparameters, and training-data construction. The noise-sensitivity analysis reports increasing diversity but an inverted-U quality pattern as σ varies.

D.1 Benchmarks

  • The study evaluates creative writing on WritingPrompts and BookMIA, math on MATH500, AMC 2023, and AIME 2024/2025, and code on HumanEval, MBPP, and their plus variants.The writing setup uses 100 held-out prompts per benchmark and 30 responses per prompt; math and code passages describe their benchmark composition.Source: §D.1 · p. 16

D.2 Models

  • Creative-writing experiments use Qwen2.5-7B-Instruct and Llama-3.1-8B-Instruct; math and code experiments use a Llama-3.1-8B-Base checkpoint after warm-up SFT.Larger Qwen and Llama instruction-tuned models are additional creative-writing reference points.Source: §D.2 · pp. 16–17

D.3 Baselines

  • Baselines include vanilla sampling, SFT, GEM, NEFTune, VS-Standard, and FIRE for creative writing, plus Base, SFT, and Diverse Beam Search for reasoning.The comparisons include training-time and inference-time approaches.Source: §D.3 · p. 17

D.4 Artifact Licenses and Terms of Use

  • The appendix documents licenses and distribution terms for the datasets and models used, including research-only conditions for WritingPrompts and BookMIA snippets.It also states that no proprietary content is redistributed.Source: §D.4 · pp. 17–18

D.5 Hyperparameters

  • Noise is applied independently to special structural-token embeddings during training and inference, with σ = 0.005 for creative writing and σ = 0.01 for reasoning.The sensitivity sweep finds monotonically rising diversity and inverted-U quality; σ values from 0.001 to 0.01 strictly Pareto-improve over no noise.Source: §D.5 · p. 18, Table 6 · p. 18

F Fintune Backbone Details

  • Training data uses submodular selection of four diverse-yet-high-quality completions per creative-writing prompt and four answers per math problem; code selection uses n-gram and AST features.Creative-writing candidates come from eight teacher models, while code data combines several sources and selection sizes.Source: F Fintune Backbone Details · pp. 18–19

G RL Training Details

The RL stage uses verl and task-specific training data for math and code. Hyperparameters are documented separately for the two tasks.

G RL Training Details

H Human Check Results

This section examines PRISM’s output diversity through narrative and mathematical case studies, including contrasting solution strategies and results. It also describes criteria and an LLM-based pipeline for assessing solution diversity.

H Human Check Results

  • Human evaluators treated code solutions as distinct when they differed in core algorithm, data structure, or time complexity, and math solutions when they relied on different theorems or proof approaches.The math criteria count only theorems or lemmas used in the final solution, not in failed attempts or side notes.Source: H Human Check Results · pp. 19–20

I LLM-As-Judge Pipeline

  • The algorithm-classification pipeline used Gemini 2.5 Pro to derive a pool of 39 algorithms from 1,000 sampled answers, then had Qwen3-8B classify each answer against that pool.The prompt directs the judge to select the closest listed algorithm and permits multiple algorithm numbers when relevant.Source: I LLM-As-Judge Pipeline · pp. 20–21

J.1 Creative Writing: A Case Study in Narrative Diversity

  • In 10 WritingPrompts samples, SFT followed a near-identical three-act story structure in 9, while PRISM produced four distinct narrative framings.The PRISM framings included epistolary revelation, a supernatural time loop, a mystery/thriller, and quiet literary meditation; the passage says this variation preserves per-sample quality.Source: §J.1 · p. 21

J.1.1 Baseline Model Generations

  • The SFT story sample follows a sentimental-reunion arc in which Sarah learns her grandfather left rather than died and reconnects with family.The sample’s resolution depicts Sarah feeling she understands her grandmother and is “at last, completely known.”Source: §J.1.1 · pp. 21–22

J.1.2 PRISM Model Generations

  • PRISM’s four story samples vary from uncovering hidden lineage through letters to a future-self time loop, a family-fraud thriller, and an unresolved meditation on inherited silence.The examples show distinct premises and narrative resolutions, including an archive of letters and a protagonist who leaves a letter in the attic.Source: §J.1.2 · p. 22

J.2 Math Reasoning: A Case Study in Solution Diversity

  • In 10 MATH500 samples, SFT used the same substitution method throughout, while PRISM attempted four structurally distinct approaches.PRISM's approaches were direct substitution in x, inverse substitution in y, a transformation t = x^2, and a vertex-based heuristic; the passage links this broader exploration to Pass@N gains.Source: §J.2 · p. 23

J.2.1 Baseline Model Generations

  • The SFT math example solves by direct substitution in x, minimizes the squared distance, and reports a + b = 8.The supplied solution passage presents the result after evaluating its critical points and simplifying the minimum distance.Source: §J.2.1 · p. 23

J.2.2 PRISM Model Generations

  • The supplied PRISM math generations include both direct substitution in x and inverse substitution in y, illustrating distinct formulations of the same optimization problem.The direct-substitution example reports a + b = 7; the inverse-substitution example rewrites x^2 in terms of y and minimizes a quadratic in y.Source: §J.2.2 · p. 23
  • The variable-transformation solution identifies t = 2 as a minimum and concludes a = 5, b = 2, and a + b = 7.It describes t = x2 as the substitution.Source: §J.2.2 · p. 24
  • The vertex-based heuristic treats the parabola’s vertex as the closest point to the origin and concludes a = 9, b = 2, and a + b = 11.The generation identifies the vertex as (0, −3 √ 2).Source: §J.2.2 · p. 24

J.3 Code Generation: A Case Study in Algorithmic Diversity

In a 256-sample Bell-number case study, PRISM generates more algorithmic approaches than SFT and achieves a higher task success rate. The examples contrast SFT’s repeated approach with PRISM’s varied solution framings.

J.3.1 Baseline Model Generations

  • Across 256 samples, SFT produces only 2 algorithm types and solves the Bell-number task correctly once.Its samples are concentrated in closed-form combinatorial summation (220) and dynamic programming (35).Source: §J.3 · p. 24

J.3.2 PRISM Model Generations

  • Across 256 samples, PRISM produces 10 algorithm types and 29 correct solutions, with correct outputs spanning three distinct framings.The listed routes include dynamic programming, memoized recursion, and combinatorial summation; further attempts use simulation, number theory, modular arithmetic, and bit manipulation.Source: §J.3 · p. 24

J.3 Code Generation: A Case Study in Algorithmic Diversity

  • Success on the case-study query rises from 0.4% to 11.6% as PRISM’s samples span structurally distinct algorithmic routes.The passage connects this richer search space with the query’s higher success rate.Source: §J.3 · p. 24

Related papers

Loading 2610.32101v1…