M3.8 — Output Head and Decoding Policy

0x00 Overview And Background

M3.7 ended with final normalized state rows:

H_final: [B, S, d_model]

Each row is a learned continuous representation. It is not yet:

a token ID
a vocabulary score vector
a probability distribution
a selected next token

The output head, also called the language-model head, generator, vocabulary projection, or unembedding, changes the final axis:

hidden width d_model
    -> output vocabulary width V_out

For all sequence rows:

[B, S, d_model] -> [B, S, V_out]

For one selected inference row:

[B, d_model] -> [B, V_out]

M3.8 owns both the projection and the local decoding policy that consumes its vocabulary scores. M3.10 owns the surrounding autoregressive state machine: appending IDs, updating cache state, stopping, and invoking the next model step.

Owner question

How does one final hidden row become a typed vocabulary distribution, and how do greedy, stochastic, truncated, and beam policies transform that distribution into a token decision without confusing model scores with generation state?

M3.8 owns:

final hidden rows -> output projection -> logits
logits -> processors -> warpers -> normalization -> selection
vocabulary identity and optional input/output weight tying
training loss compatibility at the logit boundary
local token or beam candidate metadata

It defers shifted targets and optimizer execution to M3.9, the prefix/cache loop to M3.10, serving scheduling to M5, and verified multi-token acceleration to M7.

---

Background knowledge: language modeling ends with an unembedding

Input embeddings map vocabulary addresses into model-width vectors. The output head performs the reverse type of interface: it compares a final hidden row with one learned direction per vocabulary item and produces one scalar score for each item.

hidden representation:
    [B,S,d_model]

vocabulary projection:
    [d_model] -> [V]

logits:
    unnormalized scores used by the loss or decoding policy

weight tying:
    optional reuse of the input embedding parameters as output directions

The head does not choose a token and softmax is not always part of the module itself. Training losses, generation APIs, and fused kernels may consume logits under different conventions, so naming the exact output object is part of correctness.

Why this interface became its own system

The oldest useful abstraction is multiclass classification at every target position. A recurrent or Transformer decoder produces a context-dependent state; a learned affine map assigns one energy to every output symbol; a normalized likelihood trains those energies against the observed next symbol. That view is still correct, but modern language models made each part large enough to become an engineering boundary:

learned representation:
    one model-width row summarizes the visible prefix

learned classifier:
    one output direction and optional bias per vocabulary address

training normalization:
    log-sum-exp over a vocabulary that may be sharded or fused with loss

generation policy:
    a stateful transformation of one score row before a token is selected

The vocabulary axis can contain tens or hundreds of thousands of rows. Its matrix may be tied to the input embedding, sharded across ranks, quantized differently from the trunk, or replaced by a hierarchical or clustered head. At inference time, the same learned logits can drive greedy selection, truncated sampling, adaptive feedback, constrained decoding, or search. The model checkpoint therefore does not define a complete generation behavior by itself.

This is also why the common phrase “the softmax layer generates the next token” is too compressed. Three different contracts are hidden inside it:

Contract	Learned or configured?	Output
vocabulary projection	learned checkpoint state	raw logits
normalization/loss	mathematical and runtime contract	probabilities, log-probabilities, or scalar loss
decoding policy	configured algorithm plus local state/RNG	token or candidate metadata

The original Transformer made the pre-softmax projection explicit and shared compatible embedding/projection weights. Later work studied weight tying as an architecture and optimization choice, while large-vocabulary work introduced sampled, hierarchical, and adaptive alternatives. Open-ended neural generation then exposed a second problem: maximizing likelihood during training does not imply that local likelihood maximization is the best generation policy. Nucleus, typical, eta, feedback-controlled, and contrastive policies are different responses to that inference problem, not different definitions of the learned head.

From residual coordinates to vocabulary evidence

The unembedding does more than resize a tensor. Each vocabulary row defines a direction against which the final residual state is scored:

$$z_v = \langle h, e^{\text{out}}_v \rangle + b_v$$

Consequently, every residual update can change the final logit row through the same output directions. This observation motivates output-embedding analyses and logit-lens diagnostics, but it does not make logits calibrated probabilities or prove that an intermediate residual row has the same geometry as the final normalized state.

The historical progression is therefore best read as a set of contracts:

fixed-vocabulary classifier
    -> optional input/output storage tying
        -> large-vocabulary training and sharding alternatives
            -> local generation policies over one score row
                -> runtime integration, evaluation, and rollback

M3.8 follows that progression. It does not turn the output article into the outer autoregressive loop, a serving scheduler, or a survey of speculative execution.

Immutable head-to-policy manifest

Shape compatibility is necessary but insufficient. Two artifacts can both produce a tensor shaped [B,S,V] while assigning different meanings to its rows, normalizing it differently, or consuming randomness in a different order. A reproducible experiment and a releasable runtime therefore pin an identity manifest before logits are compared.

Identity group	Fields that must be pinned	Failure if omitted
head	model width, output width, storage orientation, bias, parameter dtype, accumulation dtype	transposed or numerically different scores
vocabulary	tokenizer artifact hash, ordered ID-to-token table hash, special IDs, added-token order	correct numbers attached to wrong symbols
tying	tied/untied flag, source parameter path, alias/storage identity, tie restoration policy	duplicate parameters or silently broken sharing
loss	score type, target type, ignore index, label smoothing, reduction, vocabulary sharding	double normalization or wrong gradients
policy	ordered processors, ordered warpers, support fallback, selector, stop-token view	same logits produce different decisions
randomness	generator algorithm, seed, per-row stream/counter, draw schedule, distributed owner	irreproducible or batch-order-dependent samples
checkpoint/runtime	format/version, parameter keys, shard ranges, framework/engine version, kernel path	loader accepts an incompatible artifact

One compact manifest can be written as:

{
  "head": {
    "weight_key": "lm_head.weight",
    "stored_shape": ["V", "d_model"],
    "bias": false,
    "parameter_dtype": "bf16",
    "accumulation_dtype": "fp32"
  },
  "vocabulary": {
    "size": "V",
    "ordered_table_sha256": "...",
    "special_ids": {"bos": 1, "eos": 2, "pad": 0}
  },
  "tying": {
    "enabled": true,
    "embedding_key": "model.embed_tokens.weight",
    "checkpoint_alias_required": true
  },
  "loss": {
    "input": "raw_logits",
    "target": "class_index",
    "ignore_index": -100,
    "label_smoothing": 0.0,
    "reduction": "token_mean",
    "vocab_parallel": false
  },
  "policy": {
    "processors": ["repetition_penalty", "constraint_mask"],
    "warpers": ["temperature", "top_p"],
    "selector": "multinomial",
    "empty_support": "raise"
  },
  "rng": {
    "algorithm": "framework_generator",
    "seed": 1234,
    "stream_owner": "request",
    "draw_contract": "one categorical draw per active row"
  }
}

This is not a universal schema. It is a declaration of invariants. A production format may encode hashes, enums, and runtime keys differently, but it must answer the same questions.

Compatibility gates

Before loading or comparing two artifacts:

1. Compare ordered vocabulary identities, not vocabulary sizes alone.
2. Reconstruct parameter aliases and prove pointer/storage identity when tying is required.
3. Confirm whether the caller supplies raw logits, probabilities, or log-probabilities.
4. Compare processor and warper order, because most transforms do not commute.
5. Confirm the sampling generator and per-request draw schedule before calling two outputs deterministic equivalents.
6. For a vocabulary shard, verify global row offsets and the rank that owns each target row.
7. Refuse release when a checkpoint migration lacks a reversible vocabulary and optimizer-state mapping.

The manifest separates architecture identity from an experiment’s tunable policy parameters. Sweeping temperature or top-p does not change the checkpoint head. Reordering a constraint mask and top-p, changing token IDs, or breaking a tied alias changes the executable contract.

0x01 The local object: final normalized hidden rows

The output head begins after the final normalization expected by the checkpoint:

Transformer blocks
    -> final LayerNorm or RMSNorm
        -> H_final
            -> output head

For a decoder-only model:

H_final: [B, S, d_model]

For an encoder-decoder model, the generative head consumes decoder-owned target rows:

decoder H_final: [B, T, d_model]
encoder memory:  [B, S, d_model]

The encoder memory informs the decoder through cross-attention, but the target-vocabulary head is applied to decoder rows, not directly to encoder memory.

The local row equation is:

$$z_i = h_i W_U + b_U$$

where:

h_i: [d_model]
W_U: [d_model, V_out] in row-vector algebra
b_U: [V_out] optional
z_i: [V_out]

The vector z_i contains logits, one score for every output-vocabulary row.

---

0x02 Linear projection and storage orientation

A minimal raw-logit head is:

class LMHead(nn.Module):
    def __init__(self, d_model, vocab_size, bias=False):
        super().__init__()
        self.proj = nn.Linear(d_model, vocab_size, bias=bias)

    def forward(self, hidden_states):
        return self.proj(hidden_states)

Shape trace:

hidden_states: [B, S, d_model]
logits:        [B, S, V_out]

nn.Linear acts on the last dimension and preserves all leading dimensions. The same head can therefore process:

[d_model]       -> [V_out]
[B,d_model]     -> [B,V_out]
[B,S,d_model]   -> [B,S,V_out]

Algebra versus PyTorch weight shape

Row-vector algebra often writes:

W_U: [d_model, V_out]
z = h @ W_U + b

PyTorch stores:

nn.Linear(d_model, V_out).weight: [V_out, d_model]

and computes:

z = h @ weight.T + bias

This orientation matters for checkpoint conversion and weight tying. The PyTorch output weight has the same stored shape as a token embedding table:

embedding.weight: [V_out, d_model]
lm_head.weight:   [V_out, d_model]

It is not necessary to manually transpose the parameter before calling nn.Linear.

Logits are learned compatibility energies

Let $e_v$ denote output row $v$ in PyTorch storage. With no bias,

$$z_v=h^\top e_v.$$

The score grows when the final hidden direction aligns with the token’s output direction and with their magnitudes. This is a dot product, not necessarily cosine similarity because neither vector must be unit-normalized.

Softmax is translation invariant:

$$\operatorname{softmax}(z+c\mathbf1) =\operatorname{softmax}(z).$$

Only relative logits determine probabilities. Their pairwise difference is a multiclass log-probability ratio:

$$\log\frac{p_i}{p_j}=z_i-z_j.$$

That statement is more precise than calling each individual logit a binary log-odds value.

Cross-entropy gradient at the head

For a one-hot target $y$ and probability row $p=\operatorname{softmax}(z)$,

$$\mathcal L=-\sum_v y_v\log p_v.$$

Differentiate through log-softmax:

$$\frac{\partial\mathcal L}{\partial z_v}=p_v-y_v.$$

For a batch of flattened supervised rows $H\in\mathbb R^{R\times d}$ and output storage $E\in\mathbb R^{V\times d}$ with $Z=HE^\top$:

$$G_E=(P-Y)^\top H, \qquad G_H=(P-Y)E.$$

Every vocabulary row receives a gradient proportional to its predicted probability, while the target row receives the additional -1 term. This is why training does not need argmax and why the full vocabulary participates even though labels contain one integer per row.

With label smoothing target distribution $q$, the same result becomes

$$\frac{\partial\mathcal L}{\partial z}=p-q.$$

Deriving the complete affine-head backward

Flatten only the supervised row axes and retain vocabulary as the class axis:

$$H\in\mathbb R^{R\times d}, \qquad E\in\mathbb R^{V\times d}, \qquad b\in\mathbb R^V.$$

