8 thoughts on “Simple pseudocode for transformer decoding a la GPT”

1. The text includes the definition: “simplex(T) to denote a vector in [0, ∞)^T that sums to unity”. If all values are positive and sum to 1, wouldn’t it be simpler to say that it’s a vector in [0,1]^T. Or am I missing something simple? Also, does V need to be even for the positional embeddings to be well defined?

   Reply ↓

   • Thanks for the feedback.

     No, I don’t think you’re missing something. I agree that [0, 1]^T is clearer.

     I decided to follow Stan’s notation and call it simplex(T) instead of simplex(T – 1), even though a simplex embedded in R^T only has T – 1 dimensions as a manifold.

     As I wrote it, yes, you’d need V to be even or the indexing will overflow. This could be modified to have a more complicated boundary condition, so it’s not an intrinsic limitation on positional encodings. In practice, the sizes are chosen to maximize the ability to squeeze things into GPU, so they’re usually multiples of 32 or of 64. There are also transformer variants where the positional encodings are learned vectors, not just periodic functions.

     Reply ↓

2. This is fascinating, do you have the code somewhere where it could be copy pasted more easily? I tried to port the code over as well as I could into this gist: https://gist.github.com/ajcwebdev/f24266de249c1a62fa5c6fa4c44024c4

   Managed to get some half decent syntax highlighting by tricking GitHub into thinking it was JavaScript written in a Markdown. I also asked ChatGPT to write a JavaScript implementation (since it’s the only language I know well) using the pseudocode as guidance.

   I have no idea if this code is even remotely correct but here’s what its second output contained (it first complained it didn’t have enough information to implement half of the functions):

   “`js
   function softmax(arr) {
   const max = Math.max(…arr);
   const expArr = arr.map(x => Math.exp(x – max));
   const sum = expArr.reduce((acc, val) => acc + val, 0);
   return expArr.map(x => x / sum);
   }

   function gelu(x) {
   const cdf = 0.5 * (1.0 + Math.tanh((Math.sqrt(2 / Math.PI) * (x + 0.044715 * Math.pow(x, 3)))));
   return x * cdf;
   }

   function standardize(u) {
   const mean = u.reduce((acc, val) => acc + val, 0) / u.length;
   const std = Math.sqrt(u.map(x => Math.pow(x – mean, 2)).reduce((acc, val) => acc + val, 0) / u.length);
   return u.map(x => (x – mean) / std);
   }

   function attend(x, beta) {
   let N = x.length;
   let V = x[0].length;
   let K = beta.query[0].length;

   let q = new Array(N).fill().map(() => new Array(K).fill(0));
   let k = new Array(N).fill().map(() => new Array(K).fill(0));
   let v = new Array(N).fill().map(() => new Array(V).fill(0));

   for (let n = 0; n < N; n++) {
   for (let j = 0; j sum + xi * beta.query[idx][j], 0);
   k[n][j] = x[n].reduce((sum, xi, idx) => sum + xi * beta.key[idx][j], 0);
   }
   for (let j = 0; j sum + xi * beta.value[idx][j], 0);
   }
   }

   let y = new Array(N).fill().map(() => new Array(V).fill(0));

   for (let n = 0; n < N; n++) {
   let lp = new Array(N).fill(-Infinity);
   for (let i = 0; i sum + qj * k[i][idx], 0) / Math.sqrt(V);
   }
   let p = softmax(lp);
   let u = new Array(V).fill(0);
   for (let i = 0; i < N; i++) {
   for (let j = 0; j uj + x[n][idx]));
   }

   return y;
   }

   function feedForward(x, gamma) {
   const L = gamma[‘1’].length;
   let u = gamma[‘1’].map((g1i, idx) => g1i + x.reduce((sum, xi, jdx) => sum + xi * gamma[‘2’][jdx][idx], 0));
   let w = u.map(val => gelu(val));
   let y = gamma[‘3’].map((g3i, idx) => g3i + w.reduce((sum, wi, jdx) => sum + wi * gamma[‘4’][jdx][idx], 0));
   return standardize(x.map((xi, idx) => xi + y[idx]));
   }

   function logisticRegression(x, delta) {
   const T = delta.length;
   let log_probs = new Array(T).fill(0).map((_, idx) => x.reduce((sum, xi, jdx) => sum + xi * delta[idx][jdx], 0));
   return softmax(log_probs);
   }

   function decode(tok, alpha, betas, gammas, delta) {
   const N = tok.length;
   const V = alpha[0].length;
   let xs = new Array(N).fill().map(() => new Array(V).fill(0));

   // Embed input
   for (let n = 0; n < N; n++) {
   xs[n] = alpha[tok[n]];
   }

   const A = betas.length;

   for (let a = 0; a < A; a++) {
   xs = attend(xs, betas[a]);

   for (let n = 0; n < N; n++) {
   xs[n] = feedForward(xs[n], gammas[a]);
   }
   }

   return logisticRegression(xs[N – 1], delta);
   }
   “`

   Reply ↓

   • There are online repos that I link in the case study that provide working code. Andrej Karpathy has a nice two-hour video tutorial implementing it in PyTorch with GPU-friendly code. There’s a Google Colab notebook and a GitHub repo. So if you want to get to working code, that’s what I’d recommend.

     I only provided pseudocode, which is intended to be halway between mathematical notation and code that actually works. At least that’s the usual intent with pseudocode.

     Reply ↓

3. Here is a 2022 paper that starts: “Transformers … Descriptions are usually graphical, verbal, partial, or incremental. Despite their popularity, it seems no pseudocode has ever been published for any variant. … This report intends to rectify the situation for Transformers. … The essentially complete pseudocode is about 50 lines,”
   https://arxiv.org/abs/2207.09238

   Funny that both papers agree that 50 lines are needed.

   Reply ↓

   • Oops, I now see that Carpenter references the above paper. I thought that I was posting something he might not know about.

     Reply ↓

4. Thanks for this. I’m hoping to get a coworker into messing with transformers. This looks like a good companion piece.

   On a quick scan, is there something going on with the history variable in the Parameter estimation section? I see an unmatched bracket on the right hand side of an assignment. It makes me think I don’t understand what the context variable is quite doing either.

   Reply ↓

   • Thanks. I fixed it and uploaded a new draft and changed the link in the post to point to the new draft.

     Reply ↓