Pure rust implementation of a mini GPT

Goal

Build a minimal GPT implementation from scratch in Rust.

---

1. Tensors

2. Automatic differentiation

Attention Helpers

tensor.masked_fill(mask, value) -> Tensor

---

Layers

Embedding

Embedding::new(vocab_size, embed_dim)
Embedding::forward(tokens)

Maps token ids → vectors.

Linear

Linear::new(in_dim, out_dim)
Linear::forward(x)

Computes:

y = xW + b

---

Attention

Single Head

Attention::new(d_model)
Attention::forward(x)

Computes:

Q = XWq
K = XWk
V = XWv

scores = QKᵀ / sqrt(d)
weights = softmax(scores)

output = weightsV

Causal Mask

Prevent tokens from attending to future tokens.

---

Multi-Head Attention

MultiHeadAttention::new(
    d_model,
    num_heads
)

Split → attend → concat.

---

LayerNorm

LayerNorm::new(dim)
LayerNorm::forward(x)

Computes:

(x - mean) / sqrt(var + eps)

---

MLP

MLP::new(
    d_model,
    hidden_dim
)

Computes:

Linear
↓
GELU
↓
Linear

---

Transformer Block

TransformerBlock::new(...)
TransformerBlock::forward(x)

Architecture:

LayerNorm
↓
Attention
↓
Residual

LayerNorm
↓
MLP
↓
Residual

---

GPT Model

GPT::new(...)
GPT::forward(tokens)

Components:

Token Embedding
↓
Positional Embedding
↓
N Transformer Blocks
↓
Output Projection

---

Generation

GPT::generate(prompt)

Loop:

forward
↓
next token
↓
append
↓
repeat

Sampling strategies:

argmax()
temperature()
top_k()

---

Training (Later)

Loss

cross_entropy(logits, targets)

Autograd

loss.backward()

Optimizer

AdamW

Training loop:

for batch in dataset {
    let logits = model.forward(&x);

    let loss = cross_entropy(
        &logits,
        &targets,
    );

    loss.backward();

    optimizer.step();
    optimizer.zero_grad();
}

Run tests

cargo test