The stored-weight forward equation is

$$Z_{rv} =\sum_{j=1}^{d}H_{rj}E_{vj}+b_v.$$

Let $G_{rv}=\partial\mathcal L/\partial Z_{rv}$. A differential exposes every parameter path:

$$dZ_{rv} =\sum_j E_{vj}\,dH_{rj} +\sum_j H_{rj}\,dE_{vj} +db_v.$$

Collecting the coefficient of each differential gives:

$$\frac{\partial\mathcal L}{\partial H_{rj}} =\sum_v G_{rv}E_{vj},$$ $$\frac{\partial\mathcal L}{\partial E_{vj}} =\sum_r G_{rv}H_{rj},$$ $$\frac{\partial\mathcal L}{\partial b_v} =\sum_r G_{rv}.$$

In matrix form:

dH = G @ E             [R,d]
dE = G.T @ H           [V,d]
db = sum_rows(G)        [V]

For ordinary one-hot cross-entropy, G = P - Y. The target row receives the negative correction, while every non-target row receives probability-weighted pressure. A vocabulary projection is therefore dense in output-gradient space even when the input lookup that produced the hidden states touched only a small set of token addresses.

What the bias can represent

The bias adds one context-independent scalar per output address:

$$z_v=h^\top e_v+b_v.$$

If $h=0$, the distribution is $\operatorname{softmax}(b)$; thus $b$ can encode a global token preference. But calling it “the unigram prior” is too strong: training jointly changes $E$, the trunk, and $b$, while softmax translation invariance leaves one common-offset degree of freedom. Bias-free heads force all context-independent preference to emerge through the hidden state and output directions.

Norm and angle are separate score mechanisms

Write the dot product as

$$z_v-b_v =\lVert h\rVert_2\lVert e_v\rVert_2\cos\theta_v.$$

A large score can come from angular alignment, output-row norm, hidden-state norm, or all three. Weight tying does not turn this into nearest-neighbor cosine search. Final normalization constrains the hidden-state scale according to its own contract, but output-row norms remain learned unless an architecture normalizes or parameterizes them explicitly.

Parameter and logical-work ledger

For an untied affine head:

$$P_{\text{head}}=Vd+\mathbf 1_{\text{bias}}V.$$

For $R$ projected rows, the matrix product performs $RVd$ multiply-accumulates, and the logical logit tensor contains $RV$ elements. A tied head removes the additional $Vd$ output parameter allocation; it does not remove the matrix product or guarantee that a runtime stores only one physical copy after conversion.

M3.9 derives how shifted labels and masks create those target rows; M3.8 owns the typed gradient boundary at the output head.

---

0x03 The A15 Harvard Generator

A15 imports the Annotated Transformer style, where the generator combines the linear projection with log_softmax:

import torch.nn.functional as F

class Generator(nn.Module):
    def __init__(self, d_model, vocab):
        super(Generator, self).__init__()
        self.proj = nn.Linear(d_model, vocab)

    def forward(self, x):
        return F.log_softmax(self.proj(x), dim=-1)

The shape path is:

x:            [B,T,d_model] or [B,d_model]
self.proj(x): [B,T,V_tgt]   or [B,V_tgt]     raw logits
generator(x):[B,T,V_tgt]   or [B,V_tgt]     log-probabilities

The final shape does not identify the score type. Both logits and log-probabilities use the same vocabulary-width tensor.

Source-audit corrections

The A15 local code is useful, but several surrounding explanations should not be imported as technical claims.

A15 wording or implication	Correct contract
Final decoder state is a probability vector.	It is a hidden representation in d_model space.
Each logit is a probability or ordinary binary log-odds.	Multiclass logits are unnormalized real scores.
Softmax selects the highest-scoring token.	Softmax normalizes scores; a separate policy selects a token.
Training finds an argmax token at every row.	Training normally computes a differentiable loss against known labels; argmax is not in that gradient path.
The toy tensor named prob contains probabilities.	This Generator returns log-probabilities.
generator.lut.weight is the output projection parameter.	In the shown class it is generator.proj.weight.

The source snippet also imports torch.nn.functional as F but prints a bare log_softmax(...) call. The executable form is F.log_softmax(...) unless log_softmax was separately imported.

---

0x04 Logits, probabilities, and log-probabilities

Let one output row be:

logits z: [V_out]

Raw logits

logits = lm_head(hidden)

Logits are unconstrained real scores:

they can be positive or negative
they do not need to sum to one
they are not yet sampling weights

For multiclass output, calling each logit a binary log-odds value is generally misleading. The distribution depends on all vocabulary scores through the normalization denominator.

Probabilities

$$p_c = \frac{\exp(z_c)}{\sum_{j=1}^{V}\exp(z_j)}$$

probs = torch.softmax(logits, dim=-1)

Now:

probs[c] >= 0
sum_c probs[c] = 1

Log-probabilities

$$\log p_c = z_c - \log\sum_{j=1}^{V}\exp(z_j)$$

log_probs = F.log_softmax(logits, dim=-1)

PyTorch 2.13 documents F.log_softmax as mathematically equivalent to log(softmax(x)) but implemented with a numerically stable formulation. Do not compute the two operations separately merely to obtain log-probabilities.

One denominator derives softmax, log-softmax, and cross-entropy

The naive normalizer

$$\log\sum_v e^{z_v}$$

can overflow even when the final probabilities are ordinary. Let $m=\max_v z_v$. Multiplying numerator and denominator by $e^{-m}$ does not change softmax:

$$\frac{e^{z_v}}{\sum_j e^{z_j}} = \frac{e^{z_v-m}}{\sum_j e^{z_j-m}}.$$

The largest shifted logit is zero, so every exponential lies in $(0,1]$. The stable log-normalizer is:

$$\operatorname{LSE}(z) =m+\log\sum_j e^{z_j-m}.$$

That one scalar derives both normalized forms:

$$\log p_v=z_v-\operatorname{LSE}(z), \qquad p_v=e^{z_v-\operatorname{LSE}(z)}.$$

For class target $y$, negative log-likelihood cancels every irrelevant intermediate:

$$\mathcal L =-\log p_y =\operatorname{LSE}(z)-z_y.$$

An implementation can therefore compute loss without materializing a persistent probability tensor. A fused implementation may keep only the row maximum, denominator, target score, and whichever state its backward needs.

def stable_cross_entropy_row(logits, target):
    # Accumulate the reduction in the declared stable dtype.
    maximum = logits.max()
    shifted = logits - maximum
    log_denom = maximum + log(exp(shifted).sum())
    return log_denom - logits[target]

Why the gradient is probability minus target

Differentiate the normalizer:

$$\frac{\partial\operatorname{LSE}(z)}{\partial z_k} = \frac{e^{z_k}}{\sum_j e^{z_j}} =p_k.$$

The target term contributes $-1$ only at $k=y$:

$$\frac{\partial\mathcal L}{\partial z_k} =p_k-\mathbf1[k=y].$$

The Hessian is:

$$\nabla_z^2\mathcal L =\operatorname{diag}(p)-pp^\top.$$

It is positive semidefinite and has a zero direction along the all-ones vector. That zero mode is the differential form of translation invariance: adding the same constant to every logit changes neither probabilities nor cross-entropy.

Label smoothing is a target-distribution contract

For the convention that mixes the one-hot target with a uniform distribution over all V classes:

$$q_v=(1-\varepsilon)\mathbf1[v=y]+\frac{\varepsilon}{V},$$ $$\mathcal L_{\text{smooth}} =-\sum_v q_v\log p_v =(1-\varepsilon)(-\log p_y) +\varepsilon\left(-\frac1V\sum_v\log p_v\right),$$

and:

$$\frac{\partial\mathcal L_{\text{smooth}}}{\partial z} =p-q.$$

Another legitimate convention distributes all smoothing mass over the V-1 non-target classes. The same epsilon then produces a different target distribution and gradient. A manifest must pin the convention, not only the number.

Dtype and degenerate-row policy

Stable algebra does not choose an accumulation dtype. A runtime may project in bf16/fp16, promote row maximum and exponential reductions to fp32, and cast the loss or gradient according to its mixed-precision contract. Merely subtracting the maximum does not prevent underflow in tiny tail probabilities; underflowed tail mass can be acceptable for a loss kernel only if the target and reduction remain valid.

Special rows require explicit behavior:

Input row	Mathematical result	Runtime policy
one finite maximum, other very negative logits	concentrated valid distribution	stable max subtraction
several equal maxima	mass divided across tied maxima	tie order matters only to discrete selection
all logits negative infinity	undefined zero-over-zero normalization	reject or provide declared legal fallback
any NaN	NaN contaminates comparisons/reduction	fail telemetry; do not silently sample
masked target	loss is invalid for that row	M3.9 excludes it through the declared ignore/mask contract

Temperature belongs before the denominator:

$$p_T=\operatorname{softmax}(z/T).$$

If a training loss deliberately uses this transformed row, its logit gradient contains a factor $1/T$. Ordinary next-token training uses the checkpoint’s declared objective, usually at $T=1$; inference temperature is a policy and does not retroactively change the trained loss.

Transform relationships

probs = log_probs.exp()
probs = torch.softmax(log_probs, dim=-1)

Both recover the probability distribution mathematically because normalized log-probabilities have logsumexp(log_probs) = 0. Applying softmax to log-probabilities is therefore redundant, not a distribution-changing second softmax.

In contrast:

torch.softmax(probs, dim=-1)

does change an already normalized probability distribution and is usually a type error.

Argmax invariance

Because softmax, log_softmax, and exp are monotone in each score while sharing one row-wise denominator:

argmax(logits) == argmax(probs) == argmax(log_probs)

That equivalence is useful for greedy selection, but it does not make the tensor types interchangeable for loss functions, cumulative-probability filtering, or multinomial sampling.

---

0x05 The training interface: all supervised rows

Under teacher forcing, one model pass can produce one hidden row for every target input position:

H_final: [B, T, d_model]

The output head can project all rows in parallel:

scores: [B, T, V_tgt]

The label path supplies one target token ID per supervised position:

labels: [B, T]

M3.9 owns target shifting, padding-label exclusion, flattening, label smoothing, and loss normalization. M3.8 owns the score-type interface between the head and criterion.

Harvard path: log-probabilities into KL divergence

log_probs = generator(hidden_states)       # [B,T,V]

flat_log_probs = log_probs.reshape(B * T, V)
flat_targets = target_distribution.reshape(B * T, V)

loss = F.kl_div(
    flat_log_probs,
    flat_targets,
    reduction="sum",
)

The Annotated Transformer label-smoothing path constructs a target distribution and expects generator log-probabilities.

Modern path: raw logits into cross entropy

logits = lm_head(hidden_states)             # [B,T,V]

flat_logits = logits.reshape(B * T, V)
flat_labels = labels.reshape(B * T)

loss = F.cross_entropy(
    flat_logits,
    flat_labels,
    ignore_index=pad_or_ignore_id,
)

PyTorch 2.13 CrossEntropyLoss expects unnormalized logits. It incorporates the log-softmax/NLL-style calculation for class-index targets.

The implementation rule is:

Generator returns log-probs:
    pair with an NLL/KL-style criterion expecting log-probs

LM head returns raw logits:
    pair with cross entropy or explicitly convert as required

Training does not need to select an argmax token before computing the loss. The known label indexes the appropriate class term while gradients flow through the entire vocabulary score row.

Source-to-tensor training trace

M3.9 supplies an immutable row map:

row r
    -> batch index
    -> sequence position
    -> target token ID or ignore

