A model architecture is the blueprint defining a model’s layers, component types, and connections. For language models, choices depend on the task, data scale, and available compute. section 5.2.1
Architecture choices shape both capability and training practicality:
- Feed-forward layers provide a simple starting point; attention and Transformer components support modern language-model designs. section 5.2.1 section 2.5
- A larger or more sophisticated architecture does not guarantee better results; data quality and size can matter more than architecture or the training loop. section 5.5
- Architecture decisions interact with memory, speed, model size, tokenization, loss functions, and the training loop. section 5.5 section 9.1.3
A deliberately simple language model
A compact implementation can combine token embeddings with two linear layers and a ReLU activation, then predict vocabulary tokens using cross-entropy loss. This baseline makes later architectural changes visible rather than hiding them inside a large model. section 9.1.3
Listing 9.3 Simple model and training loop
class SimpleFeedForwardNN(nn.Module):
def __init__(self, config=MASTER_CONFIG):
super().__init__()
self.config = config
self.embedding = nn.Embedding(
config["vocab_size"], config["d_model"]
)
self.linear = nn.Sequential(
nn.Linear(config["d_model"], config["d_model"]),
nn.ReLU(),
nn.Linear(config["d_model"], config["vocab_size"]),
)
print(
f"model params: {sum([m.numel() for m in self.parameters()])}"
)
def forward(self, idx, targets=None):
x = self.embedding(idx)
logits = self.linear(x)
if targets is not None:
loss = F.cross_entropy(
logits.view(-1, self.config["vocab_size"]),
targets.view(-1),
ignore_index=tokenizer.pad_token_id,
# reduction="sum",
)
return logits, loss
else:
return logits
model = SimpleFeedForwardNN(MASTER_CONFIG).to(device)
opt_model = torch.compile(model) #1
Inspect the minimal PyTorch model to see embeddings, linear layers, activation, vocabulary logits, and padding-aware loss in one implementation. section 9.1.3
