RobertaForMaskedLM

RobertaForMaskedLM at 6 layers and hidden size 384. One published checkpoint has this shape.

Implementing it

How to read this sheet: method.

Generated from granite-embedding-30m-sparse.

x = word[ids] + position[pos] + token_type[0]   # [T, 384]
x = layernorm(x, embeddings.LayerNorm, eps=1e-12)
# pos = padding_idx + 1 + i for real tokens (`pad_token_id: 1`), and
#   padding_idx for padding: real tokens start at row 2 (see Notes).

for i in 0 .. 5:
    q = x @ layers[i].attention.self.query.T + layers[i].attention.self.query.bias
    k = x @ layers[i].attention.self.key.T   + layers[i].attention.self.key.bias
    v = x @ layers[i].attention.self.value.T + layers[i].attention.self.value.bias
    a = softmax(q @ k.T / sqrt(32), mask=padding) @ v     # bidirectional
    o = a @ layers[i].attention.output.dense.T + layers[i].attention.output.dense.bias
    x = layernorm(x + o, layers[i].attention.output.LayerNorm, eps=1e-12)   # post-norm
    h = gelu(x @ layers[i].intermediate.dense.T + layers[i].intermediate.dense.bias)   # [T, 1536]
    o = h @ layers[i].output.dense.T + layers[i].output.dense.bias
    x = layernorm(x + o, layers[i].output.LayerNorm, eps=1e-12)

# The masked-LM head, on every position:
h = gelu(x @ lm_head.dense.T + lm_head.dense.bias)
h = layernorm(h, lm_head.layer_norm, eps=1e-12)
logits = h @ word.T + lm_head.bias        # [T, 50265]; the decoder is tied to
                                          #   the word embedding
# The sparse vector (sentence-transformers `SpladePooling`, max / relu):
w = max over positions t of  log(1 + relu(logits[t] * mask[t]))   # [50265]

Notes

The head is dense → gelu → LayerNorm → decoder, not a bare projection (modeling_roberta.py L1108-L1127). The checkpoint ships lm_head.dense, lm_head.layer_norm and lm_head.bias but no decoder weight, which is tied to roberta.embeddings.word_embeddings.weight (L1017).

The output is one weight per vocabulary entry. The repository's sentence-transformers pipeline is MLMTransformer then SpladePooling with pooling_strategy: max and activation_function: relu (modules.json, 1_SpladePooling/config.json); SpladePooling masks padding, applies relu then log1p, and takes the maximum over positions (SpladePooling.py L96-L117 at sentence-transformers v5.0.0).

Position ids start at padding_idx + 1 = 2. create_position_ids_from_input_ids numbers real tokens from padding_idx + 1 and gives padding the id padding_idx (modeling_roberta.py L1537-L1550). With pad_token_id: 1, row 0 of the 514-row position table is never read, row 1 only at padding positions, and a 512-token input reads rows 2 to 513, so the 514-row table (max_position_embeddings) holds 512 usable positions. A port that starts at row 0 looks everything up successfully and shifts every position by 2.

The attention is bidirectional: only padding is masked, so every position reads the whole sequence.

This lineage is post-norm: each sublayer's LayerNorm comes after the residual add, the opposite of every decoder on these pages.

Geometry

Weight structure

The tensors one element of the repeating stack holds, by the names the published checkpoint uses.

How this was checked

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

• granite-embedding-30m-sparse: runs, recorded-forward — this checkpoint was loaded and completed a forward pass

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

• an encoder 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

Recorded pass: passes.