GraniteSpeechNarForASR

GraniteSpeechNarForASR at 40 layers and hidden size 2048. One published checkpoint has this shape.

The pass (block level)

The order of operations for one forward. ×N marks a position that fires once per layer.

6 steps, one level. Enough to name the steps, not enough to show what nests inside what. Recorded from a complete forward of this shape in the author's independent implementation — the structure is what that run emitted, not a reading of the configuration.

block
├─ ff1
├─ attn
├─ conv
├─ ff2
└─ post_ln

The same thing in canonical form:

0(1 2 3 4 5)

Numbers are positions in the pass, not layer indices. Two passes count as the same structure when these strings match.

Implementing it

One forward, written out, per checkpoint whose arithmetic actually differs. A shape groups checkpoints by class, depth, width and layer stack — none of which decides an activation, a rope base or a routing rule — so where the members of this shape disagree there is a block each, and each names the checkpoint it was generated from.

It is not the recorded pass above, which is captured from a real forward; it is what the published configuration says the arithmetic is. Where the configuration does not say, the line says that instead of guessing.

The names are the ones the published checkpoint uses where a name is published, and canonical otherwise. A checkpoint loader may store them differently — fusing a gate/up pair into one matrix, or splitting one published projection in two — and those are that loader's names, not the model's. Everything spelled here is what the download contains.

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

# Mel features in, tokens out. The encoder is a conformer; what sits on
# top of it is what distinguishes these four 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(prenorm(x))         # half-weight feed-forward
    # Blocked, not full: frames split into non-overlapping 200-frame
    # windows, attention is block-diagonal, no window reads another.
    for w in windows(x, 200):           # each window separately
        b = q @ k.T / sqrt(head_dim) + shaw_rel_pos(rel_pos_emb, w)
        w[:] = softmax(b) @ v             # Shaw relative positions:
                                          #   a learned bias per offset,
                                          #   added to the scores. No
                                          #   rope anywhere in this tower.
    x = x + to_out(attn_blocked(prenorm(x)))
    x = x + conv(x)                       # LayerNorm, pointwise, depthwise,
                                          #   BatchNorm, pointwise
    x = x + 0.5 * ff2(prenorm(x))         # the second half-weight FFN
    x = post_norm(x)
    if i in [4, 8, 12, 15]:
        captured.append(copy(x))      # NAR captures after the detour
    if i == 8:                          # self-conditioning, once, mid-stack
        x = x + out_mid(softmax(out(x)))  # the CTC head's own opinion,
                                          # fed back into the residual

#   The conformer tower, as published. Every line here is a tensor shape a
#   port allocates, and none of it is stated anywhere else on this page.
#     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`

# 1. Taps, captured after the self-conditioning detour, not before.
#    `encoder_layer_indices` indexes a tuple whose entry 0 is the
#    input_linear output and entry k is block k's output, so -1 is
#    the last block and has to be resolved against the layer count.
#    4.1-plus captures before the detour and index 8 is the detour's,
#    so the difference lands on one of the taps and is real.
taps = [x_after_block[k] for k in encoder_layer_indices]   # [4, 8, 12, -1]
h    = layer_projector(concat(taps, dim=-1))

# 2. The NAR projector — cross-attention only. No self-attention.
q = learned_query + mean_pool(window)     # queries are the learned
                                          # parameter plus the window's
                                          # own mean, not the parameter
                                          # alone
k = h + window_positions[block]           # the learned position table
                                          # is added to the keys, not
                                          # the queries. Both have the
                                          # right shape either way.
q = prenorm_cross_attn(q, k); q = silu_mlp(q)

#   The NAR projector, as published. Every line here is a tensor shape a
#   port allocates, and none of it is stated anywhere else on this page.
#     layers         2   `num_layers`
#     width       2048   `hidden_size`
#     heads         32   `num_heads`
#     downsample     5   `downsample_rate`
#     mlp bias    True   `mlp_bias`
# The last window is zero-padded to a full block, not truncated, so
# it still emits its full row count; the caller keeps only the rows
# the real audio covers (252 produced, 249 kept on one clip here).

