Text Generation: Logit Decoding, Argmax & Context Caching

Autoregressive generation operates as a closed feedback loop:

1. The Step-by-Step Cycle:

• Step 1 (Forward Pass): Feed current token sequence x_(1:T) into the transformer.
• Step 2 (Logit Slicing): Extract the logit vector for the final position z_T ∈ R^V.
• Step 3 (Token Selection): Apply sampling (greedy argmax, temperature, top-k, top-p) to choose new token x_(T+1).
• Step 4 (Sequence Concatenation): Append x_(T+1) to the prompt:

x_(1:T+1) = [x_(1:T), x_(T+1)]

• Step 5 (Iteration): Repeat until an `[EOS]` (end-of-sequence) token is emitted or maximum length T_(max) is reached.

1. The Quadratic Inference Bottleneck (O(T²)):

Without optimization, generating N tokens requires re-processing the entire growing prompt at every single step, incurring O(N²) attention FLOPs.

1. KV Cache Resolution:

Because past tokens never change in causal attention, storing past Key and Value vectors in a KV-Cache reduces per-step compute from O(T) down to O(1).

Every conversational AI interface—from ChatGPT to Claude—executes this exact loop during streaming responses. Tokens appear one by one on the user's screen as each iteration of the loop completes on the GPU.

Autoregressive generation originates in Claude Shannon's 1948 paper "A Mathematical Theory of Communication", where Shannon generated statistical text n-grams by hand using dictionary frequencies. Modern LLMs apply Shannon's exact autoregressive principle scaled across billions of neural parameters.

• Predict one, append, repeat: The three-step pulse of autoregressive generation.
• The future becomes the past: Yesterday's prediction is today's prompt context.

#### 1. Tensor Slicing Mechanics

The transformer forward pass processes batch X ∈ R^(B × T) and produces a 3D logit tensor:

Z = model(X) ∈ R^(B × T × V)

To predict the next token, all historical timesteps t ∈ 0, ṡ, T-2 are discarded; only the final timestep slice is extracted:

z_(next) = Z[:, -1, :] ∈ R^(B × V)

1. Temperature Scaling (T > 0):

pᵢ = (exp(zᵢ / T)) / (∑_(j=1)^V exp(zⱼ / T))

• T → 0: Distribution approaches a Dirac delta function ⟹ Greedy Argmax: x_(next) = argmaxᵢ zᵢ.
• T = 1.0: Standard unmodified model probabilities.
• T > 1.0: Flattens distribution toward uniform, boosting vocabulary diversity and creativity.

1. Top-K Truncation:

Retains only the K largest logits; all other indices are masked to -∞:

zᵢ' = zᵢ & if zᵢ ≥ z_((K)) \ -∞ & otherwise

1. Top-P (Nucleus) Truncation (Holtzman et al., 2019):

Sorts probabilities in descending order p_((1)) ≥ p_((2)) ≥ ṡ and identifies the smallest cutoff index k such that the cumulative mass reaches threshold p:

∑_(j=1)^k p_((j)) ≥ p

All tokens with rank > k are discarded, dynamically adjusting candidate set size based on model confidence.

Let z = [4.0, 2.0, 1.0] for tokens `['cat', 'dog', 'fish']`.

• At T = 1.0:

e⁴ ≈ 54.60, e² ≈ 7.39, e¹ ≈ 2.72 ⟹ ∑ = 64.71.

p = [54.60/64.71, 7.39/64.71, 2.72/64.71] = [0.844, 0.114, 0.042]

• At T = 0.5 (Sharpened):

z / 0.5 = [8.0, 4.0, 2.0].

e⁸ ≈ 2980.96, e⁴ ≈ 54.60, e² ≈ 7.39 ⟹ ∑ = 3042.95.

p = [2980.96/3042.95, 54.60/3042.95, 7.39/3042.95] = [0.980, 0.018, 0.002]

Notice that 'cat' rose from 84.4\% to 98.0\%, suppressing lower-confidence choices.

Production APIs combine these techniques in sequence: Logits → Temperature → Top-K → Top-P → Multinomial Sampling. For deterministic code generation and mathematical reasoning, T=0.0 (greedy) is preferred; for creative writing, T=0.7 with P=0.9 yields optimal natural fluency.

Temperature scaling originated in statistical physics (Boltzmann distribution) and was adapted to knowledge distillation by Geoffrey Hinton (2015). Ari Holtzman et al. introduced Nucleus (Top-P) sampling in 2019, solving the unnatural repetition artifacts of pure Top-K.

