GraniteSpeechNarForASR

GraniteSpeechNarForASR 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-nar.

# Mel features in, tokens out, with no autoregressive loop.

x = input_linear(features)                # [T, 1024], features 160 wide
states = [x]                              # entry 0 is input_linear's output
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_layer: 8`)
        p_mid   = softmax(out(x))         # 348-way, blank at index 0
        p_blank = p_mid[:, 0]
        x = x + out_mid(p_mid)
    states.append(x)                      # captured after the detour

#   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

# 1. The draft: a second CTC head, over the BPE vocabulary, reading the
#    final encoder state pooled in windows of 4 frames with weights
#    1 - p_blank from the mid-stack head.
pooled = posterior_pool(x, weights=1 - p_blank, window=4)   # `bpe_pooling_window`
draft  = argmax(pooled @ out_bpe.T)       # [., 100352]
draft  = drop(unique_consecutive(draft), 100257)   # collapse, then drop blank

# 2. The projector. Taps are `encoder_layer_indices: [4, 8, 12, -1]` into
#    `states`, so -1 is the last block; each tap has its own LayerNorm.
taps = [layernorm_j(states[k]) for j, k in enumerate([4, 8, 12, -1])]
h    = gelu(layer_projector(concat(taps, dim=-1)))   # 4096 -> 2048
for each window of 15 rows of h (tail zero-padded):
    q = query + mean_pool(window, 5)      # [3, 2048]: the learned query
                                          #   plus the window's own means
    kv = window + window_positions        # positions go into the keys and
                                          #   values, not the queries
    for layer in 0 .. 1:                  # cross-attention only, pre-norm
        q = q + cross_attn(layernorm(q), kv)   # 32 heads, with bias
        q = q + mlp(layernorm(q))              # silu, 2048 -> 4096 -> 2048
rows = out_linear(layernorm(q))           # 3 rows per window
rows = rows[: frames // 5]                # only the rows the real audio covers
rows = rows / 12                          # `scale_projected_embeddings`: the
                                          #   decoder multiplies by 12 again

# 3. The decoder edits the draft, once, bidirectionally.
slots = [blank, draft[0], blank, draft[1], ..., blank]   # 2n+1, at least 8
x = concat(rows, embed[slots])            # blank = 100257

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)                        # no causal mask: every row sees
                                          #   every other row of its sample
    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[slot_rows] @ embed.T           # tied (`tie_word_embeddings: true`)
logits = logits / 8                       # logits_scaling divides
tokens = drop(unique_consecutive(argmax(logits)), 100257)

Sources, all in the checkpoint's own modeling_granite_speech_nar.py at the pinned revision: the encoder, self-conditioning and BPE head; _posterior_weighted_pool; the projector and its Q-Former layer; the draft decode, slots and flat inputs; GraniteSpeechNarForASR.forward; the bidirectional decoder and attention; the final decode.

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.

There is a decoder and it is the stack this page describes. The BPE CTC head produces a draft, not the transcript: the draft is interleaved with blanks into editing slots, the projector's audio rows are prepended, and the whole sequence goes through the decoder once with no causal mask (packed samples are kept apart). A slot in the middle of the draft can use the audio that comes after it. The answer is the collapse of the argmax over the slot rows only.

The audio rows are divided by embedding_multiplier on the way in, because the decoder's embedding path multiplies every input row, audio included, by 12.

Its projector is not the 4.1 Q-Former. It reads four encoder taps, each with its own LayerNorm, has no self-attention, is pre-norm, starts its queries from the learned parameter plus the window's mean, and adds its learned position table to the encoder rows it attends over. Every one of those has the right shape the other way round.

Two CTC heads with different vocabularies and different blank ids. The mid-stack head's blank is index 0 of its 348-way output; the BPE head's and the decoder's blank is 100257 of 100,352.

Both CTC decodes collapse repeats first and drop blanks second. The other order merges two real repeats that a blank was separating.

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-nar: 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.