Architecture

mvGPT is a decoder-only transformer — the same family as GPT-2, scaled down to ~14.6M parameters so the whole thing stays legible. Everything in this page lives in base_model.py, shared by every training and inference stage.

Watch the architecture walkthrough

Data flow

Tokens become embeddings, flow through six identical transformer blocks, and a final projection turns the last position's vector into a probability over the vocabulary. The single most important detail is causal masking: no token may attend to a token that comes after it, which is what makes the model autoregressive — able to generate left to right.

%%{init: {'theme':'dark','flowchart':{'htmlLabels':false},'themeVariables':{'fontFamily':'Helvetica, Arial, sans-serif','lineColor':'#8b949e','clusterBkg':'#161b22','clusterBorder':'#30363d','edgeLabelBackground':'#161b22'}}}%%
flowchart TD
    P["PROMPT<br/>“Once upon a time”"]
    TOK["TOKENIZE: string → ids<br/>[60, 61, 7, 47]<br/>out: T tokens (T ≤ 256)"]
    WTE["EMBED_TOKENS<br/>lookup table 10,000 × 384<br/>in: T → out: T × 384"]
    WPE["EMBED_POSITIONS<br/>lookup table 256 × 384<br/>in: T → out: T × 384"]
    ADD0{{"add token + position<br/>= T × 384"}}

    P --> TOK
    TOK -->|T| WTE
    TOK -->|T| WPE
    WTE -->|+| ADD0
    WPE -->|+| ADD0

    subgraph BLOCK["TRANSFORMER BLOCK ( ×6 )"]
        direction TB
        BIN["block input: T × 384"]
        AN["ATTN_NORM (LayerNorm)<br/>T × 384 → T × 384"]
        Q["Q_PROJ<br/>T × 384 → T × 384<br/>(6 heads × T × 64)"]
        K["K_PROJ<br/>T × 384 → T × 384<br/>(6 heads × T × 64)"]
        V["V_PROJ<br/>T × 384 → T × 384<br/>(6 heads × T × 64)"]
        ATT["ATTENTION (per head, ×6)<br/>softmax(Q · K^T / √64) → T × T<br/>scores · V → T × 64"]
        MH["MERGE_HEADS<br/>6 × (T × 64) → T × 384"]
        O["O_PROJ<br/>T × 384 → T × 384"]
        ADD1{{"+ residual<br/>T × 384"}}
        MN["MLP_NORM (LayerNorm)<br/>T × 384 → T × 384"]
        UP["UP_PROJ + GELU<br/>T × 384 → T × 1536"]
        DN["DOWN_PROJ<br/>T × 1536 → T × 384"]
        ADD2{{"+ residual<br/>T × 384"}}

        BIN --> AN
        AN --> Q & K & V
        Q --> ATT
        K --> ATT
        V --> ATT
        ATT --> MH --> O --> ADD1
        BIN -. residual .-> ADD1
        ADD1 --> MN --> UP --> DN --> ADD2
        ADD1 -. residual .-> ADD2
    end

    ADD0 --> BIN
    ADD2 --> FN["FINAL NORM (LayerNorm)<br/>T × 384 → T × 384"]
    FN --> LM["LM_HEAD (tied to WTE)<br/>T × 384 → T × 10,000<br/>logits over vocab"]
    LM --> SMP["SOFTMAX over vocab<br/>→ SAMPLE TOKEN"]

    classDef blue fill:#16324f,stroke:#5aa9f8,color:#e6edf3;
    classDef pink fill:#4a2230,stroke:#ff9eb1,color:#e6edf3;
    classDef green fill:#16402a,stroke:#5fd98a,color:#e6edf3;
    class P,TOK,WTE,WPE,ADD0 blue;
    class BIN,AN,Q,K,V,ATT,MH,O,ADD1,MN,UP,DN,ADD2 pink;
    class FN,LM,SMP green;

Reading the diagram top to bottom:

1. Tokenize the prompt into integer ids (T ≤ 256).
2. Embed each id (embed_tokens) and each position (embed_positions), then add them — the model has no other notion of order.
3. Run the sum through 6 transformer blocks, each a causal self-attention sublayer followed by an MLP, both wrapped in pre-norm residuals.
4. Apply a final LayerNorm, then the LM head (tied to the token embedding) to get logits over all 10,000 tokens.
5. Softmax + sample the next token, append it, and repeat.

Where the parameters live

Of the 14.6M parameters, ~73% are in the six transformer blocks and ~27% are the token-embedding matrix (which is reused as the output head — see weight tying below, so it is counted once).

