# 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

```bash
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:

```bash
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.