M3.8 consumes only supervised rows:

H_valid: [R,d]
E:       [V,d]
b:       [V] optional

Z = H_valid @ E.T + b             [R,V]
m = row_max(Z)                    [R]
s = sum_v exp(Z - m)              [R]
t = gather_target_logit(Z, y)     [R]
loss_row = m + log(s) - t         [R]

The reduction then receives the declared valid-row count from M3.9. A token mean is:

$$\mathcal L_{\text{mean}} = \frac{\sum_{r\in\mathcal V}\mathcal L_r}{|\mathcal V|},$$

where $\mathcal V$ is the supervised-row set. Averaging each local chunk or rank independently and then averaging those means is wrong when valid counts differ.

Materialized, chunked, and fused execution

The same mathematical objective admits several schedules:

Schedule	Persistent logical score state	Exact objective?	Main tradeoff
materialized	full [R,V] logits, possibly probabilities	yes	simple APIs, largest live score tensor
row-chunked	[C,V] for chunk size C	yes	repeated launches, bounded live score state
fused linear plus cross-entropy	tiled logits/reduction state	yes when contract matches	specialized backward and kernel constraints
sampled/clustered objective	subset or hierarchy	generally a different estimator/head	changes training contract

Let $b_z$ be bytes per stored logit. Full materialization requires:

$$M_{\text{logits,logical}}=RVb_z.$$

Row chunking bounds this object by:

$$M_{\text{chunk,logical}}=CVb_z, \qquad C<R.$$

These are logical tensor sizes, not hardware peak claims. M4 combines them with allocator behavior, kernel workspaces, recomputation, and the rest of the model.

An exact row-chunked loss must accumulate a sum and count:

loss_sum = 0.0
valid_count = 0

for hidden_chunk, labels_chunk in valid_row_chunks:
    logits = linear(hidden_chunk, output_weight, output_bias)
    losses = cross_entropy(logits, labels_chunk, reduction="none")
    loss_sum += losses.sum()
    valid_count += labels_chunk.numel()

loss = loss_sum / valid_count

Do not call mean inside each differently sized chunk and average chunk means.

Backward state and recomputation

The exact backward needs the local quantity:

$$G_Z=P-Y$$

scaled by the outer loss reduction. A schedule can obtain it by saving probabilities, saving enough row-normalizer state and recomputing exponentials, or recomputing the projection and loss tile. It then accumulates:

dH += G_Z @ E
dE += G_Z.T @ H
db += row_sum(G_Z)

Strategy	Saved state	Recomputed state	Contract risk
ordinary eager autograd	framework-selected inputs/intermediates	implementation-dependent	accidental full-logit liveness
checkpointed head	hidden or upstream recompute inputs, labels	projection and normalization	replay must preserve dtype and row map
fused loss	hidden tile, weight, target, reduction/normalizer metadata	probability-like tile in backward	kernel convention must match reference
chunked loss	chunk inputs or replay source, targets	each score chunk	sum/count and gradient accumulation order

Recomputation is valid only if it recreates the same hidden rows, weights, targets, masking, and numerical policy. If upstream dropout is replayed, its RNG state belongs to the M3.9 training-execution contract. Inference sampling RNG is a different state and must never be consumed by the training loss.

Exact conformance checks

A candidate implementation should be tested against a small materialized reference for:

1. per-row loss before reduction;
2. valid-row sum, count, and final reduction;
3. gradients for hidden rows, output weight, and optional bias;
4. tied-weight accumulation when the embedding path is present;
5. extreme logits, ignored rows, label smoothing, and uneven chunks;
6. forward and backward dtypes within declared tolerances.

Matching only the scalar mean loss can hide row swaps, wrong valid counts, and compensating errors in gradients.

---

0x06 The inference interface: project the decision row

At one generation step, the policy needs one vocabulary score row per active sequence:

decision hidden state: [B,d_model]
decision logits:       [B,V_out]

In a simple unpadded full-prefix example:

last_hidden = hidden_states[:, -1, :]  # [B,d_model]
logits = lm_head(last_hidden)          # [B,V]

But hidden_states[:, -1, :] means the last rectangular column, not necessarily the last semantic token. With right padding, use the row selected by the inference contract:

lengths = attention_mask.sum(dim=-1)               # [B]
last_idx = lengths - 1                              # [B]
batch_idx = torch.arange(B, device=hidden_states.device)
last_hidden = hidden_states[batch_idx, last_idx, :] # [B,d_model]
logits = lm_head(last_hidden)                       # [B,V]

M3.10 owns the padding-side and last-valid-row logic. M3.8 requires that the head receive the correct row.

During cached one-token decode, the model may already return:

hidden_states: [B,1,d_model]

Then:

logits = lm_head(hidden_states)  # [B,1,V]
step_logits = logits[:, -1, :]   # [B,V]

The output head does not append the selected token, update the cache, stop completed rows, or choose a decoding policy.

---

0x07 A versioned raw-logit path: LlamaForCausalLM

Hugging Face Transformers v5.13.1 LlamaForCausalLM separates the base model from a bias-free vocabulary projection:

self.model = LlamaModel(config)
self.vocab_size = config.vocab_size
self.lm_head = nn.Linear(
    config.hidden_size,
    config.vocab_size,
    bias=False,
)

Its forward path obtains final normalized hidden states from LlamaModel, selects the requested sequence rows, and returns raw logits:

hidden_states = outputs.last_hidden_state

slice_indices = (
    slice(-logits_to_keep, None)
    if isinstance(logits_to_keep, int)
    else logits_to_keep
)

logits = self.lm_head(hidden_states[:, slice_indices, :])

Shape examples:

logits_to_keep = 0:
    selected hidden: [B,S,d_model]
    logits:          [B,S,V]

logits_to_keep = 1:
    selected hidden: [B,1,d_model]
    logits:          [B,1,V]

Here -0 is 0, so slice(-0, None) means all sequence positions. The option changes how many rows receive the vocabulary projection; it does not change vocabulary size or the meaning of each output column.

This modern interface returns logits because training loss, generation processors, temperature, and selection code can consume raw scores without forcing one normalization policy inside the model head.

---

0x08 The vocabulary-axis identity contract

The last axis is not an anonymous class dimension. Output column c must correspond to the checkpoint’s token ID c:

logits[..., c]
    -> score for output token ID c
    -> tokenizer/model vocabulary entry for ID c

The minimum contract is:

0 <= emitted tokenizer ID < V_out
output row ordering matches checkpoint token-ID ordering
special token IDs agree with model configuration
added token rows exist and were initialized/adapted as intended

tokenizer.vocab_size, len(tokenizer), embedding rows, and output rows are not guaranteed to be reported identically in every library because added and reserved tokens can be handled separately. The robust audit prints all of them:

print(tokenizer.vocab_size)
print(len(tokenizer))
print(model.get_input_embeddings().weight.shape)
print(model.get_output_embeddings().weight.shape)
print(model.config.vocab_size)

For encoder-decoder models:

source IDs address V_src
target decoder IDs address V_tgt
output logits usually use V_tgt

Equal row counts do not prove equal vocabulary semantics. Swapping two tokenizer IDs preserves every shape and assigns scores to the wrong tokens.

Vocabulary migration is a row-address transaction

A safe migration starts with an explicit partial bijection:

$$\pi:\mathcal V_{\text{old}}\supseteq\mathcal K \rightarrow \mathcal V_{\text{new}},$$

where $\mathcal K$ is the set of retained old token IDs. For every retained token:

$$E^{new}_{\pi(v)}=E^{old}_v, \qquad b^{new}_{\pi(v)}=b^{old}_v.$$

Untied models apply the same semantic mapping independently to input embedding rows and output-head rows. Tied models migrate one shared table and restore one alias after loading. Optimizer moments and any row-shaped master weights must follow the same address transform:

$$m^{new}_{\pi(v)}=m^{old}_v, \qquad u^{new}_{\pi(v)}=u^{old}_v.$$

New rows need declared initialization. Common choices include random checkpoint-style initialization, an average of source-subtoken rows, copying a semantically related row, or a separately trained initialization. These are design choices, not equivalent migrations. Adam-like moments for genuinely new rows are usually initialized to zero under the chosen optimizer contract.

def migrate_vocab_rows(old_rows, old_to_new, new_size, init_new):
    new_rows = init_new(new_size, old_rows.shape[1])
    copied = zeros(new_size, dtype=bool)

    for old_id, new_id in old_to_new.items():
        new_rows[new_id] = old_rows[old_id]
        copied[new_id] = True

    return new_rows, copied

Append-only growth is safer, not free

Appending tokens while preserving every old ID avoids row permutation:

old IDs: 0 ... V_old-1
new IDs: V_old ... V_new-1

The old raw logits can remain identical for a fixed hidden row if old weights and biases are unchanged. The old probabilities do not necessarily remain identical because new rows enter the softmax denominator:

$$p^{new}_v = \frac{e^{z_v}} {\sum_{j<V_{old}}e^{z_j}+\sum_{j\ge V_{old}}e^{z_j}}.$$

Badly initialized new rows can therefore steal probability mass before they have useful semantics. Append-only compatibility tests must compare mapped raw logits and separately measure distribution and generation drift.

Tied alias and optimizer restoration order

One safe construction order is:

1. Load and hash the old tokenizer, ordered table, config, checkpoint, and optimizer manifest.
2. Build the new ordered vocabulary and immutable old-to-new mapping.
3. Allocate resized embedding/head/bias/master/moment tensors.
4. Transform every row-addressed object with the same mapping.
5. Load module state and restore the required tied Parameter alias.
6. Restore or create exactly one optimizer state for each shared Parameter.
7. Assert shapes, token hashes, special IDs, pointer identity, and shard ranges.
8. Run old-token raw-logit, new-token, loss, and generation conformance tests.
9. Release tokenizer, model, optimizer, policy, and mapping manifests as one versioned bundle.

Creating the optimizer before replacing parameters can leave stale objects in its parameter groups. Re-tying after optimizer restore can leave duplicate moment slots. The loader must prove both model alias identity and optimizer state ownership.

Release and rollback ledger

Artifact	Release gate	Rollback requirement
tokenizer and ordered table	deterministic ID table hash	previous immutable tokenizer
model rows and bias	mapped raw-logit oracle, new-row initialization audit	previous checkpoint
tied aliases	storage/pointer identity assertion	previous architecture config
optimizer/master state	mapped old rows, zero/declared new rows, one state per shared parameter	previous optimizer checkpoint
shards	contiguous global row ranges and owner map	previous shard manifest
local policy	special IDs and legal-token masks match new vocabulary	previous policy manifest

Rollback after further training is not “apply the inverse permutation”: new rows may have influenced normalization and all model gradients. Reliable rollback restores the complete previous release bundle.

---

0x09 Weight tying: shared storage, not similar shape

An input embedding table and PyTorch output head can have the same stored shape:

embed_tokens.weight: [V,d_model]
lm_head.weight:      [V,d_model]

Untied model:

two separate Parameter objects
independent values and gradients

Tied model:

one shared Parameter object
used for input lookup and output projection

Conceptually:

hidden = embed_tokens(input_ids)       # row lookup from E
logits = F.linear(hidden, E)           # hidden @ E.T

An explicit PyTorch tie is:

lm_head.weight = embed_tokens.weight

assert lm_head.weight is embed_tokens.weight
assert lm_head.weight.data_ptr() == embed_tokens.weight.data_ptr()

Copying values is not tying:

lm_head.weight.data.copy_(embed_tokens.weight.data)