%%{init: {'theme':'dark','flowchart':{'htmlLabels':false},'themeVariables':{'fontFamily':'Helvetica, Arial, sans-serif','lineColor':'#8b949e','clusterBkg':'#161b22','clusterBorder':'#30363d','edgeLabelBackground':'#161b22'}}}%%
flowchart TD
    subgraph EMB["EMBEDDINGS: 3,938,304"]
        direction TB
        WTE["EMBED_TOKENS<br/>10,000 × 384 = 3,840,000"]
        WPE["EMBED_POSITIONS<br/>256 × 384 = 98,304"]
    end

    subgraph BLK["TRANSFORMER BLOCK ( ×6 )"]
        direction TB
        AN["ATTN_NORM (weight + bias)<br/>2 × 384 = 768"]
        QKV["Q_PROJ + K_PROJ + V_PROJ<br/>3 × (384×384 + 384) = 443,520"]
        O["O_PROJ<br/>384×384 + 384 = 147,840"]
        MN["MLP_NORM (weight + bias)<br/>2 × 384 = 768"]
        UP["UP_PROJ (384 → 1536)<br/>384×1536 + 1536 = 591,360"]
        DN["DOWN_PROJ (1536 → 384)<br/>1536×384 + 384 = 590,208"]
        SUB(["per-block subtotal = 1,774,464<br/>× 6 blocks = 10,646,784"])
        AN --> QKV --> O --> MN --> UP --> DN --> SUB
    end

    subgraph HEAD["OUTPUT HEAD: 768"]
        direction TB
        FN["FINAL NORM (weight + bias)<br/>2 × 384 = 768"]
        LM["LM_HEAD<br/>tied to EMBED_TOKENS → +0"]
    end

    TIE(["WEIGHT TYING<br/>LM_HEAD reuses EMBED_TOKENS<br/>(10,000 × 384), counted once above"])
    TOTAL["TOTAL PARAMETERS<br/>3,938,304 + 10,646,784 + 768<br/>= 14,585,856"]

    EMB --> BLK --> HEAD --> TOTAL
    WTE -. shared .- TIE
    LM -. shared .- TIE

    classDef blue fill:#16324f,stroke:#5aa9f8,color:#e6edf3;
    classDef pink fill:#4a2230,stroke:#ff9eb1,color:#e6edf3;
    classDef green fill:#16402a,stroke:#5fd98a,color:#e6edf3;
    classDef yellow fill:#4a3a12,stroke:#e9c46a,color:#e6edf3;
    classDef total fill:#0d1117,stroke:#e6edf3,stroke-width:3px,color:#e6edf3;
    class WTE,WPE blue;
    class AN,QKV,O,MN,UP,DN,SUB pink;
    class FN,LM green;
    class TIE yellow;
    class TOTAL total;

Components

Tokenizer

A deliberately simple word-level tokenizer (WordTokenizer). Text is pre-tokenized with a regex into runs of word characters or punctuation (PRETOK = \w+|[^\w\s]+), whitespace is dropped, and each piece is looked up in a 10,000-entry vocabulary. Three reserved specials take the first ids:

Token	Id	Purpose
[UNK]	0	out-of-vocabulary fallback
[PAD]	1	batch padding
<\|endoftext\|>	2	document / generation boundary

decode(encode(x)) recovers the words but not the exact spacing, since tokens are rejoined with single spaces.

Embeddings

self.embed_tokens = nn.Embedding(vocab_size, n_embd)     # 10000 × 384
self.embed_positions = nn.Embedding(block_size, n_embd)  # 256 × 384
x = self.drop(self.embed_tokens(idx) + self.embed_positions(pos))

Positions are learned (an embedding table), not sinusoidal, and always run 0..T-1 with no offset during training.

Causal self-attention

The only layer where tokens exchange information — and only backward. Each position projects to a query, key, and value; attention is computed per head with PyTorch's fused scaled_dot_product_attention:

\[ \text{Attention}(Q,K,V) = \text{softmax}\!\left(\frac{QK^{\top}}{\sqrt{d_k}}\right)V \]

y = F.scaled_dot_product_attention(q, k, v, dropout_p=..., is_causal=True)

is_causal=True builds the lower-triangular mask internally (6 heads × 64 dims). After attention the heads are merged and mixed back together by o_proj.

MLP

A position-wise feed-forward network — attention moves information between tokens, the MLP processes each token independently. Expand 4×, apply GELU, project back:

self.up_proj   = nn.Linear(n_embd, 4 * n_embd)   # 384 → 1536
self.down_proj = nn.Linear(4 * n_embd, n_embd)   # 1536 → 384
# forward: down_proj(gelu(up_proj(x)))

Transformer block (pre-norm residual)

LayerNorm is applied before each sublayer, not after. This keeps a clean identity path through the residual stream, so each block adds a small correction instead of replacing the representation — which is what makes deep stacks trainable.

x = x + self.attn(self.attn_norm(x))   # attention sublayer
x = x + self.mlp(self.mlp_norm(x))     # MLP sublayer

Output head & weight tying

self.lm_head = nn.Linear(n_embd, vocab_size, bias=False)  # 384 → 10000
self.embed_tokens.weight = self.lm_head.weight            # tied

The output projection shares its weight matrix with the token embedding. That saves 3.84M parameters and ties "what a token means on the way in" to "what predicting that token looks like on the way out."

Initialization

GPT-2-style: weights ~ N(0, 0.02), biases zero. The two residual output projections (o_proj, down_proj) are scaled down by 1/√(2·n_layer) so the residual stream's variance stays flat as depth grows.

Configuration

Every shape comes from a single dataclass, so a checkpoint is only loadable into a model built with the same config:

@dataclass
class GPTConfig:
    vocab_size: int = 10000
    block_size: int = 256   # max context length
    n_layer:    int = 6
    n_head:     int = 6
    n_embd:     int = 384   # must be divisible by n_head
    dropout:  float = 0.1
    bias:      bool = True

Generation

GPT.generate() samples autoregressively with temperature and optional top-k filtering, cropping the context to the last block_size tokens. It re-runs the whole context through all six blocks on every step — simple, but O(T²) over a generation. The Inference & KV Cache page shows how to fix that without changing the weights.

---

📄 Source: base_model.py