• Slice at -1: The last token's logits hold the key to the next token's future.
• Low temperature sharpens; high temperature scatters: T < 1 for logic; T > 1 for surprise.

The linear output projection head multiplies the final hidden states H ∈ R^B × T × d_(model) by the transposed embedding matrix W_(vocab) ∈ R^d_(model) × V:

Z = H W_(vocab) ∈ R^(B × T × V)

Every row z_(b, t, :) contains raw unconstrained logits spanning (-∞, +∞) for all tokens in the vocabulary.

In weight-tied architectures (such as GPT-2 and Gemma), W_(vocab) shares memory with the input token embedding table, saving hundreds of megabytes of GPU VRAM.

Introduced by Press and Wolf (2017) in 'Using the Output Embedding to Improve Language Models', proving that input and output word representations inhabit dual geometric vector spaces.

Dimension B is parallel batches; dimension T is sequence time; dimension V is the competitive arena of words.

Because causal attention prevents future tokens from affecting past tokens, the logit vector for the subsequent token is strictly:

z_(next) = Z[:, -1, :] ∈ R^(B × V)

Slicing at index T-1 extracts the only unconditioned prediction vector required for the next generation step.

This bottleneck inspired KV-caching. In production inference, earlier hidden states are stored in memory so that only a 1 × 1 × d query vector is computed at each step, slashing generation latency from O(T²) to O(T).

First highlighted in early transformer decoders (Vaswani et al., 2017) where parallel training computations contrast sharply with sequential inference bottlenecks.

All sequence history collapses into the final temporal slice: index -1 holds the key to the future.

Greedy decoding selects the index with the maximum scalar logit directly:

w_(next) = argmax_(i) zᵢ

Softmax converts logits into normalized probabilities:

P(wᵢ) = (e^(zᵢ)) / (∑_(j=1)^V e^(zⱼ))

Notice that because the exponential function is strictly monotonic, argmaxᵢ P(wᵢ) ≡ argmaxᵢ zᵢ, meaning softmax is computationally unnecessary if you only want the greedy token.

Greedy decoding is preferred in deterministic reasoning (code generation, math solving, SQL synthesis). Temperature-scaled sampling is preferred in creative writing and conversational chat.

Greedy decoding dates back to Viterbi search (1967). Its tendency to produce repetitive degenerative text in neural LMs was comprehensively documented by Holtzman et al. (2019).

Argmax takes the peak without looking; softmax maps the whole mountain into a rolling probability plain.

When input context length T_(cur) > T_(max) (e.g. 1024 tokens for GPT-2), positional embeddings cannot be indexed beyond T_(max)-1. The input tensor is cropped to keep only the most recent tokens:

x_(window) = x_T_(cur) - T_(max) : T_(cur)

Tokens prior to T_(cur) - T_(max) are completely evicted from the self-attention receptive field.

Sliding window attention in Mistral-7B and FlashAttention-2 uses this exact principle to achieve constant memory consumption during long-document processing.

Context limits originated from quadratic self-attention memory complexity O(T²). RoPE positional scaling (Su et al., 2021) and YaRN later extended effective windows from 2k to 128k+ tokens.

When memory is finite, the distant past must make way for the immediate present: crop from the left to preserve the right.

The core generation loop follows four distinct steps per iteration k ∈ [1, N]:

1. Context crop: x_(in) = x_(1:t)[-T_(max):]
2. Forward pass: Z = Model(x_(in))
3. Next token selection: w_(new) = argmax Z[-1, :]
4. Tensor concatenation: x_(1:t+1) = [x_(1:t), w_(new)]

This loop forms the foundational `generate()` method in the Hugging Face Transformers library and custom LLM inference runtimes.

Pioneered in Karpathy's `minGPT` and `nanoGPT` educational codebases as the clean mental blueprint of neural text generation.

Tokenize → Forward → Slice → Argmax → Concatenate → Repeat.

During training, dropout randomly zeroes out activations with probability p and scales survivors by (1) / (1-p):

y_(train) = (1) / (1-p) m \odot x, m ~ Bernoulli(1-p)

In evaluation mode (`model.eval()`), the dropout mask m is deactivated, setting y_(eval) = x identically. BatchNorm and LayerNorm running statistics are frozen simultaneously.

Critical in every deep learning production pipeline. Forgetting `model.eval()` introduces stochastic noise that degrades code generation and mathematical reasoning benchmarks by 15–30%.