After the next optimizer update, two copied Parameters can diverge.

What tying does and does not imply

The original Transformer shared compatible source embeddings, target embeddings, and the pre-softmax matrix under its vocabulary setup. Press and Wolf studied the output matrix as an embedding and reported reduced perplexity in their evaluated language models plus substantial parameter reduction in translation models.

Those results support weight tying as a real architecture choice, not a universal requirement:

same vocabulary semantics and dimensions:
    tying may be valid

different source and target vocabularies:
    blind three-way tying is invalid

same shapes but checkpoint trained untied:
    tying changes the model

The Transformers v5.13.1 LlamaConfig default sets tie_word_embeddings=False, even though the framework exposes tied-weight metadata for configurations/checkpoints that request tying. Read the checkpoint configuration and verify actual parameter identity rather than inferring from class names.

If the output head has a bias, tying the weight does not automatically tie or remove that separate [V] parameter.

Tied parameters receive two semantic gradient contributions

If one parameter $E$ is used for both input lookup and output projection, autograd accumulates both uses:

$$\frac{\partial\mathcal L}{\partial E} =\left.\frac{\partial\mathcal L}{\partial E}\right|_{input\ lookup} +\left.\frac{\partial\mathcal L}{\partial E}\right|_{output\ projection}.$$

The input contribution is sparse in token addresses for an ordinary embedding lookup before optimizer effects; the output contribution is generally dense across vocabulary rows because softmax assigns probability to all candidates. Tying is therefore not only parameter-count reduction. It couples how token rows are learned as inputs and as output directions.

For token-ID matrix $X$ and upstream gradient $G^{in}$ at the embedding output, the input-use contribution for row $v$ is a scatter-add:

$$\left.\frac{\partial\mathcal L}{\partial E_v}\right|_{\text{input}} = \sum_{(r,t):X_{rt}=v}G^{in}_{rt}.$$

The output-use contribution from flattened hidden rows $H$ is:

$$\left.\frac{\partial\mathcal L}{\partial E_v}\right|_{\text{output}} = \sum_r(P_{rv}-Y_{rv})H_r.$$

Autograd must add those expressions into one gradient buffer before one optimizer state is updated:

shared row E[v]
    <- sparse scatter from every input occurrence of token v
    +  dense classifier contribution for output candidate v

This also explains a migration hazard. If a loader materializes two Parameter objects, loads equal values, and creates separate optimizer slots, the next training step implements two independent equations rather than the tied one.

Tying variants are different architectures

Variant	Shared storage	Extra transform	Compatibility condition
untied	none	optional bias	input and output widths may differ
direct tying	input table equals output weight	none	same ordered vocabulary and model width
projected tying	shared table plus learned projection	maps hidden/output spaces	projection and checkpoint keys must be pinned
source-target-output tying	three uses share one table	task-dependent	source and target token identities must truly match
copied initialization	equal initial values only	none	not weight tying after optimization

Press and Wolf and Inan et al. motivate direct tying from related but distinct analyses. Their reported empirical results do not imply that every tokenizer, translation pair, or checkpoint should be tied. The architecture choice must be declared before training and reconstructed during load.

Distributed and sharded optimizers must preserve that identity. Registering two independently sharded parameters that happen to begin with equal values does not reproduce one shared parameter with one accumulated gradient and one optimizer state.

---

0x0A Output-row scope is an interface and cost boundary

The linear head can project any number of selected rows:

training all rows:
    [B,T,d_model] -> [B,T,V]

prefill last-valid rows only:
    [B,d_model] -> [B,V]

cached decode row:
    [B,1,d_model] -> [B,1,V]

Projecting all prompt rows during inference is mathematically valid if those rows are available, but often unnecessary when only the final row drives the next decision. Selecting hidden rows before the large vocabulary projection can avoid producing unused [B,S,V] tensors.

Let R be the number of hidden rows actually projected. The local ledger is:

$$N_{param,head}=Vd_{model}$$

plus V parameters if a bias exists, and:

$$\operatorname{MACs}_{head}=RVd_{model}.$$

The produced score tensor contains R*V elements. Therefore selecting rows before the projection reduces both projection work and output materialization in direct proportion to the number of omitted rows.

Worked example:

B       = 8
S       = 2048
d_model = 4096
V       = 32000
dtype   = bf16/fp16 for stored weights/logits

output-head weights:
    32000 * 4096
    = 131,072,000 parameters
    = 250 MiB at 2 bytes/parameter

all prompt-row logits:
    8 * 2048 * 32000
    = 524,288,000 elements
    = 1000 MiB at 2 bytes/element

last-row logits only:
    8 * 32000
    = 256,000 elements
    ~= 0.488 MiB

all-row projection:
    2,147,483,648,000 MACs

last-row projection:
    1,048,576,000 MACs

The ratio is exactly S=2048 for this row-selection comparison. Training needs all supervised rows, so it cannot apply the same shortcut blindly. Cached generation normally needs only the current decision row. M4.1 owns the whole-model and hardware traffic model, but this local derivation belongs here because it explains why the output-row interface is operationally important.

The [B,S,V] tensor is the logical training interface, not proof that an optimized loss kernel writes the entire tensor to HBM. A fused linear plus cross-entropy implementation can tile vocabulary work and combine projection, normalization, and loss reduction. That changes physical materialization, not the output-head class domain or the B*S*V*d_model projection arithmetic being audited here.

Vocabulary-parallel projection and loss

When $V$ is large, tensor parallelism can shard output rows. Rank $r$ owns

E_r: [V_r, d_model]
Z_r = H @ E_r.T: [R, V_r]

The global softmax denominator still spans all vocabulary shards. A numerically stable distributed log-sum-exp uses a global row maximum:

$$m=\max_r\max_{v\in\mathcal V_r}z_v,$$ $$\log\sum_{v=1}^{V}e^{z_v} =m+\log\sum_r\sum_{v\in\mathcal V_r}e^{z_v-m}.$$

This requires a max reduction followed by a sum reduction unless the backend fuses or restructures them. The target logit is local on exactly one vocabulary shard and must be selected and reduced consistently:

owner rank:
    contributes z_target

non-owner ranks:
    contribute neutral value under the chosen reduction

The per-row cross entropy is then

$$\mathcal L =-z_{target} +\log\sum_v e^{z_v}.$$

A vocabulary-parallel fused cross-entropy can compute this without all-gathering [R,V] logits. It still computes each local shard’s projection and participates in the global reductions. “Logits were not materialized globally” does not mean the vocabulary normalization disappeared.

Exact forward trace over vocabulary shards

Assume contiguous shard $r$ owns global IDs:

$$\mathcal V_r=[a_r,b_r), \qquad V_r=b_r-a_r.$$

For each supervised row:

1. Compute local logits Z_r = H E_r^T in the declared projection dtype.
2. Compute local maximum $m_r$ and all-reduce MAX to obtain $m$.
3. Form local exponentials $X_r=\exp(Z_r-m)$.
4. Sum each row locally and all-reduce SUM to obtain $s=\sum_r\sum_{v\in\mathcal V_r}X_{rv}$.
5. Determine owner mask $a_r\le y<b_r$ and local target index $y-a_r$.
6. The owner contributes its target logit; non-owners contribute zero. An all-reduce SUM yields the global target score.
7. Return $m+\log s-z_y$ before the declared outer reduction.

local_logits = linear(hidden, local_vocab_weight)

global_max = all_reduce_max(local_logits.max(dim=-1))
local_exp = exp(local_logits - global_max[:, None])
global_denom = all_reduce_sum(local_exp.sum(dim=-1))

owns_target = (shard_start <= target) & (target < shard_end)
# Non-owners gather a safe in-range index; where() discards that value.
local_target = where(owns_target, target - shard_start, 0)
target_contribution = where(
    owns_target,
    gather(local_logits, local_target),
    0.0,
)
global_target = all_reduce_sum(target_contribution)

loss_row = global_max + log(global_denom) - global_target

Using negative infinity as the non-owner value would be correct for a maximum reduction, but not for the SUM shown above. Neutral values belong to the chosen collective.

Exact backward and ownership

Each rank can normalize its local shard with the global denominator:

$$P_r=\frac{\exp(Z_r-m)}{s}.$$

For one-hot class targets:

$$(G_{Z_r})_v = (P_r)_v-\mathbf1[r\text{ owns }y]\mathbf1[v=y-a_r].$$

No rank subtracts the target outside its own shard. Local parameter gradients need no vocabulary all-gather:

$$G_{E_r}=G_{Z_r}^\top H.$$

Every shard contributes to the hidden-state gradient:

$$G_H^{(r)}=G_{Z_r}E_r, \qquad G_H=\sum_rG_H^{(r)}.$$

When $H$ is replicated across tensor-parallel ranks, that sum is a collective reduction. Other layouts can use reduce-scatter or fused collectives, but they must implement the same global derivative.

Label smoothing adds a distributed target value for every local class. For the all-class convention:

$$(Q_r)_v =(1-\varepsilon)\mathbf1[r\text{ owns }y]\mathbf1[v=y-a_r] +\frac{\varepsilon}{V},$$

and G_Zr = P_r - Q_r. A non-target-only convention uses a different formula. Mixing the convention across ranks breaks both the loss and gradient.

Collective and failure ledger

Object	Local shape	Global operation	Silent failure
local maximum	[R]	MAX	normalization differs by shard
exponential sum	[R]	SUM	local probabilities sum to more than one globally
target contribution	[R]	SUM with zero non-owner	wrong or missing target score
hidden gradient contribution	[R,d]	SUM/reduce-scatter by layout	trunk sees only one vocabulary shard
weight gradient	[V_r,d]	local until optimizer/sharding contract	row owners or optimizer shards disagree

Telemetry should include target-owner count, target IDs outside all shard ranges, non-finite local maxima, global denominator validity, and per-rank collective participation. One rank skipping an ignored row while others enter its collectives can deadlock; ignored-row compaction or masking must preserve a common distributed schedule.

Distributed inference selection

Greedy selection over vocabulary shards needs the global maximum and its global token ID. Each rank can propose (local_max, local_id); a collective maximum with deterministic tie semantics chooses the winner.

Top-k can similarly merge local candidate lists:

each rank:
    local top-k over V_r

collective or coordinator:
    top-k over at most TP*k candidates

Top-p is harder to distribute exactly because cumulative mass depends on global ordering and the global softmax denominator. Implementations may gather scores, use distributed sort/selection, overselect local candidates with an exact correction, or impose approximations. A local top-p on each shard is not equal to global top-p.

For exact global top-k, each shard’s local top-k is sufficient: a candidate ranked below k inside its own shard cannot enter the global top-k because at least k same-shard candidates already outrank it. Merge candidates using global token IDs, and pin a deterministic secondary key for equal scores.

Global multinomial sampling also needs one declared owner. Gathering the full row is simple but expensive; distributed categorical algorithms can avoid that gather by sampling from shard masses and then a local conditional distribution. They must consume randomness according to the same manifest contract and return the same global token-address semantics.

Output-head quantization and tying boundary

The head can be quantized independently, but several constraints interact:

tied input/output weight:
    one physical parameter may serve lookup and GEMM paths with different kernel needs

vocabulary rows:
    per-row scales align naturally with output classes but not every embedding lookup kernel

rare or added tokens:
    outlier rows can be disproportionately affected by shared-group quantization

fused loss:
    dequantization and projection must match the loss kernel's accepted layout

M8 owns the compression decision. M3.8 records why the output head is not just another anonymous linear layer: its row identity is externally observable through token IDs, and weight tying can couple it to input lookup.

