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.

training loss

test loss

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
from tinygrad.nn import LayerNorm
from tinygrad.nn.state import get_parameters
from tinygrad.nn.optim import AdamW
from tinygrad.helpers import CI
from tinygrad import TinyJit
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]

class MultiHeadAttention:
    def __init__(self, num_heads, embed_dim):
        self.num_heads, self.embed_dim = num_heads, embed_dim
        self.head_dim = self.embed_dim // self.num_heads

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

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

    def __call__(self, x, mask=None):
        bz = x.shape[0]
        QKV = [
            x.linear(weight, bias).reshape(
                bz, self.num_heads, -1, self.head_dim
            for weight, bias in [self.query, self.key, self.value]

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

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

        return out

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

    def __call__(self, x, mask=None):
        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 = vocab_size
        self.context_length = context_length
        self.embed_dim = embed_dim
        self.ffn_dim = ffn_dim
        self.num_heads = num_heads

        self.embed = Tensor.scaled_uniform(
            (context_length + vocab_size, embed_dim)
        self.blocks = [
            TransformerBlock(num_heads, embed_dim, ffn_dim)
            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(
    steps=1, # Only one step is needed for full batch training
    lossfn=lambda out, y: out.sparse_categorical_crossentropy(y),
    def train_step(x, y):
        # network
        out = model(x)[:,-1]
        loss = lossfn(out, y)
        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()
            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)


You must log in to answer this question.

Browse other questions tagged .