GraniteSpeechForConditionalGeneration

GraniteSpeechForConditionalGeneration at 40 layers and hidden size 4096. 2 published checkpoints share this shape.

Implementing it

How to read this sheet: method.

This section is not the recorded pass (that is on passes); it is what the published configuration and the modelling code say the arithmetic is.

From granite-speech-3.3-8b

# Mel features in, tokens out. The encoder is a conformer; what sits on
# top of it is what distinguishes the speech classes from each other.

x = input_linear(features)                # [T, 1024]
for i in 0 .. 15:                         # conformer block, five modules in this order
    x = x + 0.5 * ff1(x)                  # pre_norm -> up_proj -> silu -> down_proj
    x = x + to_out(attn_blocked(pre_norm(x)))   # defined below
    x = x + conv(x)                       # LayerNorm, pointwise, GLU,
                                          #   depthwise, BatchNorm, silu,
                                          #   pointwise
    x = x + 0.5 * ff2(x)                  # the second half-weight FFN
    x = post_norm(x)
    if i == 7:                          # after block 8 of 16: self-conditioning
        x = x + out_mid(softmax(out(x)))  # the CTC head's own distribution,
                                          # fed back into the residual

#   The conformer tower, as published:
#     layers         16   `num_layers`
#     width        1024   `hidden_dim`
#     heads           8   `num_heads`
#     head dim      128   `dim_head`
#     ffn mult        4   `feedforward_mult`
#     conv kernel    15   `conv_kernel_size`
#     output dim    256   `output_dim`

# attn_blocked, used once per block above. Blocked, not full: frames are
# split into non-overlapping 200-frame windows (the tail is zero-padded
# and masked), attention is block-diagonal, no window reads another.
def attn_blocked(h):
    out = empty_like(h)
    for w, o in zip(windows(h, 200), windows(out, 200)):
        q = to_q(w); k, v = split(to_kv(w), 2)
        b = (q @ k.T + q @ rel_pos_emb[dist].T) / sqrt(head_dim)
        o[:] = softmax(b) @ v             # Shaw relative positions: a
                                          #   learned vector per clamped
                                          #   offset, dotted with q and
                                          #   added to the scores. No rope
                                          #   anywhere in this tower.
    return out

enc = x                                   # the final residual, only

for each window of 15 encoder rows:        # windowed Q-Former, BLIP-2 block
    q = layernorm(query)                  # 3 learned queries [3, 1024], reused
                                          #   for every window, LayerNorm'd once
                                          #   before layer 0
    for layer in 0 .. 1:
        q = self_attn(q); q = cross_attn(q, window); q = ffn(q)   # post-norm
# The tail window is zero-padded to a full window, not truncated, and
# the processor sizes the audio placeholder run for the padded count.
rows = linear(q)                          # 3 rows per window, decoder-wide

#   The Q-Former projector, as published:
#     layers             2   `num_hidden_layers`
#     width           1024   `hidden_size`
#     heads             16   `num_attention_heads`
#     ffn width       4096   `intermediate_size`
#     activation      gelu   `hidden_act`
#     norm eps       1e-12   `layer_norm_eps`

x = embed[ids]; x[at_the_audio_placeholder] = rows

# The decoder is an ordinary dense Granite, and its four multipliers
# apply at different points in the pass.

x = x * 12                                # embedding_multiplier, once, on the way in
for i in 0 .. 39:
    h = rmsnorm(x, layers[i].input_layernorm, eps=1e-05) # pre-norm
    q, k, v = h @ layers[i].self_attn.q_proj.T, h @ layers[i].self_attn.k_proj.T, h @ layers[i].self_attn.v_proj.T # 32 q heads, 8 kv
    q, k = rope(q, k, theta=10000000)     # before the scores
    s = q @ k.T * 0.0078125               # attention_multiplier replaces
                                          #   1/sqrt(head_dim), not a factor
                                          #   applied beside it
    a = softmax(s, mask=causal)
    o = (a @ v) @ layers[i].self_attn.o_proj.T
    x = x + 0.22 * o                      # residual_multiplier on both
    h = rmsnorm(x, layers[i].post_attention_layernorm)
    m_ = (silu(h @ layers[i].mlp.gate_proj.T) * (h @ layers[i].mlp.up_proj.T)) @ layers[i].mlp.down_proj.T
    x = x + 0.22 * m_                     #   ...and on the MLP add too

x = rmsnorm(x, model.norm)
logits = x @ embed.T                      # tied (`tie_word_embeddings: true`)
logits = logits / 16                       # logits_scaling divides

granite-speech-3.2-8b: differs from the block above

Same arithmetic; only these values change.

Sources, transformers v4.57.1: the conformer block, its feed-forward, block attention with Shaw positions and conv module; the encoder loop and self-conditioning; the windowed Q-Former projector and the BLIP-2 Q-Former's LayerNorm on the queries before layer 0; the placeholder count; the Granite decoder's attention_multiplier, residual_multiplier, embedding_multiplier and logits_scaling.

Notes that apply to more than one block

The conformer's five modules run in that order and the two feed-forwards are half-weighted — each contributes 0.5 * to the residual. That factor is structural, carries no tensor, and is invisible in a weight listing.

The tower never runs a full attention. Frames are split into non-overlapping context_size: 200-frame windows and attention is block-diagonal — no window reads another. Inside a window the positions are Shaw relative: attn.rel_pos_emb holds a learned vector per clamped offset, dotted with the query and added to the scores. There is no rope in this tower; full attention with rope is a different function.

Self-conditioning fires once, after the middle block (num_layers // 2, counting from 1): the mid-stack CTC head's distribution is projected back into the residual stream through out_mid. A tower that skips it is a different function.

The Q-Former's last window is zero-padded to a full window rather than truncated, and the processor counts ceil(frames / 15) * 3 audio placeholders, so every row the projector emits — the padded window's included — has a slot. Truncating instead loses the end of every clip that does not divide evenly.

Geometry

Weight structure

The tensors one element of the repeating stack holds, by the names the published checkpoints use. The two checkpoints of this shape name their encoder tensors differently.

In conv.net.# and ff1.fn.fn.net.# the # is a sub-module index inside one layer, not a layer number.

granite-speech-3.2-8b names its encoder tensors differently from its class. transformers v4.57.1 builds this class's encoder as input_linear and layers (L262–L263) and declares no key mapping, so by name none of the 3.2-8b encoder tensors has a parameter to load into. The arithmetic above holds for both, and the recorded forward comes from the author's implementation, which reads the published names; a standard transformers load of this checkpoint was not run here (findings §10).

Other shapes of GraniteSpeechForConditionalGeneration

Identical across all of them: layers 40.

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 2 checkpoints of this shape: runs, recorded-forward — this checkpoint was loaded and completed a forward pass

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 conformer tower and projector 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

Recorded pass: passes.