# 3. A second CTC head over a BPE vocabulary — and what comes out of
#    it is a draft, not the answer. The pool is posterior-weighted:
#    each window of 4 frames is averaged with weights 1 - P(blank)
#    taken from the mid-stack CTC head.
rows   = posterior_pool(q, weights=1 - P_blank)
draft  = collapse_repeats(drop_blanks(argmax(rows @ out_bpe.T)))   # [., 100352]

# 4. The decoder, and it is the stack this page is about. The draft
#    is interleaved with blanks to make editing slots, the audio
#    rows are prepended, and the whole thing goes through the
#    40-layer decoder once, bidirectionally — every row attends
#    over every other, so a slot in the middle of the draft reads the
#    audio that follows it. Nothing is sampled and no KV is reused.
slots = [blank, draft[0], blank, draft[1], ..., blank]   # 2n+1, min 8, blank=100257
x     = concat(audio_rows / embedding_multiplier, embed[slots])
                                          # the projector's rows were
                                          # already scaled; the decoder
                                          # scales again, so divide out

# The decoder is an ordinary dense Granite, and its four multipliers are
# not optional. They are plain floats with no analogue in a Llama
# config, they apply at different points in the pass, and a port
# that drops them generates fluent text that is not this model's.
#
# And it runs once, bidirectionally. No causal mask, no KV reuse and
#   nothing sampled — the mask is the only difference from the loop
#   below, and a causal pass over the same rows reads just as
#   fluently while being unable to do this model's job.

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: the norm feeds the
                                          #   block, the residual below
                                          #   carries the unnormalised x
    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, not
                                          #   after. Rotating the output
                                          #   of attention is a model
                                          #   whose positions do nothing.
    s = q @ k.T * 0.0078125               # attention_multiplier replaces
                                          #   1/sqrt(head_dim) — it is a
                                          #   substitute, not a factor
                                          #   applied beside it
    a = softmax(s)                        # no mask — every row sees
                                          #   every other row
    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)                # pre-norm stacks normalise once
                                          #   more at the end. Without it
                                          #   the head reads a residual
                                          #   stream nothing scaled.
logits = x[the_slot_rows] @ embed.T       # the audio rows carry no text
                                          #   Tied to the input embedding
                                          #   (`tie_word_embeddings: true`)
logits = logits / 8                       # logits_scaling divides.
                                          #   Multiplying has the same
                                          #   shape and flattens or
                                          #   sharpens every sample.
tokens = collapse_repeats(drop_blanks(argmax(logits)))

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 is a learned bias per offset, added to the scores. There is no rope in this tower at all, and a port that runs full attention with rope is more expensive, transcribes, and is not this model.

The self-conditioning detour fires once, after block 8 of 16: the mid-stack CTC head's distribution is projected back into the residual stream. A tower that skips it runs and transcribes worse.

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. That is the point — a slot in the middle of the draft is meant to use the audio that comes after it, and a causal pass over the same rows reads just as fluently while being unable to. The answer is the argmax over the slot rows only.

The audio rows are divided by embedding_multiplier on the way in, because the projector already produced them at the decoder's scale and the decoder's embedding path multiplies again. Skipping the division leaves a model that still transcribes.

Its projector is not the 4.1 Q-Former. It has no self-attention, it is pre-norm where 4.1 is post-norm, its queries are the learned parameter plus the window's mean rather than the parameter alone, and its learned position table is added to the keys. Every one of those has the right shape the other way round.

Two CTC heads are live in one forward with different vocabularies and different blank ids — the mid-stack head's blank is index 0 of its own narrow vocabulary, the BPE head's is 100257 of 100352. Dropping the wrong one leaves a transcript that still reads.

Geometry

Weight structure

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

How this was checked

implemented, evidence grade shape-pack, per the assessment: converted checkpoints of this shape run; this checkpoint's own weights are 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