README.md

5.1 Inspecting Next-Token Candidates, Probabilities, and Sampling

This code answers three questions: which candidate tokens a language model chooses among at each step and how confident it is in each one; what temperature and top_p actually change; and when asking the same question multiple times gives different answers, whether the step probabilities show ahead of time where the model is guessing.

For each prompt you pass in, the script does three things:

1. Generates once at temperature 0 (greedy decoding, picking the highest-probability token at every step) and prints a step-by-step table of the chosen token, its probability, the runner-up candidates, and how many candidates are needed to cover 90% of the probability mass.
2. Samples 5 times at each of several temperatures (0.3, 0.7, 1.0, 1.5), flagging any generated token whose probability is below 0.3 as 【token|prob】.
3. Samples 5 more times at temperature 1.0 with top_p = 0.9 so you can compare it directly against temperature 1.0 without top_p.

Every reported probability comes from the model's raw, unscaled distribution (temperature = 1.0). That way, the same token at the same context position carries the same probability value across different temperature runs, making runs easy to compare.

Hardware requirements

This experiment runs a small open-weight model locally on either of two setups:

• Apple Silicon Mac (mps): The default model Qwen/Qwen3.5-2B takes about 4.2 GB in bfloat16; 16 GB of unified memory is recommended. On an 8 GB Mac, switch to --model Qwen/Qwen3.5-0.8B.
• NVIDIA GPU (cuda): 8 GB of VRAM is enough.

If you have neither, the script falls back to CPU (slow, but it works). None of the hosted models on AI Builder Space currently return logprobs (verified October 2026), so there is no API version of this script.

Setup and usage

cd 5_1_logprobs
uv venv && source .venv/bin/activate      # If you do not have uv: python3 -m venv .venv
uv pip install -r requirements.txt        # torch and transformers, ~1 GB
python logprobs.py                        # Downloads the model on first run; ~20s after that on an M3 Ultra
python logprobs.py --prompt "What is the capital of France? Answer in one sentence." --prompt "Recommend an independent coffee shop in Seattle with its name and street address, in one sentence."
python logprobs.py --prompt "Your own question" --prompt "Another question"
python logprobs.py --model Qwen/Qwen3.5-0.8B --temps 0.5 1.0 --samples 10

Running python logprobs.py with no --prompt argument runs the two built-in Chinese prompts stored in DEFAULT_PROMPTS (whose reference outputs are in sample_output/). Passing the two English --prompt strings above runs the two comparison prompts from Section 5.1 (whose reference outputs are in sample_output_en/).

During execution you may see a few lines saying falling back to its reference PyTorch implementation. That just means Qwen3.5's linear attention layers are running the reference PyTorch fallback rather than a custom accelerated kernel. The numerical results are valid, just slightly slower, so you can ignore the message. On an NVIDIA GPU you can run uv pip install flash-linear-attention to enable the accelerated implementation (which is how our RTX 5090 runs were recorded); macOS does not currently have an equivalent package.

Sample output files

Two folders of pre-recorded runs (Qwen/Qwen3.5-2B in bfloat16) are included for inspection if you do not run the model locally:

• sample_output_en/ contains the English-language runs from Section 5.1, recorded on 2026-10-29 on a Mac Studio M3 Ultra (logprobs_Qwen3.5-2B_mps_20261029_1750.json and temperature_step1_mps.json) and an NVIDIA RTX 5090 (cuda_rtx5090_20261029.json and cuda_rtx5090_20261029_console.txt). The prompts and generated completions in these files are in English, while the console log headers in cuda_rtx5090_20261029_console.txt and the "samples" dictionary keys inside the two logprobs JSON files ("温度 0.3", etc.) retain Chinese labels from an earlier Chinese-labelled version of the script.
• sample_output/ contains runs on the two built-in Chinese prompts in DEFAULT_PROMPTS, recorded on 2026-10-23 on the same Mac Studio M3 Ultra (mps_m3ultra_qwen3.5-2b_20261023.json, mps_m3ultra_console.txt, and temperature_step1_mps.json) and RTX 5090 (cuda_rtx5090_qwen3.5-2b_20261023.json and cuda_rtx5090_console.txt). Their prompts, model outputs, and console labels are all in Chinese.

What to look for

In the English runs in sample_output_en/, one prompt asks for a known fact ("What is the capital of France? Answer in one sentence.") and the other forces the model to invent ungrounded local details ("Recommend an independent coffee shop in Seattle with its name and street address, in one sentence."):