Introduced by Srivastava et al. (2014) in 'Dropout: A Simple Way to Prevent Neural Networks from Overfitting', demonstrating that deterministic inference approximates an ensemble of thinned models.

Train with noise to build resilience; evaluate without noise to achieve precision.

The complete text-to-text pipeline requires three coordinate transformations:

1. Encode: S_(prompt) →(completion)

The standard API interface for OpenAI `/v1/chat/completions` and Anthropic `/v1/messages`, where tokenizers handle unicode byte round-tripping seamlessly.

BPE byte-level encoding was popularized by Sennrich et al. (2016) and adopted by Radford et al. in GPT-2 (2019) to eliminate out-of-vocabulary `<unk>` errors.

Tokens are the bridge between human language and linear algebra: encode to enter the model, decode to return to the human.

With random Gaussian weights W ~ N(0, σ²), the logits z have near-zero mean and equal variance. Softmax assigns a nearly uniform probability ≈ 1/50257 ≈ 0.0000199 to every word, producing high entropy H = -∑ p log p ≈ ln(50257) ≈ 10.82 nats.

Verifying that an untrained model outputs random text with initial loss ≈ 10.82 is the universal smoke test used by ML engineers before spending thousands of dollars on cloud training.

Highlighting initialization statistics dates back to Xavier Glorot and Kaiming He, whose variance-preserving initializations prevent exploding activations before training starts.

An architecture provides the capacity to learn; only optimization provides the knowledge to speak.

The autoregressive generation loop executes in four deterministic steps:

```python

def generate(model, input_ids, max_new_tokens, context_window=1024):

for _ in range(max_new_tokens):

# 1. Truncate context to maximum context length (sliding window)

idx_cond = input_ids[:, -context_window:]

# 2. Forward pass: compute logits for all tokens in context

logits = model(idx_cond) # Shape: (batch, seq_len, vocab_size)

# 3. Focus on the very last position logits: (batch, vocab_size)

next_token_logits = logits[:, -1, :]

# 4. Sample next token (greedy or multinomial)

next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)

# 5. Check for End-of-Sequence (EOS) token

if next_token.item() == EOS_ID:

break

# 6. Append new token to ongoing context

input_ids = torch.cat((input_ids, next_token), dim=1)

return input_ids

```

Computational complexity:

Generating N tokens without KV caching requires O(N²) FLOPs because the entire sequence is reprocessed at every step. With a KV cache, complexity drops to O(N) operations.

This sliding window loop is the engine powering ChatGPT, Claude, and every command-line LLM runner (`llama.cpp`, Ollama).

Autoregressive Markov models date back to Andrei Markov (1913). Radford et al. combined deep autoregressive transformers with massive web corpora to create GPT (2018).

• Today's output becomes tomorrow's input.
• Feed the prompt, grab the last logit, sample, append, repeat.

A projection architecture diagram in square 1:1 format on warm cream parchment (#fcf0db) with a sepia frame.

Top section: Title in navy serif: "From Hidden States to Vocabulary Logits". Subtitle: "The final linear projection W_(head) mapping representation space to token scores."

Middle section: Dimensional transformation flow:

1. Final Transformer Block Output: h_L ∈ R^(B × L × D) (Hidden state of the final token at position L).
2. Final LayerNorm: ĥ_L = LayerNorm(h_L) ∈ R^D.
3. Unembedding Matrix (LM Head): W_(head) ∈ R^(V × D) (where V = 50,257).
4. Compute Raw Logits: z = W_(head) · ĥ_L ∈ R^V.

Show vocabulary logit bar chart: A vector of 50,257 real numbers. Highlight highest logits corresponding to top candidate words ("learning", "processing", "networks").

Bottom section: Operational bridge: "Logits represent unnormalized log-probabilities: p = Softmax(z). Top-k, top-p, and temperature decoding operate directly on these scores."

Clean vector projection arrows, highlighted vocabulary bars, balanced square styling.

Every transformation preserves structural invariants while mapping coordinate spaces. When operating in high dimensions, numerical precision and computational efficiency require adhering strictly to canonical mathematical laws.

Formal derivation and peer-reviewed analysis:

Radford et al. (2019), Language Models are Unsupervised Multitask Learners, OpenAI ↗

"Generating Output via Logits: The Unembedding Projection: Invariance, symmetry, and structure govern mathematical truth across continuous and discrete domains."