MLPSpeculatorPreTrainedModel

MLPSpeculatorPreTrainedModel at 4 prediction stages and stage width 4096. 2 published checkpoints share this shape.

Implementing it

One forward, written out, per checkpoint whose arithmetic actually differs. A shape groups checkpoints by class, depth, width and layer stack — none of which decides an activation, a rope base or a routing rule — so where the members of this shape disagree there is a block each, and each names the checkpoint it was generated from.

No forward was recorded for this shape, so there is no pass above to compare it against — this block is what the published configuration says the arithmetic is, not a transcript of a run. Where the configuration does not say, the line says that instead of guessing.

The names are the ones the published checkpoint uses where a name is published, and canonical otherwise. A checkpoint loader may store them differently — fusing a gate/up pair into one matrix, or splitting one published projection in two — and those are that loader's names, not the model's. Everything spelled here is what the download contains.

From granite-20b-code-instruct-accelerator

This checkpoint only — its arithmetic differs from the others of this shape.

# A draft head, not a model. It proposes the next few tokens so the base
# model can verify several in one pass. There is no attention here, no
# position and no cache: the head sees one hidden state and one token,
# and everything it knows about the sentence arrived inside that state.

# state  = the base model's last hidden row      [6144]
# token  = the id the base model just emitted

# `scale_input: false` — the state is used as it arrives, unnormalised.

state_weight = 0.5 ** (0.5 / 4)   # = 0.917004
emb_weight   = sqrt((1 - state_weight**2) * 4096 / 2)   # = 18.0512
alpha        = emb_weight / state_weight                # = 19.6849
# Both are derived from the config and neither is a tensor. They are
# chosen so the original state still accounts for half the magnitude
# by the final head. Only the ratio is applied: upstream writes
# `add(state, z, alpha=emb_weight/state_weight)` and lets the norm
# below absorb the missing state_weight — exact up to its epsilon.

for i in 0 .. 3:
    state = proj[i] @ state               # [4096]; proj[0] reads 6144, every
                                          #   stage after it is square
    state = state + alpha * emb[i][token]
    state = rmsnorm(state, ln[i].weight, ln[i].bias)
    #  ^ not a LayerNorm despite the name. `LayerNormParameterized` runs
    #    here with its mean subtraction off: x * rsqrt(mean(x^2) + eps),
    #    then weight * x + bias. It is an RMS norm with a bias.
    state = gelu(state)                   # erf gelu, not tanh
    logits = head[i] @ state              # [49152]
    token  = argmax(logits)               # greedy; the caller may take
                                          #   the top-k instead
    draft.append(token)

emb_dim (6144) and inner_dim (4096) differ, and that is not a contradiction. emb_dim is the base model's hidden width — what proj[0] reads — and inner_dim is this head's own, which every stage after the first runs at.

From granite-3.0-8b-instruct-accelerator

This checkpoint only — its arithmetic differs from the others of this shape.

# A draft head, not a model. It proposes the next few tokens so the base
# model can verify several in one pass. There is no attention here, no
# position and no cache: the head sees one hidden state and one token,
# and everything it knows about the sentence arrived inside that state.

# state  = the base model's last hidden row      [4096]
# token  = the id the base model just emitted

state = rmsnorm(state, no_weight, no_bias) / sqrt(2)
# `scale_input: true`. The norm here has no parameters at all
# (`elementwise_shift=False, elementwise_scale=False`), so its
# absence from the checkpoint is correct rather than missing, and
# the / sqrt(2) is outside it — a linear map does not normalise, so
# the factor survives into the first projection.

state_weight = 0.5 ** (0.5 / 4)   # = 0.917004
emb_weight   = sqrt((1 - state_weight**2) * 4096 / 2)   # = 18.0512
alpha        = emb_weight / state_weight                # = 19.6849
# Both are derived from the config and neither is a tensor. They are
# chosen so the original state still accounts for half the magnitude
# by the final head. Only the ratio is applied: upstream writes
# `add(state, z, alpha=emb_weight/state_weight)` and lets the norm
# below absorb the missing state_weight — exact up to its epsilon.

for i in 0 .. 3:
    state = proj[i] @ state               # [4096]; proj[0] reads 4096, every
                                          #   stage after it is square
    state = state + alpha * emb[i][token]
    state = rmsnorm(state, ln[i].weight, ln[i].bias)
    #  ^ not a LayerNorm despite the name. `LayerNormParameterized` runs
    #    here with its mean subtraction off: x * rsqrt(mean(x^2) + eps),
    #    then weight * x + bias. It is an RMS norm with a bias.
    state = gelu(state)                   # erf gelu, not tanh
    logits = head[i] @ state              # [49155]
    token  = argmax(logits)               # greedy; the caller may take
                                          #   the top-k instead
    draft.append(token)

tie_weights ships as aliases. emb[i], head[i] and ln[i] all point at index 0 and proj[i>=2] at proj[1], which is why these checkpoints publish pytorch_model.bin rather than safetensors — safetensors cannot express two names for one buffer.

Notes that apply to more than one block

Stated once here rather than under each block above.

This head predicts 4 tokens ahead, one per stage, and each stage feeds the next — so stage i drafts on top of what stage i-1 proposed. Nothing is verified here; the base model checks every proposed token and keeps the longest correct prefix.

A wrong draft is not a wrong answer. The base model verifies every token, so this head's arithmetic cannot change what is emitted — only how many tokens get accepted per pass. That is why f16 is fine here where it would not be in a tower. It is not a licence to skip the reference: "it can only cost speed" and "it costs no speed because it is wrong" are the same observation from different ends, and the failure shows up only as an acceptance rate nobody has a prior for.

The norm is an RMS norm with a bias, not a LayerNorm. Upstream's LayerNormParameterized is instantiated with its mean subtraction off. Subtracting the mean loads, runs, and drafts slightly wrong tokens — visible only as a lower acceptance rate, which is the least detectable failure this head can have.

state_weight and emb_weight are derived from n_predict and inner_dim, not stored — 0.917004 and 18.0512 here — and only their ratio (19.6849) is applied, because the following norm absorbs the rest. A port that scales both terms separately is the same function up to that norm's epsilon and is not the published arithmetic.

top_k_tokens_per_head: [4, 3, 2, 2] with n_candidates: 4 is the search the head was trained for: stage 0 keeps its top 4, and the branching narrows with depth. Greedy argmax at every stage is one path through that tree and drafts fewer accepted tokens rather than wrong ones.

Geometry

Other shapes of MLPSpeculatorPreTrainedModel

Only the columns that differ are shown; a cell with several values means the checkpoints of that shape disagree.

How this was checked

implemented, evidence grade own-pack, per the assessment, for each checkpoint of this shape: a converted weight file for this checkpoint exists, but no forward was recorded from it. The implementation meant here and below is an unpublished independent inference implementation by the author.

What stands behind the block above, beyond the published configuration it is read from:

• a draft head implemented from the published configuration
• a checkpoint loader, which is where published tensor names are read
• an independent reference implementation of this architecture in Python, driving the published modelling code — an executable statement of what the model should compute, written against the publication rather than against any one implementation of it

Where an independent implementation and the published configuration disagree, the configuration is what this page reports and the disagreement is what it says.

Checkpoints with this architecture