---

0x0B The decoding-policy pipeline

The output head produces scores. A policy composes transformations in a declared order:

raw logits
    -> processors
       hard constraints, history penalties, token biases
    -> warpers
       temperature, top-k, top-p, typical/min-p filters
    -> normalization when needed
       softmax or log_softmax
    -> selector
       argmax, multinomial draw, beam candidate top-k

Names differ across libraries, but ordering is semantic. Temperature preserves rank for positive $T$, yet changes top-p cumulative mass. A repetition penalty can change rank, so applying it before versus after top-k can change support. A probability-space multiplier is not equal to an additive logit bias.

def local_policy(
    logits,
    *,
    prefix_ids,
    processors=(),
    do_sample=False,
    temperature=1.0,
    top_k=None,
    top_p=None,
    generator=None,
):
    scores = logits
    for processor in processors:
        scores = processor(prefix_ids, scores)

    if not do_sample:
        return scores.argmax(dim=-1)

    if temperature <= 0:
        raise ValueError("sampling temperature must be positive")

    scores = scores / temperature
    if top_k is not None:
        scores = top_k_filter(scores, top_k)

    probs = torch.softmax(scores, dim=-1)
    if top_p is not None:
        return sample_top_p(probs, top_p, generator)

    return torch.multinomial(probs, 1, generator=generator)

This is a coherent teaching pipeline, not a universal API order. A release contract should record:

score type before every transform
processor order
normalization point
support and tie rules
selector and beam scoring
random generator and seed ownership

---

Policy state is typed and partly mutable

The outer request owns the generated prefix and invokes the next model step. The local policy reads a snapshot of that prefix and can own additional state:

State	Examples	Mutation owner
immutable policy config	processor order, temperature bounds, k/p thresholds, fallback, tie rule	release/configuration
read-only request view	prefix IDs, active/finished flag, current step, special IDs	M3.10 outer loop
constraint state	grammar/FSM state, legal token set, stop-sequence prefix	policy transition returned to M3.10
adaptive controller	target entropy/surprisal, running error, current temperature or k	local policy
randomness	algorithm, request stream/key, counter/draw index	local selector
beam metadata	parent IDs, cumulative scores, finished candidates	local search policy; outer loop reorders all model state

A processor transforms scores using model-independent or prefix-dependent rules. A warper changes the sampling distribution or support. Libraries draw that naming boundary differently; the executable contract is the ordered composition and score type at every edge.

PolicyInput:
    raw logits [B,V]
    read-only prefix/step/special-ID view
    constraint/controller/RNG state

PolicyOutput:
    selected global token IDs or ranked candidates
    selected score/log-probability metadata
    updated constraint/controller/RNG state
    diagnostic and fallback events

Most policy transforms do not commute

An additive bias $b$ and temperature $T$ give:

$$\frac{z+b}{T} \ne \frac{z}{T}+b \quad\text{unless }T=1\text{ or the bias is rescaled}.$$

Applying top-k before a repetition penalty can permanently remove a token that the penalty would have promoted relative to others. Applying top-p before a hard grammar mask can select a nucleus containing no legal token. A safe grammar-oriented ordering is commonly:

raw scores
    -> mandatory token suppression / vocabulary legality
    -> prefix-dependent penalties and biases
    -> adaptive or fixed score warping
    -> support truncation
    -> final legal-support assertion
    -> normalization
    -> selection

That ordering is a design, not a universal law. Some algorithms define their thresholds in probability space and require an earlier normalization. The manifest must pin the exact graph rather than relying on the words “processor” and “warper.”

Hard constraints define a conditional support

Let $A(s_t)$ be the legal token set in constraint state $s_t$. A logit-space mask is:

$$\tilde z_v = \begin{cases} z_v,&v\in A(s_t),\\ -\infty,&v\notin A(s_t). \end{cases}$$

The resulting distribution is the model policy conditioned on that legal set, after any other declared transforms. It is not the original full-softmax distribution. The state transition must use the emitted token:

$$s_{t+1}=\delta(s_t,y_t).$$

Before normalization:

if not any(isfinite(score) for score in constrained_scores):
    if empty_support_policy == "raise":
        raise EmptySupportError(constraint_state, step)
    constrained_scores = declared_fallback(...)

Silently restoring every raw logit defeats the constraint. Silently forcing EOS is also wrong unless EOS is legal and the fallback contract declares it.

Request-owned randomness and replay

A global generator consumed in rectangular batch order couples samples to batch composition. Removing a finished row or changing scheduler order can shift every later draw. Request-owned streams instead key randomness to a stable request identity and explicit draw schedule:

rng identity:
    algorithm/version
    base seed or key
    request/beam owner
    decoding step
    draw counter

draw contract:
    how many draws a selector consumes
    whether rejected or fallback paths consume draws
    how beam forks derive child streams

The strongest replay artifact records the processed support and selected draw metadata in addition to the seed. Floating-point or kernel changes can alter probabilities enough that the same uniform variate crosses a different cumulative boundary.

Local policy trace and assertions

def policy_step(raw_logits, request_view, state, config):
    scores = assert_finite_or_declared(raw_logits)
    trace = [summary("raw", scores)]

    for transform in config.ordered_transforms:
        scores = transform(scores, request_view, state)
        trace.append(summary(transform.name, scores))

    assert scores.shape == raw_logits.shape
    assert at_least_one_finite_score_per_active_row(scores)

    distribution = normalize_if_required(scores, config.selector)
    decision, rng_state = config.selector(distribution, state.rng)
    next_constraint = transition(state.constraint, decision)

    return decision, state.replace(
        rng=rng_state,
        constraint=next_constraint,
    ), trace

Useful trace fields are top-score margin, entropy, finite support size, legal support size, removed target/special IDs, selected raw and processed score, selected log-probability, RNG counter, and fallback reason. These are policy observations, not claims that the selected token is semantically correct.

---

0x0C Greedy decoding and temperature

Greedy is a local argmax

Greedy decoding selects

$$\hat y_t=\operatorname*{argmax}_v z_{t,v}.$$

Softmax is unnecessary because it preserves order:

$$\operatorname*{argmax}_v z_v =\operatorname*{argmax}_v\operatorname{softmax}(z)_v.$$

The sequence objective is

$$\log p(y_{1:T}\mid x) =\sum_{t=1}^{T}\log p(y_t\mid y_{<t},x).$$

Choosing the best token at each step need not maximize that sum over completed sequences. This is search myopia, not an implementation bug in argmax.

Temperature rescales logit gaps

For $T>0$,

$$p_T(v)=\frac{\exp(z_v/T)}{\sum_j\exp(z_j/T)}.$$

Pairwise log-probability ratios become

$$\log\frac{p_T(i)}{p_T(j)}=\frac{z_i-z_j}{T}.$$

Thus $T<1$ magnifies every gap and $T>1$ contracts it. Positive temperature does not change ranking, but it changes stochastic mass and any probability-dependent support filter.

Temperature also gives an exact entropy result. Let

$$H(T)=-\sum_v p_T(v)\log p_T(v).$$

Writing inverse temperature $\beta=1/T$ gives

$$H=\log Z(\beta)-\beta\mathbb E_{p_\beta}[z].$$

Differentiate using $\frac{d}{d\beta}\mathbb E[z]=\operatorname{Var}(z)$:

$$\frac{dH}{dT} =\frac{\operatorname{Var}_{p_T}(z)}{T^3}\ge0.$$

Entropy increases monotonically with positive temperature unless all supported logits are equal under the distribution. This derives “higher temperature is flatter” rather than leaving it as a slogan.

T=0 is not valid division. A user API may interpret it as greedy, but the implementation must branch to argmax.

Dynamic temperature is another policy

A policy may choose $T_t=f(z_t,h_t)$ from entropy, divergence, or generated history. That can adapt exploration to the current score row, but it defines a history-dependent distribution and requires its own calibration, bounds, and reproducibility tests.

The figure is a source-reported method design, not evidence that divergence-guided temperature universally improves generation.

---

0x0D Multinomial, top-k, and nucleus sampling

Multinomial needs nonnegative mass

probs = torch.softmax(scores, dim=-1)
token_ids = torch.multinomial(probs, num_samples=1, generator=generator)

Raw logits may be negative and are invalid as the weight argument to torch.multinomial. A categorical API explicitly parameterized by logits performs its own normalization; that is a different interface.

An explicit torch.Generator controls this draw. End-to-end reproducibility additionally requires stable upstream logits, processor order, batch/RNG assignment, kernels, and model state.

Top-k: fixed cardinality

Top-k retains the k highest scores and removes the rest:

def top_k_filter(logits, k):
    if not 0 < k <= logits.size(-1):
        raise ValueError("k must be in [1, V]")

    values, indices = torch.topk(logits, k, dim=-1)
    filtered = torch.full_like(logits, -torch.inf)
    return filtered.scatter(-1, indices, values)

After filtering, softmax renormalizes over retained finite entries. Threshold-based implementations can retain more than k entries under ties; torch.topk itself does not promise stable tied indices. Exact reproducibility must define tie behavior.

Top-p: dynamic probability mass

Nucleus sampling sorts probabilities and retains the smallest prefix whose cumulative mass reaches $p$.

def sample_top_p(probs, p, generator=None):
    if not 0.0 < p <= 1.0:
        raise ValueError("p must be in (0, 1]")

    sorted_p, sorted_ids = probs.sort(dim=-1, descending=True)
    cumulative = sorted_p.cumsum(dim=-1)

    # Keep the item that crosses p.
    remove = cumulative - sorted_p >= p
    sorted_p = sorted_p.masked_fill(remove, 0.0)

    z = sorted_p.sum(dim=-1, keepdim=True)
    if torch.any(z <= 0):
        raise RuntimeError("empty nucleus")
    sorted_p = sorted_p / z

    sampled_sorted = torch.multinomial(sorted_p, 1, generator=generator)
    return sorted_ids.gather(-1, sampled_sorted)

For sorted probabilities [0.52, 0.25, 0.13, 0.06, 0.04] and p=0.80, the retained prefix contains three tokens because 0.52+0.25=0.77 is below the threshold and adding 0.13 crosses it.

The nucleus paper contrasts flat and peaked distributions: one fixed k may be too small in a flat context and too permissive in a peaked context.

Active temperature, top-k, and top-p generally change the model’s full-softmax distribution. They are policy choices, not distribution-preserving acceleration.

One row through the full sampling pipeline

Take vocabulary rows [A,B,C,D,E] and raw logits

z = [3.0, 2.0, 1.0, 0.0, -1.0]

Suppose C is forbidden, T=2, k=3, and p=0.80.

1. Constraint processing:

   [3.0, 2.0, -inf, 0.0, -1.0]
   

2. Temperature:

   [1.5, 1.0, -inf, 0.0, -0.5]
   

3. Top-k retains A,B,D:

   [1.5, 1.0, -inf, 0.0, -inf]
   

4. Softmax over retained support is approximately:

   A: 0.547
   B: 0.331
   D: 0.122
   

5. Top-p with p=0.80 keeps A and B because their cumulative mass is 0.878, then renormalizes:

   A: 0.623
   B: 0.377
   

The final draw is not from the original model distribution, nor from temperature alone, nor from top-k alone. It is from the composition in the declared order.

---

Sampling-family contracts beyond top-k and top-p

Most truncation methods create a support set $S(z)$ and then sample from the renormalized distribution:

$$q(v) = \frac{p(v)\mathbf1[v\in S]} {\sum_{j\in S}p(j)}.$$

