nnterp · encyclopedia doge · nnterp.families.doge

family doge

Doge

Llama's pre-norm block whose residual stream is multiplied by a learned per-channel gate before each add, and whose attention adds a learned bias per key, computed from the values, to the scores.

Scaled residual addsQuery/key normsDynamic maskGrouped-query attention

The block

01 · the residual stream

One block of the stack. Hover any part for its nnterp name; click to pin the card.

embed_tokenstoken_embeddings
→
layers× 24
→
normfinal_norm_output
→
lm_headlm_head_output
→
logitsthe output
i = 0

identity post_attention_residual * (input_residual * layers[i].input + self_attn.attention_output) + mlp.mlp_output == layer_output

The API

02 · support()

Values, by host

⚠︎ marks a value this checkpoint has only under some load: hover it for the condition.

model DogeForCausalLMlayoutwhat it isread at
logits Logits [batch seq vocab_size] The logits the model returns, after anything it does past lm_head (a softcap, a scale)

the module's own output
token_embeddings Residual [batch seq hidden_size] The embedding module's output; layers[0].input is what enters the first block

the child module embed_tokens's output
final_norm_output Residual [batch seq hidden_size] The residual stream after the final norm, what lm_head reads

the child module norm's output
lm_head_output Logits [batch seq vocab_size] The head's logits before anything the model does past lm_head (a softcap, a scale)

the child module lm_head's output
next_token_probs NextTokenProbs [batch vocab_size] The next-token distribution at the last position; derived, read-only

the module's own output
input_ids Tokens [batch seq] The token ids the model was called with

the module's own input
attention_mask Tokens [batch seq] The attention mask the model was called with; zeros are padding

the module's own input
input_size scalar [batch, seq] of the current call; read-only

the module's own input
model.layers[i] DogeDecoderLayerlayoutwhat it isread at
layer_input Residual [batch seq hidden_size] The residual stream entering the block

the module's own input
attention_norm_output Residual [batch seq hidden_size] The normed stream entering the attention

the child module input_layernorm's output
layer_mid Residual [batch seq hidden_size] The residual stream after the attention add

the child module post_attention_layernorm's input
mlp_norm_output Residual [batch seq hidden_size] The normed stream entering the MLP

the child module post_attention_layernorm's output
mlp_input Residual [batch seq hidden_size] What the MLP is called with

the child module mlp's input
layer_output Residual [batch seq hidden_size] The residual stream leaving the block, a tensor even when the block returns a tuple

the module's own output
model.layers[i].self_attn DogeAttentionlayoutwhat it isread at
⚠︎attention_queries Queries [batch num_heads seq qk_head_dim] The queries entering attention

inside the forward: the arguments of attention_interface_1, element 1
⚠︎attention_keys Keys [batch num_kv_heads seq qk_head_dim] The keys entering attention

inside the forward: the arguments of attention_interface_1, element 2
⚠︎attention_values Values [batch num_kv_heads seq head_dim] The values entering attention

inside the forward: the arguments of attention_interface_1, element 3
⚠︎attention_scores Pattern [batch num_heads query key] The attention scores entering the softmax, masked

inside the forward: the input of nn_functional_softmax_0 inside attention_interface_1
attention_premix HeadOutputs [batch seq num_heads head_dim] Every head's output side by side, as the output projection takes it

the child module o_proj's input
attention_output Residual [batch seq hidden_size] What the attention adds to the residual stream

the module's own output
⚠︎attention_probabilities Pattern [batch num_heads query key] The attention pattern the values are mixed with

inside the forward: the output of nn_functional_dropout_0 inside attention_interface_1
⚠︎attention_head_outputs HeadOutputs [batch seq num_heads head_dim] The per-head outputs before the output projection

inside the forward: the output of attention_interface_1, element 0
model.layers[i].mlp DogeMLPlayoutwhat it isread at
mlp_activation Neurons [batch seq intermediate_size] The activation function's output

the child module act_fn's output
mlp_neurons Neurons [batch seq intermediate_size] What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise

the child module down_proj's input
mlp_output Residual [batch seq hidden_size] What the MLP adds to the residual stream

the module's own output

Printout

What print(model) shows on SmallDoge/Doge-260M: the native modules, standard names and every standard value with its layout.

