DiffSampling: Enhancing Diversity and Accuracy in Neural Text Generation

Code and results of the experiments presented in DiffSampling: Enhancing Diversity and Accuracy in Neural Text Generation, published in Transactions on Machine Learning Research.

TL;DR on DiffSampling

DiffSampling is a family of decoding strategies that leverage mathematical analysis to avoid incorrect tokens and possibly increase diversity with safe text. After sorting the probability distribution of the next token in descending order, we compute the forward difference approximation of its discrete derivative, and we use it to remove tokens after its minimum value (possibly after imposing a total probability lower bound or a minimum probability relaxation).

Then, we can also sample at higher temperatures among the preserved tokens to increase output diversity while maintaining its accuracy.

A transformers-based implementation is reported in diffsampling/diffsampling.py.

Results

The directory results/ contains the generated outputs and relative scores for math problem resolution, extreme summarization, Divergent Association Task, and WritingPrompts-based story generation computed per the instructions below.

To easily visualize the results as reported in the paper, we suggest using the show_results.ipynb notebook, which also contains the code used to generate the distribution shown above and an illustrative example on how our methods improve upon existing strategies.

To test new configurations, you can just modify one of the config files and run the experiments as detailed here.

Math Problem Resolution

To generate the responses with all the generation configs listed in the config file, you can use:

python diffsampling/main_math.py --model_name metamath --size 7 --seed 1 --dataset metamathqa --num_of_seqs 1000

python diffsampling/main_math.py --model_name metamath --size 7 --seed 1 --dataset gsm8k --num_of_seqs -1

python diffsampling/main_math.py --model_name metamath --size 7 --seed 1 --dataset hendrycks --num_of_seqs -1

repeating with different seeds to obtain more robust results (in our experiments, we used 1, 42, and 121).

Then, to compute the relative accuracy and both across-input and against-greedy diversity scores, just run:

python diffsampling/compute_math_scores.py --path ./results/math/metamathqa/ --model_name metamath --dataset metamathqa

python diffsampling/compute_math_scores.py --path ./results/math/gsm8k/ --model_name metamath --dataset gsm8k

python diffsampling/compute_math_scores.py --path ./results/math/hendrycks/ --model_name metamath --dataset hendrycks

By default, the accuracy is computed using the final answer extraction method described in the original MetaMath paper. However, if you want to test DiffSampling on different models (e.g., Llama3-70B), we suggest adding --modality hf to use the math-verify developed by HuggingFace.

Extreme Summarization

To generate the responses with all the generation configs listed in the config file, you can use:

python diffsampling/main_xsum.py --model_name llama2 --size 7 --seed 1 --num_of_seqs 1000

python diffsampling/main_xsum.py --model_name llama2 --size 7 --seed 1 --num_of_seqs 1000 --use_instruct_version

repeating with different seeds to obtain more robust results.

Then, to compute the relative accuracy (in terms of coherence, rouge-1, and cosine similarity) and diversity scores against both the original dataset and greedy strategy, just run:

python diffsampling/compute_xsum_scores.py --path ./results/xsum/llama2_7/ --model_name llama2

python diffsampling/compute_xsum_scores.py --path ./results/xsum/llama2_7_instruct/ --model_name llama2

Divergent Association Task

To generate the responses with all the generation configs listed in the config file, you can use:

python diffsampling/main_dat.py --model_name llama3 --size 8 --seed 1 --num_of_seqs 100

python diffsampling/main_dat.py --model_name llama3 --size 8 --seed 1 --num_of_seqs 100 --use_instruct_version

Then, to compute the actual scores, run:

python diffsampling/compute_dat_scores.py --path ./results/dat/llama3_8/

python diffsampling/compute_dat_scores.py --path ./results/dat/llama3_8_instruct/

NOTE: To compute the DAT score, you first need to download glove.840B.300d.txt from here and place it inside the DiffSampling folder.

Story Generation

To generate stories based on WritingPrompts with all the generation configs listed in the config file, you can use:

python diffsampling/main_wp.py --model_name llama3.2 --size 3 --seed 1 --repeated_outs 5 --num_of_seqs 500

python diffsampling/main_wp.py --model_name llama3.2 --size 3 --use_instruct_version --seed 1 --repeated_outs 5 --num_of_seqs 500

Then, to compute the actual scores, run:

python diffsampling/compute_wp_scores.py --model_name llama3.2 --path ./results/wp/llama3.2_3/

python diffsampling/compute_wp_scores.py --model_name llama3.2 --path ./results/wp/llama3.2_3_instruct/

We also verify that the stories generated using DiffSampling are not perceived as unsafer than those generated by the baselines through the Llama-Guard model. To run safety checks, just use:

python diffsampling/safety_check.py --path ./results/wp/llama3.2_3/

python diffsampling/safety_check.py --path ./results/wp/llama3.2_3_instruct/

Docker

The code presented here requires the packages reported in requirements.txt. If you prefer using a Docker container, we provide you with a Dockerfile for Nvidia GPUs (Dockerfile.nv).

How to Cite

@article{franceschelli2025diffsampling,
      title={DiffSampling: Enhancing Diversity and Accuracy in Neural Text Generation},
      author={Giorgio Franceschelli and Mirco Musolesi},
      journal={Transactions on Machine Learning Research},
      issn={2835-8856},
      year={2025},
      url={https://openreview.net/forum?id=kXjHbMvdIi},
      note={}
}