They differ in how $S$ is constructed. Once active, none preserves the original full-softmax policy.

Method	Controlled quantity	Support rule	Persistent state
top-k	cardinality	k highest scores	none
top-p	cumulative probability mass	smallest probability-ranked prefix reaching p	none
locally typical	distance from expected surprisal	probability mass ranked by surprisal deviation	none
eta sampling	entropy-dependent probability floor	probability above adaptive eta threshold	none
min-p	probability relative to current maximum	probability at least alpha times maximum	none
top-n-sigma	distance below maximum in logit standard deviations	logit at least max minus n sigma	none
Mirostat	observed token surprisal/perplexity	controller-adjusted truncation	running controller value
dynamic temperature	entropy/divergence/task signal	full or subsequently truncated support after T_t	controller inputs/state

Locally typical support

For model distribution $p$, token surprisal and row entropy are:

$$I(v)=-\log p(v), \qquad H(p)=\sum_vp(v)I(v).$$

Locally typical sampling ranks tokens by:

$$d_v=|I(v)-H(p)|$$

and retains a probability-mass-limited prefix in that ranking. Unlike top-p, the maximum-probability token need not rank first: its surprisal can be farther from the expected information content than a moderately probable token.

Eta sampling

The desmoothing interpretation treats low-probability mass as partly induced by a smoothing distribution. In the notation used by Hewitt et al., an entropy-sensitive threshold can be written:

$$\eta(p) = \min\left( \varepsilon, \sqrt{\varepsilon}\,e^{-H(p)} \right).$$

Candidates below the threshold are removed. In a high-entropy row, $e^{-H(p)}$ is small and the rule relaxes the absolute floor; in a low-entropy row it behaves more like a fixed epsilon cutoff. The epsilon definition and strict/equality convention remain policy fields.

Min-p and relative logit gaps

Min-p retains:

$$p(v)\ge\alpha\,p_{\max}, \qquad 0<\alpha\le1.$$

Because the softmax denominator cancels, this is equivalent to:

$$z_v\ge z_{\max}+\log\alpha.$$

After temperature scaling, the equivalent raw-logit condition is:

$$z_v\ge z_{\max}+T\log\alpha.$$

Thus increasing temperature widens the admissible raw-logit gap for fixed alpha. The method has an intuitive implementation and a primary proposal with reported evaluations, but a later critical analysis disputes superiority claims from that evidence. M3.8 therefore treats min-p as a mechanism, not an established winner.

Top-n-sigma and affine invariance

Top-n-sigma, reconstructed from A15 and its primary paper, keeps:

$$z_v\ge z_{\max}-n\,\sigma(z).$$

Under a positive affine score transformation $z'=az+c$:

