nucleus_sampling_analysis

Research code for analyzing failure modes of nucleus (top-p) sampling and comparing it against a probability-threshold alternative on WikiText-103.

What this repo now contains

• nucleus_sampling_analysis/: the new experiment package
• scripts/run_experiment.py: CLI entrypoint for the full study
• scripts/plot_results.py: regenerate plots from saved run artifacts
• tests/: unit, synthetic, and integration coverage
• papercode_with_threshold_sampling/: original prototype kept as reference

Study design

The pipeline:

1. Loads WikiText-103 validation and test prompts.
2. Tunes one global threshold tau per temperature on validation prompts.
3. Compares greedy, temperature, top-k, nucleus, and threshold samplers.
4. Logs per-step distribution diagnostics such as entropy, effective support, retained set size, and sampled-token probability.
5. Saves aggregate metrics and plots for repetition, diversity, perplexity, and stressed peaked/flat regimes.

Quick start

Create or update an environment with compatible dependencies first. The current machine in this workspace has torch 2.0.0 with transformers 4.57.6, which is not compatible for model loading. A working setup is:

python -m pip install -e .

Run the full experiment:

python scripts/run_experiment.py --model-name distilgpt2 --output-dir artifacts/distilgpt2_run

Regenerate plots from an existing run:

python scripts/plot_results.py artifacts/distilgpt2_run

Run the local LLM-as-a-judge stage on an existing run:

python scripts/run_llm_judge.py --run-dir artifacts/distilgpt2_run --judge-model-name your-local-judge-model

Run the local test suite:

python -m unittest discover -s tests -v

Main outputs

Each run writes:

• config.json
• dataset_manifest.json
• threshold_tuning.jsonl
• generation_records.jsonl
• step_stats.jsonl
• run_summaries.jsonl
• plot PNGs in the same run directory

The judge stage writes these under RUN_DIR/judge/:

• judge_pairs.jsonl
• judge_results.jsonl
• judge_summary.json

Notes on threshold tuning

The implementation does not hard-code the paper's 12.38 human perplexity target. Instead, it tunes tau on validation prompts using a combined objective:

• lower generated negative log-likelihood
• lower repetition
• mild penalty when retained support falls outside a target band

This keeps the threshold rule simple, reproducible, and aligned with the project plan.