Transformer Block

One layer of a transformer: attention and a position-wise FFN, each wrapped in a residual connection with normalization.

Equation

Post-norm (original): $x\leftarrow \mathrm{LN}(x+\mathrm{Attn}(x)),x\leftarrow \mathrm{LN}(x+\mathrm{FFN}(x))$

Pre-norm (modern): $x\leftarrow x+\mathrm{Attn}(\mathrm{LN}(x)),x\leftarrow x+\mathrm{FFN}(\mathrm{LN}(x))$

Pre-norm leaves an unnormalized identity path (the residual stream) from input to logits, so gradients reach early layers without warmup. FFN with SwiGLU:

$\mathrm{FFN}(x)=\left(\mathrm{SiLU}(x{W}_g)\odot x{W}_u\right)W_d$

From scratch

import torch, torch.nn as nn, torch.nn.functional as F
 
class SwiGLU(nn.Module):
    def __init__(self, d, mult=8/3):            # 8/3 keeps params ~= 4d^2 of 2-matrix FFN
        super().__init__()
        h = int(mult * d) // 64 * 64            # round hidden dim for kernel efficiency
        self.wg, self.wu = nn.Linear(d, h, bias=False), nn.Linear(d, h, bias=False)
        self.wd = nn.Linear(h, d, bias=False)
    def forward(self, x):                       # (b, n, d)
        return self.wd(F.silu(self.wg(x)) * self.wu(x))
 
class Block(nn.Module):
    def __init__(self, d, n_heads):
        super().__init__()
        self.n1, self.n2 = nn.RMSNorm(d), nn.RMSNorm(d)
        self.attn, self.ffn = MHA(d, n_heads), SwiGLU(d)
    def forward(self, x):                       # (b, n, d) -> (b, n, d)
        x = x + self.attn(self.n1(x))           # pre-norm residual
        return x + self.ffn(self.n2(x))

Notes

• Params per block $\approx 12d^2$: $4d^2$ attention, $8d^2$ FFN (expansion 4, or 8/3 with three SwiGLU matrices).
• Post-norm needs LR warmup and careful init; pre-norm trains stably but the residual stream variance grows with depth, so scale output projections by $1/\sqrt{2L}$ at init.
• Sandwich/QK norm (extra norm on attention output or on q/k) is used to stop logit blowup at large scale.
• Attention moves information between positions, FFN processes each position independently; that division is the whole architecture.

Related: Transformers, Multi-Head Attention, LayerNorm, MoE Routing