$$z'_v\ge z'_{\max}-n\sigma(z') \iff az_v+c\ge az_{\max}+c-na\sigma(z) \iff z_v\ge z_{\max}-n\sigma(z).$$

The support is therefore invariant to positive temperature scaling when the standard deviation is measured on the same transformed score row. It is not invariant to arbitrary token-wise biases or penalties.

Feedback-controlled surprise and temperature

Mirostat targets a chosen surprisal $\tau$ rather than one fixed support parameter. A schematic controller observes emitted-token surprise $e_t$ and updates:

$$\mu_{t+1} = \mu_t-\gamma(e_t-\tau).$$

The controller then maps $\mu_t$ to the next truncation parameter. The exact Mirostat version, Zipf estimate, units (bits or nats), gain, initialization, clamps, and update timing are part of the policy state. Resetting $\mu$ at every token reduces the method to a misnamed stateless heuristic.

Dynamic-temperature methods instead set $T_t=f(p_t,\text{history})$. A15 surveys KL-guided, rule-based code, and entropy-based variants. They use different signals and costs: one may require a second model distribution, another token-type rules, and another the current row entropy.

Sensitivity and evaluation contract

Sampler comparisons must sweep controls rather than compare arbitrary defaults:

fixed checkpoint and tokenizer
fixed prompt/evaluation manifest
fixed processor and constraint order
request-owned seeds and enough repeated draws
matched output-length and stopping policy
quality, diversity, repetition, calibration, and latency reported separately
support size / entropy / fallback telemetry

Failure geometry differs by method:

Failure	Likely cause
support collapses to one token	aggressive k/p/epsilon/alpha/n or concentrated row
incoherent tail enters	permissive truncation or high temperature
controller oscillates	excessive gain, delayed observation, missing clamp
adaptive state leaks across requests	controller stored globally
no legal token remains	truncation composed incorrectly with hard constraint
paper result does not reproduce	tokenizer, model, prompt, metric, or implementation mismatch

No one scalar ranks policies universally. A policy that improves creative diversity can be inappropriate for exact extraction, constrained code, or target-policy verification in M7.

0x0E Score processors, repetition, and hard constraints

Processors can depend on read-only prefix state even though M3.10 owns advancing that prefix.

Additive token bias

For bias vector $b$,

$$z'_v=z_v+b_v.$$

In probability space this multiplies unnormalized mass by $e^{b_v}$:

$$p'(v)\propto p(v)e^{b_v}.$$

Setting $b_v=-\infty$ forbids a token. At least one finite score must remain.

Sign-sensitive repetition penalty

One common implementation applies penalty $r>1$ to previously generated token rows as

$$z'_v= \begin{cases} z_v/r, & z_v>0,\\ r z_v, & z_v<0. \end{cases}$$

Both branches reduce the penalized token’s relative score: positive values shrink; negative values become more negative. Dividing every logit regardless of sign can increase a negative score toward zero and accidentally reward it.

seen_scores = scores.gather(-1, prefix_ids)
seen_scores = torch.where(
    seen_scores < 0,
    seen_scores * penalty,
    seen_scores / penalty,
)
scores = scores.scatter(-1, prefix_ids, seen_scores)

This particular rule is a library/paper convention, not the only repetition control. Frequency penalties, presence penalties, no-repeat n-gram constraints, and contrastive penalties define different functions.

Hard constraints and finite-state policy

Minimum length may forbid EOS until a step threshold. Bad-word or grammar constraints may mask tokens based on a finite-state machine. Forced-prefix decoding can leave exactly one legal token at a step.

constraint state + prefix
    -> legal vocabulary mask [V]
    -> processed scores [V]

The local policy consumes the state; the outer loop owns updating and checkpointing it. Beam search requires one constraint state per beam and must reorder it with parent IDs.

---

0x0F Beam search: joint parent/token candidates

For K active beams:

local logits:   [B,K,V]
beam scores:    [B,K]
candidate grid: [B,K,V]

Using log probabilities,

$$c_{b,k,v}=s_{b,k}+\log p(v\mid y_{b,k,<t},x).$$

Flatten the parent and vocabulary axes and retain K candidates:

def beam_step(logits, beam_scores, beam_width):
    B, K, V = logits.shape
    log_probs = torch.log_softmax(logits, dim=-1)
    candidates = beam_scores[..., None] + log_probs
    next_scores, flat_ids = candidates.reshape(B, K * V).topk(
        beam_width, dim=-1
    )
    parent_ids = flat_ids // V
    token_ids = flat_ids % V
    return next_scores, parent_ids, token_ids

The quotient/remainder derivation follows the parent-major flattened index:

$$i=kV+v, \qquad k=\left\lfloor\frac{i}{V}\right\rfloor, \qquad v=i\bmod V.$$

Parent IDs are mandatory metadata. Prefix tokens, attention masks, position state, constraint state, and every K/V cache layer must follow the same reorder.

First-step and finished-set contracts

If all initial beams represent the same prefix, activate only one:

beam_scores = torch.full((B, K), -torch.inf, device=logits.device)
beam_scores[:, 0] = 0.0

Otherwise the first expansion can duplicate candidates from identical parents. EOS candidates usually move to a finished set while additional live candidates fill the beam frontier. Stopping when the first EOS appears is not generally safe.

Length penalty

Raw log-probability sums favor shorter sequences because each additional term is non-positive. One generic normalized score is

$$s_{norm}(y)=\frac{\sum_t\log p(y_t\mid y_{<t},x)}{g(|y|)},$$

where $g$ may be $|y|^\alpha$ or a smoothed function such as

$$g(|y|)=\left(\frac{5+|y|}{6}\right)^\alpha.$$

The exact formula, when it is applied, and how it enters early-stopping bounds are framework contracts. A length penalty changes ranking; a no-repeat n-gram rule changes legal continuations.

Finite K remains approximate. Even K=V is not exhaustive after multiple steps because the number of prefixes grows as $V^t$.

---

0x10 Search policies and alternative output architectures

Sampling changes one local conditional distribution. Search instead retains or revisits multiple prefixes, and alternative output architectures change how the vocabulary distribution is parameterized or trained. These are separate design axes.

Search policies are objectives over prefixes

For prefix $y_{1:t}$, the unnormalized sequence score is:

$$g(y_{1:t}) = \sum_{i=1}^{t}\log p(y_i\mid y_{<i},x).$$

Search	Frontier/order	Extra model state	Objective caveat
greedy	one live prefix	one cache	local argmax, not sequence optimum
beam	top K prefixes per depth	K caches and metadata	finite-width pruning and length bias
diverse beam	grouped beams plus diversity penalty	group/frontier state	modified ranking objective
best-first beam	priority queue across depths	bounded queue plus caches	priority and stopping rule must be pinned
A*-style search	path score plus heuristic	queue, heuristic state	guarantees require heuristic assumptions
best-of-N/verifier	complete rollouts ranked by verifier	N rollouts and verifier	not model-likelihood beam search

Encoder-decoder search can add source coverage. The GNMT-style family combines log-likelihood, length normalization, and an attention-coverage term:

$$S(y,x) = \frac{\log P(y\mid x)}{\operatorname{lp}(y)} +\operatorname{cp}(x,y).$$

The exact length exponent, coverage clipping, attention tensor, and stopping bound are one search contract. Coverage has no direct analogue when the model does not expose encoder-decoder attention.

Best-first beam changes expansion order. Instead of synchronously expanding all live prefixes at depth t, it pops the highest-priority item from a bounded queue, expands it, and reinserts candidates. Priority can use path score, length normalization, or a heuristic:

$$\operatorname{priority}(y)=g(y)+h(y).$$

Calling a search “A*” does not establish optimality; the heuristic and stopping conditions must satisfy the paper’s assumptions.

Policies that request more than one finished logit row

Contrastive search combines model confidence with a representation-degeneration penalty:

$$S(v) =(1-\alpha)\log p(v\mid y_{<t}) -\alpha\max_{j<t}\cos(h_v,h_j).$$

Evaluating candidate hidden state $h_v$ for several tokens requires additional model work or a specialized batched path. It is not merely a function of the already finished score row.

DoLa contrasts a mature-layer distribution with one selected premature-layer distribution. A schematic contrast is:

$$s_v = \log p_{\text{mature}}(v) -\log p_{\text{premature}}(v).$$

Layer selection, candidate restriction, normalization, and fallback are part of the method. This reaches into intermediate states and must not be confused with ordinary final-layer temperature.

Self-evaluation decoding methods detect uncertain or “chaotic” positions, construct alternatives, and ask the model or another evaluator to score them. They add candidate-generation and evaluation calls, their own prompt/template, and extra cache state.

Test-time best-of-N, verifier-guided rollouts, filler/pause tokens, and latent reasoning likewise change the number or type of model steps. They cross from a single local policy into M3.10 execution or M7 proposal/verification and are kept as boundaries here.

Alternative output heads change training and inference

Architecture/objective	Training path	Full-vocabulary inference	Contract change
full affine softmax	score and normalize all V rows	direct full distribution	baseline used above
hierarchical softmax	target path through a class/tree hierarchy	traverse or score hierarchy	token probability is product of path decisions
adaptive softmax	frequent shortlist plus rare-token clusters	evaluate head and required tails	learned cluster assignment and dimensions
sampled/importance softmax	estimate objective from target plus sampled negatives	usually still needs full or specialized head	estimator, proposal, and correction
shortlist/importance NMT	train against a selected vocabulary subset	selected or full vocabulary by deployment	subset construction is part of objective
mixture of softmaxes	mix several normalized components	several projections/softmaxes	increases expressivity and compute

In a hierarchical head, a word probability is a product over decisions on its root-to-leaf path:

$$p(w\mid h) = \prod_{n\in\operatorname{path}(w)} p(d_n(w)\mid h,n).$$

This can replace an O(V) normalization with path work, but tree construction affects both computation and learned confusion structure.

Adaptive softmax uses token-frequency imbalance. A frequent word can be scored in a short head path; a rare word uses a cluster probability and conditional tail probability:

$$p(w\mid h) = p(c(w)\mid h)\, p(w\mid c(w),h).$$

Cutoffs, cluster ownership, tail dimensions, and weight tying are checkpoint architecture. This is not a drop-in inference warper over full-softmax logits.

Sampled and importance objectives avoid evaluating every negative class during training. Unless the estimator and correction recover the declared full objective, they optimize an approximation. At inference, the system still needs a way to rank the intended output vocabulary.

Why a single softmax has an expressivity bottleneck

For context matrix $H\in\mathbb R^{N\times d}$ and output matrix $E\in\mathbb R^{V\times d}$, unnormalized scores are:

$$A=HE^\top, \qquad \operatorname{rank}(A)\le d.$$

The log-probability matrix subtracts one row normalizer:

$$\log P = A-c\mathbf1^\top, \qquad c_i=\operatorname{LSE}(A_i).$$

Therefore:

$$\operatorname{rank}(\log P)\le d+1.$$

The softmax-bottleneck analysis asks whether the target family of conditional log distributions requires higher rank. Mixture of Softmaxes combines several component distributions:

$$p(v\mid h) = \sum_{m=1}^{M}\pi_m(h)\, \operatorname{softmax}(z^{(m)}(h))_v.$$

The logarithm of that mixture is not one low-rank affine score matrix minus one normalizer, so it can represent a richer family at the cost of multiple heads, normalizations, and mixture state. It is an architecture/training choice, not a sampling method.

Paper-by-contract comparison

Method	Changes weights/training?	Needs extra forward work?	Preserves base local distribution?	Primary audit question
beam / best-first	no	multiple live prefixes	ranks sequence likelihood unless penalties modify it	are queue, score, stop, and cache reorder exact?
contrastive search	possibly contrastive training; decode rule changes	candidate hidden states	no	is degeneration score computed for the intended candidates?
DoLa	no base retraining required by method	intermediate-layer projections	no	are mature/premature layers and candidate mask pinned?
SED	no base retraining required by method	self-evaluation calls	no	is extra work and evaluator prompt included?
adaptive/hierarchical softmax	yes, head architecture	specialized head path	defines a different normalized model	can checkpoint/token probabilities be reconstructed?
sampled objective	yes, training estimator	fewer training classes	estimator-dependent	are proposal and correction unbiased/declared?
mixture of softmaxes	yes	multiple component heads	defines a different model	are mixture gates and component normalizers correct?

Logit diagnostics are observations, not decoding rules

A score row exposes several useful diagnostics:

top-1 margin:
    z_(1) - z_(2)

entropy:
    -sum_v p_v log p_v

effective support / perplexity of the row:
    exp(H(p))

normalizer:
    logsumexp(z)

energy-style score:
    -T * logsumexp(z / T)

The top-1 margin is invariant to adding a constant to all logits but changes under temperature scaling. Entropy depends on the full distribution. A large margin or low entropy means the policy is concentrated; it does not establish semantic correctness or calibrated confidence.

Calibration asks whether confidence frequencies match empirical outcomes under a declared task and grouping. Token-level calibration is complicated by multiple valid continuations and sequence dependence. Temperature scaling fitted on held-out data can improve one calibration metric while changing generation diversity; the fitted temperature is a deployment policy artifact, not a new base-model weight.

Logit lens and intermediate unembedding

Because every residual row has width d_model, one can apply the final norm/head approximately or directly to an intermediate layer:

$$z^{(l)}=\operatorname{Head}(N_{final}(h^{(l)})).$$

This logit lens asks which vocabulary directions are already visible. It is an analysis intervention:

intermediate row
    -> final-normalization convention
    -> shared output head
    -> diagnostic logits

The actual model still sends $h^{(l)}$ through later attention, MLP, residual, and normalization operations. Intermediate top tokens are not promises about the final prediction.

Tuned lenses learn a translator for each layer before unembedding. Better predictive agreement may mean the translator compensates for representation mismatch; it does not prove that the unmodified intermediate state was literally using the final vocabulary basis.

Vocabulary bias and score preprocessing

Frequent-token priors, tokenizer segmentation, output bias, and embedding norms can all influence raw scores. Subtracting an empirical prior or calibrating logits defines a modified policy:

$$z'_v=z_v-\lambda\log\pi_v,$$

where $\pi_v$ is a reference token frequency. This can be useful in a studied setting, but it is not a universally correct “debiasing” operation. The reference corpus, zero-frequency smoothing, tokenizer, and $\lambda$ become part of the policy contract.

Likewise, a source article may visualize compressed information in logits or propose preprocessing based on a secondary distribution. Such a transformation must be evaluated as an explicit processor with baseline, ablation, and task-specific metrics rather than imported as an assumed improvement.

Hidden reasoning and filler tokens cross the local boundary

Pause tokens, chain-of-thought tokens, and latent-reasoning methods spend additional model steps before the final answer. They can change the hidden state that reaches the output head:

ordinary policy:
    one state -> one output decision

filler / visible reasoning:
    select extra tokens -> run additional model steps -> later decision

latent reasoning:
    evolve hidden state under an architecture-specific recurrence -> later head

These methods are not equivalent to lowering entropy on the same fixed score row. Their evaluation must account for additional tokens/compute, training distribution, answer extraction, and whether intermediate reasoning is visible or latent.

---

0x11 Shape and type ledger

Object	Shape	Type / meaning
final hidden matrix	[B,S,d_model]	Continuous states after final norm.
one decision row	[B,d_model]	Last valid or otherwise selected row.
algebraic output matrix W_U	[d_model,V]	Row-vector notation.
PyTorch lm_head.weight	[V,d_model]	Stored nn.Linear orientation.
optional output bias	[V]	Added to every output row.
all-row logits	[B,S,V]	Raw unnormalized scores.
one-step logits	[B,V] or [B,1,V]	Raw scores for decision rows.
probabilities	same as logits	softmax(logits, dim=-1).
log-probabilities	same as logits	log_softmax(logits, dim=-1).
class-index labels	[B,S] or flattened [B*S]	Integer target IDs.
distribution labels	[B,S,V] or flattened [B*S,V]	Label-smoothed/soft targets.
selected token IDs	[B] or [B,1]	Local policy output; M3.10 appends and advances state.
beam candidate scores	[B,K,V]	Parent sequence score plus local log-probability.
flattened beam candidates	[B,K*V]	Joint parent/token candidate axis.
parent beam IDs	[B,K]	Prefix/cache reorder metadata.
next beam token IDs	[B,K]	Retained vocabulary IDs.

Score-type contract

Tensor	Valid direct consumers
raw logits	cross entropy, argmax, temperature, score processors, softmax/log-softmax
log-probabilities	NLL/KL-style losses, argmax, sequence log-score sums, exp
probabilities	multinomial sampling, cumulative top-p mass, probability reporting

---

Official implementation reconciliation

The mathematics above is implementation-independent; runtime behavior is not. This reconstruction pins three official source snapshots:

Runtime evidence	Revision	Contract observed
Hugging Face Transformers	v5.13.1, commit 4626421	LLaMA bias-free head, tied-weight metadata, row slicing, ordered generation processors
NVIDIA Megatron-LM	main snapshot 5adbee2	global MAX/SUM vocabulary-parallel cross-entropy and saved softmax shard
vLLM	main snapshot 7b50d2c	float32 sampler path with explicit legality, processor, penalty, temperature, min-p, top-k/top-p, and selector order

The Transformers snapshot appends sampling transforms in an explicit sequence that includes temperature, top-h, top-k, top-p, min-p, typical, and epsilon cutoffs after its earlier processors. It explicitly places min-p after temperature. The vLLM snapshot uses a different staged pipeline: whitelist and bad-word suppression, non-argmax-invariant processors, penalties, greedy candidate, temperature, min-p, top-k/top-p, then random/greedy selection.

Neither order is the universal mathematical definition. They demonstrate why framework, version, and ordered transforms belong in the release manifest.

The pinned Megatron snapshot implements:

local logits -> global MAX
target owner/masked local target
target-score SUM
denominator SUM
loss and local normalized shard
saved normalized shard + target mask/index
owner-only subtraction in backward

Its label-smoothing path uses implementation-specific algebra. A release that enables smoothing must compare per-row loss and local shard gradients with the declared global-V convention rather than assuming the same parameter name implies the same target distribution.

Evaluation ladder

Layer	Required test	Main invariant
head unit	tiny affine oracle and finite differences	logits and dH/dE/db
normalization	extreme logits, translation, smoothing	stable loss and P minus Q
training schedule	materialized versus chunked/fused	per-row loss, valid sum/count, gradients
vocabulary migration	permutation/append/tied alias	mapped raw logits and one optimizer owner
distributed loss	full versus sharded oracle	loss, local gradients, target owner
processor pipeline	order, masks, fallback	nonempty legal support and declared score types
stochastic policy	repeated seeded draws and distribution test	request-owned replay plus statistical agreement
search	toy exhaustive tree for small V/T	parent/token IDs, score, stopping, state reorder
end to end	fixed prompts and policy manifest	output IDs/log-probs within declared determinism tier

Statistical policies need more than one golden string. For a fixed processed row, draw many samples and compare empirical frequencies or a categorical goodness-of-fit statistic under a predeclared tolerance. End-to-end generation then adds prompt, stopping, length, batch, and scheduler dimensions.

Release, canary, and rollback

A releasable output-policy bundle contains:

tokenizer and ordered vocabulary hash
model/head/checkpoint and tying manifest
loss and sharding manifest
processor/warper/constraint graph
search or selector configuration
RNG algorithm and request-stream policy
runtime/framework/kernel revisions
evaluation corpus/split and acceptance thresholds
previous known-good bundle identifier

Release stages:

1. Run all local numerical and migration oracles.
2. Compare pinned reference and candidate runtimes on fixed hidden rows.
3. Run deterministic prompt cases and statistical sampler cases.
4. Canary by request class with old/new policy telemetry side by side.
5. Gate on invalid-support, NaN, fallback, stop, latency, quality, diversity, and constraint-violation metrics.
6. Promote the immutable bundle, retaining a one-action rollback target.

Rollback restores tokenizer, checkpoint, aliases, policy, RNG semantics, and runtime together. Rolling back only temperature while keeping a changed vocabulary or transform order does not restore the old behavior.

Claim and visual evidence contract

This article uses four evidence labels:

Source-local:
    reconstructed from A2/A15; may include a source correction

Externally verified:
    supported by a primary paper or pinned official implementation

Rewrite interpretation:
    derived synthesis, engineering rule, or audit recommendation

Uncertain / source-reported:
    empirical or implementation-sensitive claim kept within its source scope

Reused A15 figures are retained when they teach an owned mechanism. Captions state whether the figure is a generic tensor contract, a paper-specific algorithm, or one experiment. Figures about pause tokens, latent reasoning, and multi-step proposal systems are not promoted into M3.8’s local policy spine.

0x12 Failure lab

Hidden-state and projection bugs

1. Treating final hidden states as logits. [B,S,d_model] is not vocabulary-indexed unless d_model == V by meaningless coincidence.

2. Skipping the final norm expected by a Pre-LN checkpoint. The head receives same-shaped but wrong hidden states.

3. Applying the target LM head directly to encoder memory. Encoder-decoder generation projects decoder-owned target rows.

4. Transposing the PyTorch output weight manually. nn.Linear already uses stored [V,d_model] weights as weight.T in the forward calculation.

5. Loading a bias-free checkpoint into a biased output head. An extra [V] parameter changes logits even though the main weight shape matches.

6. Projecting an MLP intermediate. The head expects model-width final states, not [B,S,d_ff] branch activations.

Score-type bugs

7. Calling logits probabilities. Raw scores can be negative and do not sum to one.

8. Calling every multiclass logit binary log-odds. Vocabulary probabilities depend on the full row normalization.

9. Softmax over the sequence axis. dim=-1 must normalize vocabulary candidates for each state row.

10. Applying softmax to probabilities again. This changes the distribution.

11. Claiming softmax over normalized log-probabilities changes the distribution. It recovers probabilities mathematically; the operation is merely redundant.

12. Computing log(softmax(logits)) as separate low-precision operations. Use the numerically stable log_softmax implementation.

13. Sampling raw logits with torch.multinomial. Multinomial weights must be non-negative; probability conversion belongs before sampling.

14. Using A15’s bare log_softmax without importing that symbol. With import torch.nn.functional as F, call F.log_softmax.

Training-interface bugs

15. Passing probabilities into CrossEntropyLoss. The official contract expects unnormalized logits.

16. Passing raw logits into a KL/NLL path expecting log-probabilities. The head and criterion conventions disagree.

17. Taking argmax before the training loss. Discrete selection breaks the intended differentiable vocabulary-score path.

18. Flattening [B,T,V] as if T were the class axis. Preserve V as the final class dimension: [B*T,V].

19. Shifting labels inside both the model and data pipeline. Output-head shape survives while supervision moves two positions.

20. Training PAD or ignored rows as ordinary target IDs. M3.9’s label mask or ignore_index must agree with the loss.

Inference-row bugs

21. Using hidden[:, -1] with right-padded unequal prompts. Shorter rows project a PAD-position state.

22. Applying argmax over sequence length. The returned index is a time position, not a vocabulary ID.

23. Computing all prompt logits and accidentally selecting the first row. Every tensor is valid, but the decision conditions on the wrong prefix position.

24. Reading logits_to_keep=0 as zero logits. In the Transformers v5.13.1 LLaMA slice convention it selects all positions.

Vocabulary and tying bugs

25. Tokenizer ID order differs from output-row order. Scores decode to the wrong strings while shapes remain valid.

26. Adding tokenizer entries without resizing/adapting the output rows. New IDs are missing or use untrained vocabulary weights.

27. Resizing only the input embedding in an untied model. Lookup accepts new IDs but the head cannot score the same output space.

28. Copying embedding values and claiming the weights are tied. Separate Parameters diverge after optimization.

29. Tying source, target, and output matrices with incompatible vocabularies. Equal shapes do not establish shared token identity.

30. Assuming every LLaMA-style checkpoint ties embeddings. The architecture supports configuration-specific behavior; verify actual parameter identity.

31. Creating the optimizer before replacing one Parameter with a tied reference. The optimizer can retain a stale parameter object or duplicate state.

32. Treating tied output rows as cosine-nearest neighbors. The head computes learned dot products, optionally plus bias; vectors are not necessarily unit-normalized.

Boundary bugs

33. Putting top-k, top-p, beam state, or EOS handling inside the learned LM-head module. M3.8 owns the adjacent policy interface, but the head itself remains a learned projection and the outer loop owns persistent state.

34. Treating the local projection ledger as a complete runtime model. M3.8 counts output-head parameters, MACs, logical logits, and local policy objects. M4.1 combines those counts with the rest of the model and analyzes materialization and hardware traffic.

Policy and search bugs

35. Passing raw logits to torch.multinomial. Convert to nonnegative probability weights or use an API explicitly accepting logits.

36. Dividing by temperature=0. Greedy is a branch or a limit, not literal zero-temperature arithmetic.

37. Assuming temperature can change top-k membership. Positive temperature preserves ranking unless another transform intervenes; it can still change top-p support.

38. Dropping the top-p threshold-crossing token. The remaining support can have cumulative mass below p.

39. Returning a top-p sorted position as the vocabulary ID. Gather through the sort permutation.

40. Assuming threshold top-k retains exactly k tokens under ties. An equality policy can keep more than k, while selected tied indices can be unstable.

41. Applying one repetition formula to positive and negative logits. Dividing a negative value moves it toward zero and can increase its relative score.

42. Masking every token to negative infinity. No categorical distribution remains; raise or provide an explicit legal fallback.

43. Ranking beam tokens without parent sequence scores. Search then chooses local tokens rather than scored continuations.

44. Using % K to decode a flattened beam candidate token. The fast-changing axis has width V; token is flat_id % V.

45. Reordering beam token IDs without all beam-owned state. Cache, positions, masks, constraints, RNG state, and prefix metadata must follow parent IDs.

46. Initializing every duplicate first-step beam at zero. Identical parents can occupy the frontier with duplicate candidates.

47. Stopping at the first EOS without a valid score bound. A live hypothesis may still outrank the finished one under the configured length score.

48. Calling active top-k/top-p distribution-preserving. They intentionally truncate the original full-softmax policy.

49. Treating a seeded local draw as complete reproducibility. Batch order, upstream numerics, processor order, and random-stream ownership also matter.

50. Calling test-time best-of-N or verifier-guided rollout ordinary beam search. These methods add complete rollouts and an external or learned verifier objective.

Extended contract failures

51. Averaging chunk means instead of accumulating loss sum and valid count. Uneven chunks or ignore masks receive the wrong weight.

52. Saving only the scalar loss as a fused-kernel conformance check. Row permutations and wrong gradients can cancel in the mean.

53. Using vocabulary size instead of ordered vocabulary identity. Equal shapes can attach every score to the wrong token.

54. Appending rows and claiming old probabilities are unchanged. New rows enter the softmax denominator even when old raw logits match.

55. Migrating model rows but not optimizer moments or master weights. The next update applies old state to the wrong token address.

56. Restoring tied values without restoring one Parameter and optimizer owner. Subsequent gradients and moments diverge.

57. Normalizing each vocabulary shard independently. Local shard probabilities do not form the global model distribution.

58. Subtracting the target on every vocabulary rank. Only the rank that owns the global target row applies the one-hot term.

59. Skipping a distributed loss collective only on ignored rows of one rank. Collective schedules diverge and can deadlock.

60. Applying top-p separately on each vocabulary shard. Local cumulative supports are not the global probability-ranked nucleus.

61. Sharing one adaptive controller across requests. Mirostat or dynamic temperature state leaks history between users.

62. Treating min-p benchmark superiority as settled. The mechanism is valid, while published evaluation claims have been critically disputed and remain source-scoped.

63. Calling sampled softmax an inference sampler. It is a training-objective approximation/family, not top-k-style decoding.

64. Calling adaptive or hierarchical softmax a drop-in warper. They change checkpoint architecture and token-probability factorization.

65. Assuming an A-named search is optimal without heuristic conditions.* Queue priority and stopping guarantees depend on the actual heuristic.

66. Comparing two runtimes without pinning policy order and log-prob mode. The same named knobs can produce different transformed scores and metadata.

---

0x13 Handoff to training, inference, and acceleration

M3.8 contains the complete local learned-to-policy boundary:

hidden: [B,d_model]
    -> lm_head
logits: [B,V]

The policy can now return a token or beam decision:

logits: [B,V]
    -> optional score processors / temperature
    -> probability conversion when required
    -> greedy, top-k, top-p, or beam policy
selected token IDs: [B] or [B,1]

M3.9 consumes all supervised logit rows to train the model. M3.10 appends selected IDs, stops finished rows, calls the next model step, and maintains cache state. M7 owns speculative and multi-token acceleration while preserving its declared target policy.

---

0xFF References

Local source and neighboring contracts

• Exploring the Transformer Series (2) — Overall Architecture
• Exploring the Transformer Series (15) — Sampling and Output
• M3.9 — Training Execution
• M3.10 — Single-Request Inference Execution
• M2.1 — Tokenization and Sequence Contracts
• M2.2 — Embeddings and Initial Model State
• M3.7 — Residual Stream and Block Assembly
• M4.1 — Transformer Cost Model

Output geometry, objectives, and large vocabularies

• Vaswani et al., Attention Is All You Need
• Press and Wolf, Using the Output Embedding to Improve Language Models
• Inan et al., Tying Word Vectors and Word Classifiers: A Loss Framework for Language Modeling
• Grave et al., Efficient Softmax Approximation for GPUs
• Jean et al., On Using Very Large Target Vocabulary for Neural Machine Translation
• Yang et al., Breaking the Softmax Bottleneck: A High-Rank RNN Language Model

Sampling and adaptive decoding policies

• Holtzman et al., The Curious Case of Neural Text Degeneration
• Meister et al., Locally Typical Sampling
• Hewitt et al., Truncation Sampling as Language Model Desmoothing
• Basu et al., Mirostat: A Neural Text Decoding Algorithm that Directly Controls Perplexity
• Nguyen et al., Turning Up the Heat: Min-p Sampling for Creative and Coherent LLM Outputs
• Schaeffer et al., Min-p, Max Exaggeration: A Critical Analysis of Min-p Sampling in Language Models
• Li et al., EDT: Improving Large Language Models’ Generation by Entropy-based Dynamic Temperature Sampling
• Tang et al., Top-n-sigma: Not All Logits Are You Need

Search and multi-pass decoding

• Wu et al., Google’s Neural Machine Translation System: Bridging the Gap between Human and Machine Translation
• Meister et al., Best-First Beam Search
• Su et al., A Contrastive Framework for Neural Text Generation
• Chuang et al., DoLa: Decoding by Contrasting Layers Improves Factuality in Large Language Models
• Wang et al., SED: Self-Evaluation Decoding Enhances Large Language Models for Better Generation
• Belrose et al., Eliciting Latent Predictions from Transformers with the Tuned Lens

Pinned implementations and API contracts

• Hugging Face Transformers v5.13.1, Generation configuration
• PyTorch 2.13, torch.multinomial
• PyTorch 2.13, torch.topk
• PyTorch 2.13, torch.nn.CrossEntropyLoss
• PyTorch 2.13, torch.nn.functional.log_softmax
• Hugging Face Transformers v5.13.1 (4626421d), LlamaForCausalLM
• Hugging Face Transformers v5.13.1 (4626421d), GenerationMixin
• Hugging Face Transformers v5.13.1 (4626421d), logits processors and warpers
• Hugging Face Transformers v5.13.1, LlamaConfig
• NVIDIA Megatron-LM (5adbee2), vocabulary-parallel cross-entropy
• vLLM (7b50d2c), V1 sampler pipeline