The baseline block

The previous posts built the pieces that turn text into vectors and let positions read each other. A decoder stacks the same block a few times on top of those vectors, and then turns the final vectors back into token scores. Day 1's baseline has two blocks, width 128 and four heads, and it follows the small GPT-2 layout.

The forward pass, with shapes

For a batch of $B$ sequences of $T$ token IDs, with the Day 1 settings $C = 128$ and $V = 1{,}024$:

Every block keeps the shape [B, T, C]. That's what lets blocks stack, the output of one is a valid input to the next.

def forward(self, token_ids: Tensor, all_depths: bool = False) -> Tensor:
    if token_ids.ndim != 2:
        raise ValueError("token IDs must have shape [batch, sequence]")
    length = token_ids.shape[1]
    if self.position_embedding is not None and length > self.config.context_length:
        raise ValueError("sequence exceeds context_length")
    rope = None
    if self.config.position == "rope":
        rope = rope_tables(
            length,
            self.config.head_size,
            self.config.rope_scale,
            token_ids.device,
            self.config.rope_base,
        )
    x = self.dropout(self._embed(token_ids))
    for block in self.blocks:
        x = block(x, rope)
    hidden = self.final_norm(x)
    logits = self.lm_head(hidden)
    if not all_depths:
        return logits
    return torch.stack([logits, *(self.lm_head(a(hidden)) for a in self.mtp_adapters)])

Each block adds to the vector, it doesn't replace it

A block doesn't compute a new vector from scratch. It computes a correction and adds it:

The running vector $x$ is called the residual stream. And the addition actually helps training. The gradient of $x + f(x)$ with respect to $x$ is $1 + f'(x)$, so every block passes at least the identity backward. A deep stack of additions gives the gradient a direct path from the loss to the embedding, and it doesn't shrink as you add more blocks.

class TransformerBlock(nn.Module):
    def __init__(self, config: DecoderConfig) -> None:
        super().__init__()
        self.pre_norm = config.residual == "pre"
        self.attention_norm = make_norm(config)
        self.attention = CausalSelfAttention(config)
        self.feed_forward_norm = make_norm(config)
        self.feed_forward = FeedForward(config)

    def forward(self, x: Tensor, rope: tuple[Tensor, Tensor] | None = None) -> Tensor:
        if self.pre_norm:
            x = x + self.attention(self.attention_norm(x), rope)
            return x + self.feed_forward(self.feed_forward_norm(x))
        x = self.attention_norm(x + self.attention(x, rope))
        return self.feed_forward_norm(x + self.feed_forward(x))

    def forward_cached(
        self,
        x: Tensor,
        rope: tuple[Tensor, Tensor] | None,
        cache: tuple[Tensor, Tensor] | None,
    ) -> tuple[Tensor, tuple[Tensor, Tensor]]:
        if self.pre_norm:
            attended, updated = self.attention.forward_cached(self.attention_norm(x), rope, cache)
            x = x + attended
            return x + self.feed_forward(self.feed_forward_norm(x)), updated
        attended, updated = self.attention.forward_cached(x, rope, cache)
        x = self.attention_norm(x + attended)
        return self.feed_forward_norm(x + self.feed_forward(x)), updated

The norm sits inside the branch, before attention and before the feed-forward layer. That's called pre-norm. The 2017 paper put it after the addition, which is post-norm, and Day 2 tested that choice in where the norm goes.

Normalize each vector across its own features

LayerNorm rescales one position's vector to mean 0 and variance 1, then applies a learned gain $\gamma$ and bias $\beta$ per feature:

The statistics come from the $C$ features of one vector. They never mix positions or sequences. The Day 1 reading plan made this point on purpose, because BatchNorm, the older method, averages across the batch, and that would make one sequence's statistics depend on another's.

So what happened ?

• The output always has mean 0 and standard deviation 1, whatever the input is.
• "Add 1 to every entry" leaves the output unchanged, LayerNorm removes any shift.
• "Double every entry" also leaves it unchanged, LayerNorm removes any scale.
• The order of the entries survives. Bigger inputs stay bigger.

