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.

The layer package provides Linear for learned linear projections (used throughout attention, the FFN, and the unembedding head) and Embeddings for token-to-vector lookup (the input embedding table that maps integer token IDs to dense vectors). Both layers hold a single weight matrix w initialized with initw().

Linear

func Linear(xDim, hiddenDim int) *LinearT
Constructs a bias-free linear projection layer. This is the fundamental building block of the model: every weight matrix in the attention projections (Wq, Wk, Wv, Wo), the SwiGLU FFN (W, V, O), and the final unembedding head is a Linear layer.
xDim
int
required
Input feature dimension — the size of the last axis of the input tensor. Determines the number of rows in the weight matrix.
hiddenDim
int
required
Output feature dimension. Determines the number of columns in the weight matrix and the size of the last axis of the output tensor.

Parameters

ParameterShapeInitial value
w(xDim, hiddenDim)initw(xDim, hiddenDim)
No bias term is stored. This matches the convention used in modern LLM implementations (LLaMA, etc.) where bias-free projections reduce parameter count with minimal quality impact.

Forward pass

y = x @ w   // matrix multiplication
The operation is a batched matrix multiply: for a 3-D input of shape (B, C, xDim), the output is (B, C, hiddenDim).

Shape examples

// Attention projection (square): (B, C, 512) → (B, C, 512)
l := layer.Linear(512, 512)

// Attention Q/K/V with multiple heads: (B, C, 512) → (B, C, 512)
l := layer.Linear(embedDim, numOfHeads*headDim)

// Unembedding head: (B, C, 192) → (B, C, vocabSize)
l := layer.Linear(192, 1000)

Where Linear is used

LocationxDimhiddenDim
Attention Wq / Wk / WvembedDimnumOfHeads * headDim
Attention WonumOfHeads * headDimembedDim
SwiGLU W / V (gate + value)embedDimembedDim * 8 / 3
SwiGLU O (output projection)embedDim * 8 / 3embedDim
Unembedding headembedDimvocabSize

Embeddings

func Embeddings(xDim, embedDim int) *EmbeddingsT
Constructs a token embedding lookup table. During the forward pass, each integer token ID is used as a row index into the weight matrix w, producing a dense vector of dimension embedDim.
xDim
int
required
Vocabulary size — the total number of distinct token IDs. Determines the number of rows in the embedding table. Input token IDs must be in the range [0, xDim).
embedDim
int
required
Embedding dimension — the length of each token’s dense vector representation. Determines the number of columns in the embedding table and the size of the last axis of the output tensor.

Parameters

ParameterShapeInitial value
w(xDim, embedDim)initw(xDim, embedDim)

Forward pass

The forward pass flattens the input token IDs, performs a batched row-gather (F.GetItem) on w, and reshapes the result to append embedDim to the original input shape.
ids := /* integer token IDs from x[0].Data */
w   := l.Parameters["w"]
y   := F.Reshape(append(x[0].Shape(), l.embedDim)...)(F.GetItem(0, ids)(w))

Shape example

embed := layer.Embeddings(1000, 192)
// Input ids: shape (B, C), values in [0, 999]
// Output:    shape (B, C, 192)
For an input batch of shape (B, C)B sequences each of length C — each integer is independently replaced by its row in the embedding table, producing an output of shape (B, C, embedDim).

Embedding table vs. unembedding head

The Embeddings layer and the final unembedding Linear layer are separate and do not share weights (no weight tying). The unembedding Linear(embedDim, vocabSize) has its own independently initialized w parameter.
Weight tying (sharing the embedding table with the unembedding projection) is a common technique that reduces parameter count, but this implementation keeps them separate for simplicity. If you need weight tying, you can manually assign the embedding weight to the unembedding layer after construction.

Weight initialization

Both Linear and Embeddings initialize their weight matrices through the same internal helper:
func initw(x, y int) *variable.Variable {
    return variable.From(tensor.Normal([]int{x, y}, 0, 0.02))
}
Weights are drawn from a normal distribution with mean 0 and standard deviation 0.02. This is the same initialization strategy used in the original GPT-2 paper and keeps the initial activations small enough to avoid gradient saturation at the start of training.
The standard deviation 0.02 is fixed and not scaled by layer depth. Some implementations apply 1/sqrt(2 * numLayers) scaling to residual projections (Wo and the FFN output) to account for the accumulation of variance across blocks. This model does not apply that scaling.

Build docs developers (and LLMs) love