Test the Model

Running entirely in your browser using ONNX Runtime WebGL/WASM

1.9M Params 4-Layer LLaMA-style WASM + INT8 (2.3MB)
Loading tokenizer...
Advanced Generation Settings

Tokenizer Playground

See how the Byte-Pair Encoding (BPE) algorithm splits text into subwords.

Live BPE Encoding

Merge Steps Visualizer:

Start typing to see character merges...

Final Token IDs (0 tokens):

How BPE Works

The tokenizer first uses a regex pattern to split text into words (e.g. 'Hello, world!' → ['Hello', ',', ' world', '!']). It then iteratively merges adjacent characters based on a trained frequency list until no more merges can be made. Decoding is simply looking up the IDs and joining the strings.

Architecture & Math

Explore the mathematical foundations of the Transformer architecture through interactive data flow.

Model Overview

The data flows through several stages to predict the next token. Click any stage to jump to its interactive explanation below!

↓
↓
↓
↓

1. Token Embedding

Each token ID acts as an index to look up a dense continuous vector from the embedding matrix. This vector represents the semantic meaning of the token.

Token ID: 42 → [0.12, -0.45, 0.88, ..., 0.03]

2. Rotary Positional Embedding (RoPE)

RoPE injects positional information by rotating the Query and Key vectors in a 2D plane based on their position in the sequence.

3. Root Mean Square Normalization (RMSNorm)

Unlike LayerNorm, RMSNorm doesn't center the mean to zero, it simply scales the vector by its Root Mean Square. Adjust the input variance below to see the bars normalize!

Raw Input Vector

➔

Normalized Output (Constant Scale)

4. Causal Self-Attention

Words attend to previous words to build context. Future tokens are masked out (grayed out). Hover or click a token to see its attention fan!

The
cute
cat
purred
The
cute
cat
purred

5. SwiGLU FFN

The Feed-Forward Network splits the input into two branches. One branch applies the non-linear SiLU activation to act as a "gate", controlling how much of the other branch's linear projection passes through via a Hadamard (element-wise) product.



SiLU(Gate)
0.00
⊗
Linear
10.0
↓
Output
0.00

6. Output Head

The final vector is projected back to the vocabulary size (1978 tokens). We apply Temperature scaling and Softmax to get a probability distribution, then sample the next token!

Reference: Consolidated Mathematical Equations
$$ \begin{aligned} H_0 &= \text{Embed}(X) \\ H_l &= F_l(H_{l-1}) \quad \text{for } l \in \{1, 2, \dots, N\} \\ \text{Logits}(X) &= \text{RMSNorm}_{\text{final}}(H_N) W_{\text{head}}^T \\ F_l(H) &= \tilde{H} + \left[ \text{SiLU}\left( \text{RMSNorm}_{\text{ffn}}(\tilde{H}) W_1^T \right) \odot \left( \text{RMSNorm}_{\text{ffn}}(\tilde{H}) W_3^T \right) \right] W_2^T \\ \text{where} \quad \tilde{H} &= H + A_l\big(\text{RMSNorm}_{\text{attn}}(H)\big) \\ A_l(x) &= \text{softmax}\left(\frac{\text{RoPE}(x W_q^T)\text{RoPE}(x W_k^T)^T}{\sqrt{d_{\text{head}}}} + M\right)(x W_v^T) W_o^T \end{aligned} $$

Training Pipeline

How the model learns from raw data.

Data Flow

The corpus is tokenized, chunked into sequences, and batched onto the GPU.

corpus.txt
Tokenizer
Batch (x,y)
GPU VRAM
Model

The Autoregressive Training Loop

Step through a single backpropagation iteration to see the exact PyTorch code executed.

1. Zero Gradients
2. Forward Pass (Logits)
3. Reshape & Loss
4. Backward Pass (Grads)
5. Optimizer Step
optimizer.zero_grad()
Clears old gradients from the previous step so they don't accumulate.

Hyperparameters

Optimizer
AdamW
Learning Rate
3e-4
Batch Size
64
Epochs
10

Loss Curve (Convergence)

Source Reference

Explore the actual codebase implementation.

File Tree

  • model.py

    Defines the PyTorch architecture (RoPE, RMSNorm, SwiGLU, Attention). View on GitHub

  • data.py

    Handles loading corpus.txt and batching. View on GitHub

  • tokenizer.py

    Builds the BPE tokenizer and exports tokenizer.json. View on GitHub

  • train.py

    The main training loop. View on GitHub