GraniteSpeechPlusForConditionalGeneration

GraniteSpeechPlusForConditionalGeneration at 40 layers and hidden size 2048. One published checkpoint has 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.

Generated from granite-speech-4.1-2b-plus.

# 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 == 2:                            # `cat_hidden_layers: [3]` counts from 1:
        captured.append(x)                #   the 3rd block's output
    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    348   `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

# Plus: the projector reads more than the final residual. The captured
# state is concatenated with the final one, so the Q-Former's encoder
# input is 2 x 1024 wide (`encoder_hidden_size: 2048`), not 1024.
enc = concat(captured + [x], dim=-1)      # [T, 2048]

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 # 16 q heads, 4 kv
    q, k = rope(q, k, theta=10000)     # 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 / 8                       # logits_scaling divides

Sources: the encoder and projector follow plain Granite Speech, whose sources apply; the capture and concatenation are GraniteSpeechPlusCTCEncoder.forward in transformers v5.8.0, the first release with this class.

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.

What "plus" changes against granite-speech-4.1-2b. cat_hidden_layers: [3] — the 3rd block's output, counting from 1 — is concatenated with the final residual, so the projector's input is twice as wide. The capture is taken before the self-conditioning rejoin (which happens after block 8). The configurations also differ in one more place: plus ties its output head to the input embedding (tie_word_embeddings: true), where plain 4.1 ships a separate lm_head. That 205.5M-parameter matrix, less the roughly 4M the wider Q-Former cross-attention adds, is the 201M by which plus is the smaller checkpoint.

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-speech-4.1-2b-plus: runs, shape-pack — stored checkpoints of this shape run; this one's own is not among them

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.