# Why doesn't my toy transformer model "grok"?

I'm working on reproducing the results by Neel Nanda on teaching a small transformer to perform modular addition: (operand_1+operand_2)%mod_value.

The expectation for this demo is for the train loss to quickly decrease while the test loss remains high due to memorization. Then the model "groks", meaning learns the generalized solution, and slowly transitions from memorization to this solution.

As a result, the test loss plummets much farther into training I'm implementing the transformer from scratch with TinyGrad. I prefer TinyGrad since I understand it a lot better than PyTorch(debugging is easier) and it offers better METAL accelerator support. This model is simpler than the typical transformer -- it doesn't have biases or normalization, and we train over the entire training set in a single batch. The optimizer is AdamW with a very high weight decay to encourage faster generalization.

As far as I can tell, I've properly matched the parameters stated in the original material. I assume it's a bug with my Transformer code -- this implementation was based off this, but I removed the unnecessary components. It handles positional embeddings differently, so this may be the cause.

While the train loss should immediately decrease to close to 0 after just a few hundred epochs, the model is unable to get a better score than about 1.5 for both the training and test loss. I'll paste plots for the train and test loss over 5000 epochs below.

My code:

from tinygrad import Tensor
import random
import numpy as np
from tqdm import trange

mod = 113

operand_1 = Tensor.arange(113).unsqueeze(0).repeat([113,1]).flatten(0)
operand_2 = Tensor.arange(113).unsqueeze(0).repeat([113,1]).T.flatten(0)
equals = Tensor.full_like(operand_1, 113)

dataset = Tensor.stack([operand_1, operand_2, equals], dim=1)
targets = Tensor((dataset[:,0].numpy() + dataset[:,1].numpy()) % mod).unsqueeze(1)

train_split = .3
cutoff = int((mod**2)*train_split)
indices = np.random.permutation(mod**2)
train_indices = Tensor(indices[:cutoff])
test_indices = Tensor(indices[cutoff:])

train_ds = dataset[train_indices]
train_targets = targets[train_indices]
test_ds = dataset[test_indices]
test_targets = targets[test_indices]

self.query = (
Tensor.scaled_uniform((self.embed_dim, self.embed_dim)),
Tensor.zeros(self.embed_dim),
)
self.key = (
Tensor.scaled_uniform((self.embed_dim, self.embed_dim)),
Tensor.zeros(self.embed_dim),
)
self.value = (
Tensor.scaled_uniform((self.embed_dim, self.embed_dim)),
Tensor.zeros(self.embed_dim),
)

self.out = (
Tensor.scaled_uniform(self.embed_dim, self.embed_dim),
Tensor.zeros(self.embed_dim),
)

bz = x.shape[0]
QKV = [
x.linear(weight, bias).reshape(
)
for weight, bias in [self.query, self.key, self.value]
]

A = (QKV[0] @ QKV[1].transpose(-2, -1) / self.head_dim**0.5).softmax(
-1
)

out = (
(A @ QKV[2])
.transpose(1, 2)
.reshape(bz, -1, self.embed_dim)
.linear(*self.out)
)

return out

class TransformerBlock:
# no bias or layernorm!
self.ffn1 = (
Tensor.scaled_uniform((embed_dim, ffn_dim)),
)
self.ffn2 = (
Tensor.scaled_uniform((ffn_dim, embed_dim)),
)

x = x + self.attention(x, mask)
x = x + x.linear(*self.ffn1).relu().linear(*self.ffn2)
return x

class Transformer:
def __init__(
self,
vocab_size,
context_length,
num_layers,
embed_dim,
ffn_dim,
):
self.vocab_size = vocab_size
self.context_length = context_length
self.embed_dim = embed_dim
self.ffn_dim = ffn_dim

self.embed = Tensor.scaled_uniform(
(context_length + vocab_size, embed_dim)
)
self.blocks = [
for _ in range(num_layers)
]
self.unembed = Tensor.scaled_uniform((embed_dim, vocab_size))

def __call__(self, x):
B, T = x.shape
assert (
T <= self.context_length
), "Input shape must be (batch, context_length)"

positional_embeddings = Tensor.eye(T).unsqueeze(0).expand([B, T, T])

x = x.one_hot(self.vocab_size)
x = positional_embeddings.cat(x, dim=2).flatten(end_dim=1)
x = (x @ self.embed).reshape(B, T, self.embed_dim)
x = x.sequential(self.blocks)
x = (x.reshape((-1, self.embed_dim)) @ self.unembed).log_softmax()
return x.reshape((B, T, self.vocab_size))

def loss_fn(logits : Tensor, labels):
log_probs = logits.log_softmax(axis=-1)
correct = log_probs.gather(idx=labels, dim=-1)[:,0]
return -correct.mean()

def train(
model,
X_train,
Y_train,
X_test,
Y_test,
optim,
steps=1, # Only one step is needed for full batch training
lossfn=lambda out, y: out.sparse_categorical_crossentropy(y),
allow_jit=True,
):
def train_step(x, y):
# network
out = model(x)[:,-1]
loss = lossfn(out, y)
loss.backward()
optim.step()
return loss.realize()

def test_step(x, y):
out = model(x)[:,-1]
loss = lossfn(out, y)
return loss.realize()

if allow_jit:
train_step = TinyJit(train_step)

with Tensor.train():
train_losses = []
test_losses = []
for i in (t := trange(steps, disable=CI)):
train_loss = train_step(X_train, Y_train).numpy()
test_loss = test_step(X_test, Y_test).numpy()
train_losses.append(train_loss)
test_losses.append(test_loss)
t.set_description("train loss: %.2f test loss: %.2f" % (train_loss, test_loss))
return [train_losses, test_losses]

model = Transformer(114,3,1,128,512,4)
optimizer = AdamW(get_parameters(model), lr=.001, b1=.9, b2=.98, wd=1)
train_losses, test_losses = train(model, train_ds, train_targets, test_ds, test_targets, optimizer, 5000, lossfn=loss_fn)