DogeForCausalLM(
  (lm_head): Linear(in_features=768, out_features=128256, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Embedding(128256, 768, padding_idx=128001)
  (layers): ModuleList(
    (0-23): 24 x DogeDecoderLayer(
      (input_layernorm): DogeRMSNorm((768,), eps=1e-06)
      (self_attn): DogeAttention(
        (q_proj): Linear(in_features=768, out_features=768, bias=False)
        (k_proj): Linear(in_features=768, out_features=384, bias=False)
        (v_proj): Linear(in_features=768, out_features=384, bias=False)
        (dt_proj): Linear(in_features=384, out_features=3, bias=False)
        (o_proj): Linear(in_features=768, out_features=768, bias=False)
        (q_norm): DogeRMSNorm((128,), eps=1e-06)
        (k_norm): DogeRMSNorm((128,), eps=1e-06)
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_premix) -> HeadOutputs [batch seq num_heads head_dim]: Every head's output side by side, as the output projection takes it
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
      )
      (post_attention_layernorm): DogeRMSNorm((768,), eps=1e-06)
      (mlp): DogeMLP(
        (gate_proj): Linear(in_features=768, out_features=2304, bias=False)
        (up_proj): Linear(in_features=768, out_features=2304, bias=False)
        (down_proj): Linear(in_features=2304, out_features=768, bias=False)
        (act_fn): SiLUActivation()
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
  )
  (norm): DogeRMSNorm((768,), eps=1e-06)
  (logits) -> Logits [batch seq vocab_size]: The logits the model returns, after anything it does past lm_head (a softcap, a scale)
  (token_embeddings) -> Residual [batch seq hidden_size]: The embedding module's output; layers[0].input is what enters the first block
  (final_norm_output) -> Residual [batch seq hidden_size]: The residual stream after the final norm, what lm_head reads
  (lm_head_output) -> Logits [batch seq vocab_size]: The head's logits before anything the model does past lm_head (a softcap, a scale)
  (next_token_probs) -> NextTokenProbs [batch vocab_size]: The next-token distribution at the last position; derived, read-only
  (input_ids) -> Tokens [batch seq]: The token ids the model was called with
  (attention_mask) -> Tokens [batch seq]: The attention mask the model was called with; zeros are padding
  (input_size): [batch, seq] of the current call; read-only
)

Config

model.num_layers24
model.hidden_size768
model.num_heads6
model.num_kv_heads3
model.head_dim128
model.qk_head_dim128
model.vocab_size128256
model.intermediate_size2304
architecturesDogeForCausalLM
num_hidden_layers24
hidden_size768
intermediate_size2304
num_attention_heads6
num_key_value_heads3
vocab_size128256
max_position_embeddings4096
hidden_actsilu
rms_norm_eps1e-06
tie_word_embeddingsTrue
dtypebfloat16
num_experts16384
num_experts_per_tok64
norm_topk_probFalse
attention_biasFalse

The Notes

03 · what to know before you run

The block, in order

h   = input_residual * x + self_attn(input_layernorm(x))
out = post_attention_residual * h + mlp(post_attention_layernorm(h))

input_residual and post_attention_residual are [hidden_size] parameters of the block. They multiply the stream the block carries, not what the sublayers add: attention_output is o_proj's output and mlp_output the MLP's, unscaled.

The stream is gated, so the identity carries the gates

The plain sum layers[i].input + attention_output + mlp_output does not give layer_output; the identity multiplies the stream by each gate where the block does. It is exact in float32 on every Doge-260M block:

with model.trace(prompt):
    x = model.layers[1].input.save()
    attn = model.layers[1].self_attn.attention_output.save()
    mlp = model.layers[1].mlp.mlp_output.save()
    out = model.layers[1].layer_output.save()

block = model.layers[1]._module
h = block.input_residual * x + attn
torch.testing.assert_close(block.post_attention_residual * h + mlp, out)

On Doge-260M the gates are below one almost everywhere: input_residual averages 0.73 to 0.94 per block (its few negative entries, 0.09% of them, are on block 0), and post_attention_residual 0.87 to 0.96. A term added at block i therefore reaches the last block's output multiplied by every later gate, channel by channel, which a direct attribution to the logits has to include. Zeroing attention_output does not leave the block's input in place: the stream still goes on as input_residual * x.

The last block erases a massive channel

On Doge-260M channel 524 of the stream is a massive activation: at the last block's input it holds 57,662 at the first position and 5,000 to 6,700 at the others, at least 85% of the stream's squared norm at every position of a 30-token prompt. The last block's post_attention_residual is 0.0008 on that channel and at least 0.72 on every other, so the block removes it before the final norm. Steering or patching that adds to channel 524 is erased the same way at the last block.

The dynamic mask is a bias per key, computed from the values

Each attention projects its values through dt_proj to one number per key and key/value head, and adds exp(A * softplus(dt)) to every query's score for that key. A is a per-head parameter, so a head with A > 0 raises the keys dt_proj picks out and a head with A < 0 lowers them. attention_scores carry the bias, and it can be rebuilt from attention_values:

import torch.nn.functional as F

attn = model.layers[1].self_attn
with model.trace(prompt):
    values = attn.attention_values.save()          # [batch, num_kv_heads, seq, head_dim]
    scores = attn.attention_scores.save()          # [batch, num_heads, query, key]

module = attn._module
v = values.transpose(1, 2).flatten(2)              # [batch, seq, num_kv_heads * head_dim]
dt = F.softplus(module.dt_proj(v))                 # [batch, seq, num_kv_heads]
bias = torch.exp(module.A * dt).transpose(1, 2)    # [batch, num_kv_heads, key]

On Doge-260M A ranges from −0.29 to 0.29 and the bias from 0 to 8.4. Removing it from the scores moves 0.02 to 0.47 of each pattern row's mass, averaged per block, on a 30-token prompt; at block 10 the mass on the first token, averaged over heads and queries, is 0.65 with the bias and 0.34 without it. The pinned tiny checkpoint has A at zero, where the bias is 1 on every key and the pattern is unchanged.

Past keep_window_size keys, the bias drops keys

When the prompt has more than keep_window_size keys (1024 on Doge-260M, 2048 on Doge-20M, 4096 on Doge-40M), each query keeps the keep_window_size keys it can see with the largest bias and masks the rest. The choice reads the bias alone, not the query's score. A dropped key has the dtype's minimum in attention_scores and zero in attention_probabilities. Below that length, every causal key is kept.

Loading

The attention interior needs attn_implementation="eager". The default sdpa load adds the same bias and the causal mask: on Doge-260M in float32 the two loads' logits agree to 1.2e-5. Doge-20M, Doge-40M and Doge-260M load in transformers. In the Instruct checkpoints and in Doge-60M, -160M and -320M, A and dt_proj are twice num_key_value_heads wide and the attention has no q_norm or k_norm, so transformers' DogeForCausalLM does not load them (they are greyed out in the selector). Doge-260M's tokenizer prepends <|begin_of_text|> (id 128000).

The cross-domain mixture

On the is_moe checkpoints mlp is DogeCDMoE: a dense SwiGLU MLP of intermediate_size plus num_experts one-neuron experts (rows of down_embed and up_embed) that router_gate retrieves by product keys, num_experts_per_tok per token. It returns (hidden_states, router_logits), and the block adds the hidden states. nnterp serves mlp_output, the dense MLP's output plus the routed experts', and reports every mixture value unavailable: the product-key routing has no router or experts module to read them at.

The readout

logits is lm_head.output, with no cap or scale, and lm_head shares its weight with embed_tokens on every listed checkpoint.

From the family module

04 · nnterp/families/doge.py

Doge (DogeForCausalLM, SmallDoge).

Llama's tree with a gated residual: the block scales the stream it carries, not what it adds. Each sublayer's output is added to the residual times a learned per-channel gate (input_residual, then post_attention_residual, both [hidden_size] parameters of the block)::

h = input_residual * x + self_attn(input_layernorm(x))
out = post_attention_residual * h + mlp(post_attention_layernorm(h))

What each sublayer adds is its module's output, unscaled, so the base holds for attention_output and mlp_output; the contribution identity is layer_output == post_attention_residual * (input_residual * input + attention_output) + mlp_output, which is the plain one only where both gates are one (their initial value). The gates are layers[i]._module.input_residual / post_attention_residual; nothing is computed or divided for you. The attention runs the shared eager forward with q/k norms and a dynamic mask: dt_proj of the values, through A, gives each key a per-head bias, added to the causal mask (and, past keep_window_size keys, used to drop all but the top keys), so attention_scores carry it. The MLP is a gated DogeMLP, or the cross-domain mixture DogeCDMoE under is_moe, which returns (hidden_states, router_logits).

Doge-20M, -40M and -260M load and run. Of the is_moe checkpoints, Doge-40M-MoE loads and runs; it needs transformers 5.18 or later, whose block takes the first element of DogeCDMoE's tuple. Where it runs, mlp_output is the mixture's hidden states (dense MLP plus routed experts); the six mixture values stay unavailable, since the product-key routing has no router or experts module to read them at.