Without a norm, activations can grow block by block, and softmax and the feed-forward layer behave differently at different scales. The norm gives every block an input at the same scale. Day 2 compares this with RMSNorm, which drops the mean.

A per-position two-layer network

Attention moves information between positions. The feed-forward layer then processes each position alone. It widens the vector to $4C$, applies a nonlinearity, and projects back:

For Day 1 that's 128 to 512 and back to 128. The two matrices hold 131,072 parameters per block, twice the attention projections. Most of a Transformer's parameters live here, not in attention, which surprised me the first time.

GELU is basically a smooth ReLU. It weights each input by the probability that a standard normal variable is below it, $\operatorname{GELU}(x) = x \cdot \Phi(x)$. Big positive inputs pass through, big negative ones go to 0, and near 0 the curve is smooth and dips a bit below zero.

class FeedForward(nn.Module):
    def __init__(self, config: DecoderConfig) -> None:
        super().__init__()
        hidden = config.hidden_size
        self.gated = config.feed_forward == "swiglu"
        self.up = nn.Linear(config.d_model, hidden, bias=False)
        self.gate = nn.Linear(config.d_model, hidden, bias=False) if self.gated else None
        self.down = nn.Linear(hidden, config.d_model, bias=False)
        self.dropout = nn.Dropout(config.dropout)

    def forward(self, x: Tensor) -> Tensor:
        if self.gate is None:
            return self.dropout(self.down(F.gelu(self.up(x))))
        return self.dropout(self.down(F.silu(self.gate(x)) * self.up(x)))

With feed_forward = "gelu" this is the Day 1 layer. The gate branch is SwiGLU, added on Day 2.

Tied weights, the embedding read backwards

The head turns each final vector into $V$ scores. octlm uses the embedding matrix for this, transposed:

The score for token $v$ is the dot product between the final hidden vector and that token's embedding row. So the model predicts token $v$ when its output vector points the same way as $v$'s input vector. One matrix does both ends, which saves $V \times C$ parameters, 131,072 at Day 1 size and 4.2 million at Day 4 size, and gives rare tokens' rows a gradient from both directions. In code it's one line, self.lm_head.weight = self.token_embedding.weight.

The bug that stopped the model from learning

Problem, the first test run of the baseline learned really poorly. The test that overfits one block failed.

The cause was the initialization. PyTorch's nn.Embedding draws its weights from a standard normal, $\sigma = 1$. As an input that's fine, because the first LayerNorm rescales it. But with a tied head, the same rows are also the output weights. The final hidden vector has entries around unit size after the last norm, and its dot product with a row of $C$ unit-variance numbers has a standard deviation of about $\sqrt{C} \approx 11$. Softmax over scores that far apart puts nearly all the probability on one random token.

So what happened ?

• At $\sigma = 1$ the top token takes most of the probability and the starting loss is way above $\ln V$. The model starts confident and wrong, and its first job is to undo that.
• At $\sigma = 0.02$ the scores spread by about 0.23, the distribution is nearly uniform, and the loss starts at $\ln V$.

The fix draws every linear and embedding weight from $\mathcal{N}(0, 0.02^2)$, the GPT-2 value, and sets biases to zero. The overfit test passed after that. Decision, keep.

@staticmethod
def _initialize(module: nn.Module) -> None:
    if isinstance(module, (nn.Linear, nn.Embedding)):
        nn.init.normal_(module.weight, mean=0.0, std=0.02)
    if isinstance(module, nn.Linear) and module.bias is not None:
        nn.init.zeros_(module.bias)

Where the 541,952 parameters live

Per block, attention holds $4C^2 = 65{,}536$ weights, the feed-forward layer $8C^2 = 131{,}072$, and the two norms $4C = 512$. The linear layers have no biases. The token table adds $VC$, the position table $T_{\max}C$, and the final norm $2C$.

What we didn't build, and why

• Biases on the linear layers. octlm follows current decoders and uses none. The norms keep their bias on Day 1 because PyTorch's LayerNorm has one.
• Dropout. The layers accept it and every config sets it to 0. These models are undertrained, not overfit.
• Scaled init for residual projections. GPT-2 also shrinks the output projections by $1/\sqrt{2L}$. At two layers the difference is small, and the change would need its own comparison.