§5 (i–ii)
Project attention output to d if needed, add X, then LayerNorm over the last dimension. Mean should be ≈0 and variance ≈1 per token before learned γ/β change it.
§6 optional (i–ii)
FFN(x)=W₂·ReLU(W₁x+b₁)+b₂, d→dff→d; apply independently to all n rows; add residual and normalize.
§7 (i–ii)
Unembedding WU∈ℝd×|V| maps each token state to logits; softmax along vocabulary yields probabilities. Decode greedy argmax or sample with temperature/top-k.
§8 recommended complete module
class MyTransformer(nn.Module):
def __init__(self, vocab, d=64, heads=4, dff=128):
super().__init__()
self.emb = nn.Embedding(vocab,d)
self.attn = nn.MultiheadAttention(d,heads,batch_first=True)
self.n1,self.n2 = nn.LayerNorm(d),nn.LayerNorm(d)
self.ff = nn.Sequential(nn.Linear(d,dff),nn.ReLU(),nn.Linear(dff,d))
self.out = nn.Linear(d,vocab,bias=False)
def forward(self, ids):
x = self.emb(ids) + sinusoidal_pe(ids.size(1),self.emb.embedding_dim).to(ids.device)
mask = torch.triu(torch.ones(ids.size(1),ids.size(1),device=ids.device,dtype=torch.bool),1)
a,_ = self.attn(x,x,x,attn_mask=mask)
x = self.n1(x+a)
x = self.n2(x+self.ff(x))
return self.out(x)
Final result: input (batch,n) → logits (batch,n,|V|). The source’s def MyTransformer(nn.Module) must be class MyTransformer(nn.Module).
Check. Shift training: compare logits[:,:-1] with targets ids[:,1:]; never train a token to predict itself.