Skip to main content

Documentation Index

Fetch the complete documentation index at: https://mintlify.com/itsubaki/gpt/llms.txt

Use this file to discover all available pages before exploring further.

Each transformer block in itsubaki/gpt follows a pre-normalization design: the residual stream is normalized with RMSNorm before entering each sub-layer, causal self-attention is applied and the result is added back to the stream, then a second RMSNorm feeds a SwiGLU feed-forward network whose output is added via another residual connection. This pre-norm arrangement stabilizes training without requiring careful learning-rate warmup.

​
Block constructor and struct

func Block(embedDim, numOfHeads int, rope function.RoPEFunc, useCache ...bool) *BlockT {
    headDim := int(embedDim / numOfHeads)
    return &BlockT{
        Layers: L.Layers{
            "norm1": RMSNorm(embedDim),
            "norm2": RMSNorm(embedDim),
            "attn":  MultiHeadAttention(embedDim, numOfHeads, headDim, rope, useCache...),
            "ffn":   SwiGLU(embedDim),
        },
    }
}

type BlockT struct {
    L.Layers
}
headDim is derived by integer division: embedDim / numOfHeads. The RoPEFunc closure is passed in from the parent GPT model and shared across all blocks so the cos/sin tables are computed once.

​
Forward pass — the residual stream

func (l *BlockT) Forward(x ...*variable.Variable) []*variable.Variable {
    x0 := l.Layers["norm1"].First(x...)  // pre-norm
    x1 := l.Layers["attn"].First(x0)    // causal self-attention
    x2 := F.Add(x[0], x1)               // residual add

    x3 := l.Layers["norm2"].First(x2)   // pre-norm
    x4 := l.Layers["ffn"].First(x3)     // SwiGLU FFN
    x5 := F.Add(x2, x4)                 // residual add

    return []*variable.Variable{x5}
}
The two F.Add calls implement skip connections: x[0] (the original input) is added to the attention output, and x2 (after the first residual) is added to the FFN output. Both norms operate on the residual stream before it enters a sub-layer, not after.

​
RMSNorm instead of LayerNorm

itsubaki/gpt uses Root Mean Square normalization in place of the standard LayerNorm.

​
Why RMSNorm

LayerNorm subtracts the per-token mean and divides by the standard deviation, then scales and shifts with learnable gamma and beta. RMSNorm omits both the mean subtraction and the beta shift: 
rms = sqrt(mean(x²) + ε)
y   = gamma * x / rms
This removes two operations per forward pass (mean subtraction, bias addition) and eliminates the beta parameter entirely, making RMSNorm slightly faster and fewer parameters to tune while empirically matching LayerNorm quality in autoregressive language models.

​
RMSNorm implementation

func RMSNorm(embedDim int) *RMSNormT {
    p := make(L.Parameters)
    p.Add("gamma", variable.Ones(embedDim))

    return &RMSNormT{
        eps:        1e-5,
        Parameters: p,
    }
}
The gamma vector is initialized to ones and learned during training. The eps value 1e-5 guards against division by zero. The forward pass computes: 
func (l *RMSNormT) Forward(x ...*variable.Variable) []*variable.Variable {
    shape := x[0].Shape()
    last := len(shape) - 1
    shape[last] = 1

    // rms = sqrt(mean(x^2) + eps)
    // y = gamma * x / rms
    x2  := F.Pow(2)(x[0])
    ms  := F.Reshape(shape...)(F.Mean(last)(x2)) // keepdims
    rms := F.Pow(0.5)(F.AddC(l.eps, ms))
    gamma := l.Parameters["gamma"]
    y := F.Mul(gamma, F.Div(x[0], rms))
    return []*variable.Variable{y}
}
The F.Reshape with keepdims-style shape ensures correct broadcasting when dividing the (B, C, EmbedDim) tensor by the scalar RMS computed over the last dimension.
Compare with LayerNorm(embedDim) in layer/layer_norm.go which additionally subtracts the mean and adds a learnable beta bias: y = gamma * (x - mean) / sqrt(variance + eps) + beta. Use RMSNorm for new models; LayerNorm is retained for reference.

​
SwiGLU feed-forward network

​
Constructor

func SwiGLU(xDim int) *SwiGLUT {
    hiddenDim := int(xDim * 8 / 3)
    return &SwiGLUT{
        Layers: L.Layers{
            "W": Linear(xDim, hiddenDim),
            "V": Linear(xDim, hiddenDim),
            "O": Linear(hiddenDim, xDim),
        },
    }
}
The hidden dimension is computed as hiddenDim = xDim * 8 / 3. For xDim = 192 that gives hiddenDim = 512; for xDim = 394 it gives 1050. This 8/3 ratio approximately matches the parameter count of a standard 4× FFN while fitting the three-matrix SwiGLU design.

​
Gated activation

func (l *SwiGLUT) Forward(x ...*variable.Variable) []*variable.Variable {
    a := l.Layers["W"].First(x...)
    b := l.Layers["V"].First(x...)

    gated := F.Mul(F.Mul(a, F.Sigmoid(a)), b)
    o := l.Layers["O"].First(gated)
    return []*variable.Variable{o}
}
The formula computed is: 
SwiGLU(x) = (W·x ⊙ sigmoid(W·x)) ⊙ (V·x)
W·x ⊙ sigmoid(W·x) is the SiLU (Swish) activation applied to the first linear branch, and the result is element-wise multiplied by the second linear branch V·x. The gated output is then projected back to xDim by O. This gating mechanism lets the network learn to suppress or amplify individual hidden units based on the input, often outperforming GELU-based FFNs.

​
Alternative FFN with GELU

The repository also provides a classic two-layer FFN for reference: 
func FFN(xDim, hiddenDim int) *FFNT {
    return &FFNT{
        Layers: L.Layers{
            "l1": Linear(xDim, hiddenDim),
            "l2": Linear(hiddenDim, xDim),
        },
    }
}

func (l *FFNT) Forward(x ...*variable.Variable) []*variable.Variable {
    x0 := l.Layers["l1"].First(x...)
    x1 := F.GELU(x0)
    x2 := l.Layers["l2"].First(x1)
    return []*variable.Variable{x2}
}
FFN uses a standard expand-then-contract shape with GELU activation and only two weight matrices compared to SwiGLU’s three. The default Block constructor uses SwiGLU; swap in FFN if you want a simpler baseline.

​
Cache propagation

BlockT.ClearCache() delegates directly to the attention layer inside the block:
func (l *BlockT) ClearCache() {
    l.Layers["attn"].(*MultiHeadAttentionT).ClearCache()
}
Calling GPT.ClearCache() walks every block[i] and calls this method, resetting the KV cache and token-offset counter in each attention layer.

Build docs developers (and LLMs) love

Get started for freeTalk to us