• Known-fact prompt: At temperature 0, the model generates "The capital of France is Paris." across 7 steps. Step 1 gives The a probability of 0.98 (with Paris at 0.02, a wording choice between starting a full sentence or answering with a bare noun), and steps 2 through 7 stay at 0.99 to 1.00 (Paris is 1.00), needing just 1 candidate to cover 90% of the probability mass at every step. At temperature 0.7 (and 0.3 on RTX 5090), all 5 sampled runs return that exact same sentence (1 distinct reply; on the Mac Studio seed-0 run at 0.3, 4 of 5 runs match and 1 draws a stray tail token). At temperature 1.0, 2 distinct replies appear across 5 runs ("Paris is the capital of France." alongside "The capital of France is Paris.").
• Ungrounded-detail prompt: At temperature 0, the 24-token greedy reply is "For a cozy atmosphere, visit **The Roasted Bean** located at 1234 Pike Street in Seattle." As soon as the model has to name a shop, the probabilities collapse: step 7 (**) is 0.12 and requires 1,836 candidates on Mac Studio (mps) or 1,770 on RTX 5090 (cuda) to cover 90% probability mass; step 8 (The) is 0.29 (480 candidates on mps, 483 on cuda); and step 9 (Ro) is 0.15 (998 candidates on mps, 1,004 on cuda). Once Ro is picked, finishing the familiar phrase "The Roasted Bean" jumps to 0.43 (asted, 8 candidates) and 0.84 (Bean, 8 to 9 candidates), even though the shop is fabricated. The street number 1234 splits into four single-digit tokens (1 at 0.53, 2 at 0.36 on mps / 0.38 on cuda, 3 at 0.44 on mps / 0.42 on cuda, and 4 at 0.63) followed by Pike (0.43). Even at temperature 0.3 or 0.7, sampling 5 times produces 5 distinct replies with different invented shops or addresses. The prose reads just as smoothly as the France answer, while the step probabilities expose where the model is guessing.
• Temperature and top_p: At temperature 1.5, flattening the logits pushes substantial mass into the long tail. Once a draw lands on a near-zero tail token, subsequent steps condition on that corrupted prefix and spiral into multilingual gibberish containing dozens of tokens near 0.00 probability (all 5 Seattle runs at temperature 1.5, and run 1 of France on Mac Studio with 48 low-probability tokens). Pairing temperature 1.0 with top_p = 0.9 trims away the bottom 10% tail before sampling, pulling all 5 France capital runs back to "The capital of France is Paris."
• Ordering detail: transformers applies temperature scaling before top_p filtering, so temperature 1.5 combined with top_p = 0.9 can still produce gibberish. Once high temperature flattens the distribution, reaching 90% cumulative mass takes far more candidates, leaving top_p almost ineffective. Different inference engines and APIs may apply temperature and top_p in a different order.

Even with the same model at temperature 0, switching machines can change the output. Parallel bfloat16 floating-point additions round slightly differently across MPS and CUDA. When the top two candidates at a step sit only a few percentage points apart, that small rounding shift is enough to swap first and second place and steer the rest of the reply onto a different trajectory. On the English Seattle prompt in sample_output_en/, no top-1 tie occurred across its 24 steps, so both machines picked the same greedy tokens while showing small bfloat16 shifts in probabilities and 90%-mass pool counts (1,836 on MPS vs 1,770 on CUDA at step 7). In the Chinese runs in sample_output/ (mps_m3ultra_* vs cuda_rtx5090_*), you can see a direct temperature-0 flip: the France capital prompt matches closely across both machines (0.69 vs 0.70 on step 1, 0.99 on Paris), whereas at step 3 of the Shanghai coffee shop prompt the top two tokens sit at 0.24 vs 0.21 on Mac Studio and reverse to 0.24 vs 0.22 on RTX 5090, leading the two machines to invent two different shop names and street addresses (0.02 vs 0.01 probability on the first character of the shop name, needing 5,371 vs 5,061 candidates for 90% mass). Temperature 0 only picks whichever candidate scores highest at each step; near a tie, greedy decoding remains sensitive to hardware, precision, and batch size.

Sampling is random, so your sampled completions will not match ours word for word; pass --seed to try a different random seed. Across multiple seeds, the pattern stays the same: the ungrounded coffee shop prompt produces 5 different answers out of 5 at every temperature, the France capital prompt stays consistent below temperature 0.7, and temperature 1.5 sometimes derails into gibberish and sometimes does not.

Try your own prompts

Use --prompt to substitute real prompts from your own projects, ideally one with a single verifiable answer and one that asks the model for specific names, numbers, or citations. Check two things: in the temperature-0 step table, which positions drop below 0.3 probability and require hundreds or thousands of candidates to cover 90% mass; and across multiple sampled runs, whether the output keeps changing at those exact positions. Those spots usually mark facts you should verify in code or supply directly in the context window.

Full JSON output is written to results/, including the greedy step-by-step candidate table and the raw and marked text of every sampled run.

Data for the temperature distribution chart

temperature_step1.py extracts the full-vocabulary step-1 logits (248,320 output logits on Qwen/Qwen3.5-2B, loaded the same way as logprobs.py in bfloat16 with thinking mode turned off), divides them by 0.5, 1.0, and 1.5, applies softmax, and records the top 5 candidates along with "rest" (the combined probability of all remaining tokens) in <out>/temperature_step1_<device>.json.

The temperature chart in Section 5.1 reads sample_output_en/temperature_step1_mps.json (October 2026, Mac Studio M3 Ultra, on "What is the capital of France? Answer in one sentence.", where "The" has raw logit 22.88 and probability 100.0% at temperature 0.5, 97.6% at temperature 1.0, and 72.5% at temperature 1.5 while "rest" rises from 0.4% at 1.0 to 21.1% at 1.5). To regenerate that file:

python temperature_step1.py --prompt "What is the capital of France? Answer in one sentence." --out sample_output_en

Running python temperature_step1.py with no arguments runs DEFAULT_PROMPTS[0] and writes to sample_output/temperature_step1_<device>.json (matching the pre-recorded Chinese step-1 file sample_output/temperature_step1_mps.json). In both cases, the temperature = 1.0 column matches step 1 of logprobs.py on the same prompt and device.