To kick my very own LLM model that can be train and learned from simple dataset - i started off with a bare minimum model that allows me to learn really simple stuff like 1 + 1 = 2, 2 + 2 = 4.
So first we outline our vocabulary and size - how we are representing this information to the LLM model
Vocabulary and token
Then we tokenize those input for training and inference. This model only understands these vocab
Special: <pad>, <start>, <eos>Operators: +, -, =Numbers: 1, 2, 3, 4Simple transformer with:
- Embedding Layer: Converts token IDs to 32-dim vectors and this is the layer where we handle our vocab
# Token embedding
self.embedding = nn.Embedding(vocab_size, d_model, padding_idx=0)
- Positional Encoding: Learnable position embeddings
# Positional encoding (learnable)
self.pos_embedding = nn.Embedding(seq_len, d_model)
- Transformer Encoder: 1 layer with 2 attention heads
- Output Linear: Predicts next token (10 possible tokens)
Why do we have embedding and position layer and both uses nn.Embedding?
Embedding layer allow our model to be able to tell which token id - for example token id 6 which is 1.
Positional embedding provides details about whereabout or position that the token sits in a input. For example, "1 + 2 = 3" - when we have token 6 => value 1, we also get position 1.
Token Embedding vs Positional Embedding
| Embedding Type | Input | Purpose | Encodes |
|---|
| Token Embedding | Which token it is (ID) | Semantic meaning | WHAT is the token |
| Positional Embedding | Where it is (position) | Order information | WHERE is the token |
Without positional embedding, the transformer can't tell the difference between "1 + 1" and "1 1 +"
Positional embedding adds order information so the model knows:
- Token "1" is at position 0
- Token "+" is at position 1
- Token "1" is at position 2
Input: "1 + 1" (positions 0, 1, 2)
Token Embedding:
Position 0: Token "1" → [0.2, -0.5, 0.8, ...] (embedding for "1")
Position 1: Token "+" → [-0.1, 0.4, -0.2, ...] (embedding for "+")
Position 2: Token "1" → [0.2, -0.5, 0.8, ...] (same "1" again)
Positional Embedding:
Position 0: [0.9, 0.1, -0.3, ...] (position 0)
Position 1: [0.5, -0.2, 0.7, ...] (position 1)
Position 2: [0.3, 0.6, -0.1, ...] (position 2)
Final (Added Together):
Position 0: [0.2+0.9, -0.5+0.1, 0.8-0.3, ...] = [1.1, -0.4, 0.5, ...]
Position 1: [-0.1+0.5, 0.4-0.2, -0.2+0.7, ...] = [0.4, 0.2, 0.5, ...]
Position 2: [0.2+0.3, -0.5+0.6, 0.8-0.1, ...] = [0.5, 0.1, 0.7, ...]
And now, the transformer would be able to correctly idenified what information we are feeding it. 👍
TransformerEncoderLayer
This is out setup where we use 2 atttention head. We have our model with 32 dimension. nHead is 2 where 32 / 2 = 16 = so we are getting 16 for our attention. FeedForward is 64 which gives the model space to learn.
# Single transformer block
self.transformer_layer = nn.TransformerEncoderLayer(
d_model=d_model,
nhead=2, # 2 attention heads
dim_feedforward=64,
dropout=0.1,
batch_first=True,
activation='relu'
)
Here is the code that allows me to do that.
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import random
# ============================================================================
# DATASET
# ============================================================================
class ArithmeticDataset(Dataset):
"""
Simple dataset for arithmetic: "1 + 1 = 2", "2 + 2 = 4"
Tokenizes into: [num1, +, num2, =, result]
"""
def __init__(self, examples):
self.examples = examples
# Simple vocab: 0=PAD, 1=+, 2=-, 3==, 4=EOS, 5+=<start>, 6-9=numbers 1-4
self.vocab = {
'<pad>': 0,
'+': 1,
'-': 2,
'=': 3,
'<eos>': 4,
'<start>': 5,
'1': 6,
'2': 7,
'3': 8,
'4': 9,
}
self.vocab_size = len(self.vocab)
self.seq_len = 7 # <start> + num + op + num + = + result + <eos>
def tokenize(self, example):
"""Convert "1 + 1 = 2" to token ids"""
tokens = [self.vocab['<start>']]
tokens.extend([self.vocab[t] for t in example.split()])
tokens.append(self.vocab['<eos>'])
# Pad to fixed length
while len(tokens) < self.seq_len:
tokens.append(self.vocab['<pad>'])
return tokens[:self.seq_len]
def __len__(self):
return len(self.examples)
def __getitem__(self, idx):
example = self.examples[idx]
tokens = self.tokenize(example)
# Input: all but last token, Target: all but first token
input_seq = torch.tensor(tokens[:-1], dtype=torch.long)
target_seq = torch.tensor(tokens[1:], dtype=torch.long)
return input_seq, target_seq
# ============================================================================
# MODEL: Barebone Transformer
# ============================================================================
class BareboneLLM(nn.Module):
"""
Minimal LLM with:
- Token embedding
- Single transformer block (1 layer)
- Simple attention
- Output linear layer
"""
def __init__(self, vocab_size, d_model=32, seq_len=7):
super().__init__()
self.vocab_size = vocab_size
self.d_model = d_model
self.seq_len = seq_len
# Token embedding
self.embedding = nn.Embedding(vocab_size, d_model, padding_idx=0)
# Positional encoding (learnable)
self.pos_embedding = nn.Embedding(seq_len, d_model)
# Single transformer block
self.transformer_layer = nn.TransformerEncoderLayer(
d_model=d_model,
nhead=2, # 2 attention heads
dim_feedforward=64,
dropout=0.1,
batch_first=True,
activation='relu'
)
# Output layer: predict next token
self.output_linear = nn.Linear(d_model, vocab_size)
def forward(self, x):
# x shape: (batch_size, seq_len)
batch_size, seq_len = x.shape
# Token embedding
x_emb = self.embedding(x) # (batch_size, seq_len, d_model)
# Add positional encoding
positions = torch.arange(seq_len, device=x.device).unsqueeze(0).expand(batch_size, -1)
pos_emb = self.pos_embedding(positions)
x_emb = x_emb + pos_emb
# Transformer encoding
x_transformed = self.transformer_layer(x_emb) # (batch_size, seq_len, d_model)
# Output logits
logits = self.output_linear(x_transformed) # (batch_size, seq_len, vocab_size)
return logits
# ============================================================================
# TRAINING LOOP
# ============================================================================
def train_model(model, train_loader, num_epochs=500, learning_rate=0.01, device='cpu'):
"""Train the model"""
model.to(device)
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
criterion = nn.CrossEntropyLoss(ignore_index=0) # Ignore padding token
print("Starting training...")
for epoch in range(num_epochs):
total_loss = 0
for input_seq, target_seq in train_loader:
input_seq = input_seq.to(device)
target_seq = target_seq.to(device)
# Forward pass
logits = model(input_seq) # (batch_size, seq_len, vocab_size)
# Compute loss
# Reshape for loss computation
loss = criterion(
logits.reshape(-1, model.vocab_size),
target_seq.reshape(-1)
)
# Backward pass
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_loss += loss.item()
if (epoch + 1) % 100 == 0:
avg_loss = total_loss / len(train_loader)
print(f"Epoch {epoch + 1}/{num_epochs} - Loss: {avg_loss:.4f}")
print("Training complete!")
return model
# ============================================================================
# INFERENCE
# ============================================================================
def generate(model, prompt_tokens, vocab_inv, max_len=10, device='cpu'):
"""Generate text token by token"""
model.eval()
model.to(device)
generated = prompt_tokens.copy()
with torch.no_grad():
for _ in range(max_len):
# Pad to sequence length
seq = generated + [0] * (7 - len(generated))
seq = seq[:7]
input_tensor = torch.tensor([seq], dtype=torch.long, device=device)
logits = model(input_tensor)
# Get last token logits
last_logits = logits[0, len(generated) - 1, :]
next_token = last_logits.argmax(dim=-1).item()
generated.append(next_token)
# Stop at EOS token
if next_token == 4:
break
return generated
# ============================================================================
# MAIN
# ============================================================================
if __name__ == "__main__":
# Set seed for reproducibility
torch.manual_seed(42)
random.seed(42)
# Device
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {device}\n")
# Create dataset with simple examples
examples = [
"1 + 1 = 2",
"2 + 2 = 4",
]
dataset = ArithmeticDataset(examples)
print("Dataset examples:")
for ex in examples:
tokens = dataset.tokenize(ex)
print(f" '{ex}' -> {tokens}")
print()
# Create data loader
train_loader = DataLoader(dataset, batch_size=1, shuffle=True)
# Initialize model
model = BareboneLLM(
vocab_size=dataset.vocab_size,
d_model=32,
seq_len=7
)
print(f"Model vocab size: {dataset.vocab_size}")
print(f"Model parameters: {sum(p.numel() for p in model.parameters())}\n")
# Train
model = train_model(model, train_loader, num_epochs=500, learning_rate=0.01, device=device)
# Inference
print("\n" + "="*50)
print("INFERENCE")
print("="*50)
# Create inverse vocabulary
vocab_inv = {v: k for k, v in dataset.vocab.items()}
# Test: generate "1 + 1 ="
print("\nPrompt: '1 + 1 ='")
prompt = [dataset.vocab['<start>'], dataset.vocab['1'], dataset.vocab['+'], dataset.vocab['1'], dataset.vocab['=']]
generated = generate(model, prompt, vocab_inv, device=device)
output = ' '.join([vocab_inv.get(t, '?') for t in generated if t != dataset.vocab['<pad>']])
print(f"Generated: {output}")
# Test: generate "2 + 2 ="
print("\nPrompt: '2 + 2 ='")
prompt = [dataset.vocab['<start>'], dataset.vocab['2'], dataset.vocab['+'], dataset.vocab['2'], dataset.vocab['=']]
generated = generate(model, prompt, vocab_inv, device=device)
output = ' '.join([vocab_inv.get(t, '?') for t in generated if t != dataset.vocab['<pad>']])
print(f"Generated: {output}")
Some trials and errors
Setting nHead to 4 and my model fail to learn and generated incomplete results when i prompt it "world is ". Using 4 head attention that is too small, where 32 / 4 = 8 can bite us back where the model fails to pay enough attention.
Setting feedforward to 128 - generate the same output but notice that the training loss is higher.
Feedforward 64
Starting training...
Epoch 100/500 - Loss: 0.2964
Epoch 200/500 - Loss: 0.2640
Epoch 300/500 - Loss: 0.4904
Epoch 400/500 - Loss: 0.2587
Epoch 500/500 - Loss: 0.4174
Training complete!
Feedforard 32
Starting training...
Epoch 100/500 - Loss: 0.2489
Epoch 200/500 - Loss: 0.3762
Epoch 300/500 - Loss: 0.2842
Epoch 400/500 - Loss: 0.2166
Epoch 500/500 - Loss: 0.1879
Training complete!
Comments