MLPSpeculatorPreTrainedModel

MLPSpeculatorPreTrainedModel at 5 prediction stages and stage width 4096. 3 published checkpoints share this shape.

Implementing it

How to read this sheet: method.

No forward was recorded for this shape; the block is read from the published configuration and from MLPSpeculator in fms-extras at commit 16339f7.

From granite-8b-code-instruct-accelerator

# A draft head, not a model: it proposes the next 5 tokens for the base
# model to verify. No attention, no positions, no cache: each stage sees
# one state vector and one token id.

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

state_weight = 0.5 ** (0.5 / 5)   # = 0.933033
emb_weight   = sqrt((1 - state_weight**2) * 4096 / 2)   # = 16.2823
alpha        = emb_weight / state_weight                # = 17.4509
# Derived from n_predict and inner_dim; neither is a tensor. Only the
# ratio is applied: the norm below absorbs the missing state_weight,
# up to the norm's eps (scaling x by c turns eps into eps / c**2).

for i in 0 .. 4:
    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, eps=1e-06)
    #  ^ `LayerNormParameterized` with its mean subtraction off:
    #    x * rsqrt(mean(x^2) + eps), then weight * x + bias.
    state = gelu(state)                   # exact (erf) GELU, not tanh
    logits = head[i] @ state              # [49152]
    token  = argmax(logits)               # greedy; see Notes for the
                                          #   candidate tree
    draft.append(token)

Differs from the block above

The same arithmetic with these values changed. Everything not listed is as above.

emb_dim and inner_dim can differ. On granite-3b-code-instruct-accelerator, emb_dim (2,560) is the base model's hidden width, which proj[0] reads, and inner_dim (4,096) is the head's own width, which every later stage runs at (speculator.py L62-L67).

Notes

This head drafts 5 tokens ahead, one per stage, each stage feeding the next. It verifies nothing: in fms-extras' greedy speculative_generate, the base model keeps the longest prefix of a candidate that matches its own next tokens (generation.py L242-L244). A wrong draft therefore costs accepted tokens per step, not correctness.

The norm is an RMS norm with a bias, not a LayerNorm. MLPSpeculator builds each ln[i] as LayerNormParameterized(inner_dim, elementwise_shift=True, elementwise_scale=True) (speculator.py L71-L78), whose defaults are use_mean=False, eps=1e-06 (fms layernorm.py L27-L34, L59-L72 at v0.0.6). Subtracting the mean would load and run, and draft different tokens. The activation is nn.GELU(), PyTorch's exact form (L86).

state_weight and emb_weight are computed, not stored (speculator.py L83-L85), and only their ratio is applied: upstream writes torch.add(state, z, alpha=self.emb_weight / self.state_weight) and relies on the following norm (L162-L165). A port that scales both terms separately differs only through that norm's epsilon, but it is not the published arithmetic.

tie_weights: true (granite-3b-code-instruct-accelerator) shares tensors across stages. emb[i], head[i] and ln[i] all point at stage 0's tensors and proj[i] for i ≥ 2 at proj[1] (speculator.py L88-L106). A loader has to resolve those aliases rather than expect one tensor per stage.

top_k_tokens_per_head: [4, 3, 2, 2, 2] with n_candidates: 5 configures generate_suffixes, which keeps the top k tokens at each stage and returns the 5 most probable candidate sequences (speculator.py L116-L187). Greedy argmax at every stage is one path through that tree.

Where the configuration omits scale_input or tie_weights, MLPSpeculatorConfig defaults both to false (modeling_mlp_speculator.py L25-L26); the table says where that default is being used.

Geometry

Weight structure

granite-8b-code-instruct-accelerator and granite-7b-instruct-accelerator publish safetensors holding five each of speculator.emb.#.weight, speculator.proj.#.weight, speculator.ln.#.weight, speculator.ln.#.bias and speculator.head.#.weight, one per stage, which is what the block reads. granite-3b-code-instruct-accelerator publishes only pytorch_model.bin, whose tensor list cannot be read without downloading the file, so its names, and the stage sharing that tie_weights: true implies, cannot be checked against the published headers.

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

Status and evidence grade in the assessment (see what the grades mean):

• each of the 3 checkpoints of this shape: runs, own-pack — a stored copy of this checkpoint exists; nothing recorded from it

Parts of the author's unpublished implementation that exist for this shape:

• 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

Checkpoints with this architecture