Switch to Adam optimizer:

WIP

shader/center.wgsl

@@ -0,0 +1,24 @@
+struct MatrixBatch {
+  size: vec4<u32>,
+  data: array<vec4<f32>>,
+}
+
+@group(0) @binding(0)
+var<storage, read> m: MatrixBatch;
+@group(0) @binding(1)
+var<storage, read_write> centered: MatrixBatch;
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+  centered.size = m.size;
+
+  let vc =  m.size.x / 4;
+  let ro = (global_id.z * m.size.y + global_id.y) * vc;
+
+  var s = vec4(0f);
+  for (var x = 0u; x < vc; x++) {
+    s += m.data[ro + x];
+  }
+
+  centered.data[ro + global_id.x] = m.data[ro + global_id.x] - vec4(s.x + s.y + s.z + s.w) / f32(m.size.x);
+}

shader/clip.wgsl

@@ -0,0 +1,28 @@
+struct ScalarBatch {
+  size: vec4<u32>,
+  data: array<f32>,
+}
+
+struct Scalar {
+  size: vec4<u32>,
+  data: array<f32>,
+}
+
+struct MatrixBatch {
+  size: vec4<u32>,
+  data: array<vec4<f32>>,
+}
+
+@group(0) @binding(0)
+var<storage, read> norm: ScalarBatch;
+@group(0) @binding(1)
+var<storage, read> threshold: Scalar;
+@group(0) @binding(2)
+var<storage, read_write> m: MatrixBatch;
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+  m.size = m.size;
+
+  m.data[(global_id.z * m.size.y + global_id.y) * m.size.x / 4 + global_id.x] *= min(1f, threshold.data[0] / norm.data[global_id.z]);
+}

shader/cpow.wgsl

@@ -0,0 +1,21 @@
+alias number = f32;
+
+struct MatrixBatch {
+  size: vec4<u32>,
+  data: array<vec4<number>>,
+}
+
+@group(0) @binding(0)
+var<storage, read> l: MatrixBatch;
+@group(0) @binding(1)
+var<storage, read> r_or_result: MatrixBatch;
+@group(0) @binding(2)
+var<storage, read_write> result: MatrixBatch;
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+  result.size = l.size;
+
+  var resi = (global_id.z * result.size.y + global_id.y) * result.size.x / 4 + global_id.x;
+  result.data[resi] = pow(l.data[resi], r_or_result.data[resi]);
+}

shader/csqr.wgsl

@@ -0,0 +1,19 @@
+alias number = f32;
+
+struct MatrixBatch {
+  size: vec4<u32>,
+  data: array<vec4<number>>,
+}
+
+@group(0) @binding(0)
+var<storage, read> v_or_result: MatrixBatch;
+@group(0) @binding(1)
+var<storage, read_write> result: MatrixBatch;
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+  result.size = v_or_result.size;
+
+  var resi = (global_id.z * result.size.y + global_id.y) * result.size.x / 4 + global_id.x;
+  result.data[resi] = v_or_result.data[resi] * v_or_result.data[resi];
+}

shader/csqrt.wgsl

@@ -0,0 +1,19 @@
+alias number = f32;
+
+struct MatrixBatch {
+  size: vec4<u32>,
+  data: array<vec4<number>>,
+}
+
+@group(0) @binding(0)
+var<storage, read> v_or_result: MatrixBatch;
+@group(0) @binding(1)
+var<storage, read_write> result: MatrixBatch;
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+  result.size = v_or_result.size;
+
+  var resi = (global_id.z * result.size.y + global_id.y) * result.size.x / 4 + global_id.x;
+  result.data[resi] = sqrt(v_or_result.data[resi]);
+}

shader/csum3.wgsl

@@ -0,0 +1,23 @@
+alias number = f32;
+
+struct MatrixBatch {
+  size: vec4<u32>,
+  data: array<vec4<number>>,
+}
+
+@group(0) @binding(0)
+var<storage, read> l: MatrixBatch;
+@group(0) @binding(1)
+var<storage, read> m: MatrixBatch;
+@group(0) @binding(2)
+var<storage, read> r_or_result: MatrixBatch;
+@group(0) @binding(3)
+var<storage, read_write> result: MatrixBatch;
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+  result.size = l.size;
+
+  var resi = (global_id.z * result.size.y + global_id.y) * result.size.x / 4 + global_id.x;
+  result.data[resi] = l.data[resi] + m.data[resi] + r_or_result.data[resi];
+}

shader/dlds.wgsl

@@ -1,48 +0,0 @@
-struct Scalar {
-  size: vec4<u32>,
-  data: array<f32>,
-}
-
-struct MatrixBatch {
-  size: vec4<u32>,
-  data: array<f32>,
-}
-
-@group(0) @binding(0)
-var<storage, read> ds: Scalar;
-@group(0) @binding(1)
-var<storage, read> wa: MatrixBatch;
-@group(0) @binding(2)
-var<storage, read> v: MatrixBatch;
-@group(0) @binding(3)
-var<storage, read> otdldo: MatrixBatch;
-@group(0) @binding(4)
-var<storage, read_write> dlds: MatrixBatch;
-
-@compute @workgroup_size(1)
-fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
-  dlds.size = wa.size;
-
-  var s = 0f;
-  var li = (global_id.z * v.size.y + global_id.x) * v.size.x;
-  var ri = global_id.z * otdldo.size.x * otdldo.size.y + global_id.y;
-
-  for (var i = 0u; i < v.size.x; i++) {
-    var cs = 0f;
-    var lj = global_id.z * wa.size.y * wa.size.x + global_id.y * wa.size.x;
-    var mj = global_id.z * v.size.x * v.size.y + i;
-    for (var j = 0u; j < wa.size.x; j++) {
-      cs += wa.data[lj] * v.data[mj];
-      lj++;
-      mj += v.size.x;
-    }
-
-    s += (v.data[li] - cs) * otdldo.data[ri];
-
-    li++;
-    ri += otdldo.size.x;
-  }
-
-  let resi = (global_id.z * dlds.size.y + global_id.y) * dlds.size.x + global_id.x;
-  dlds.data[resi] = ds.data[0] * wa.data[resi] * s;
-}

shader/gradinp.wgsl

@@ -1,47 +0,0 @@
-struct MatrixBatch {
-  size: vec4<u32>,
-  data: array<f32>,
-}
-
-@group(0) @binding(0)
-var<storage, read> ilogits: MatrixBatch;
-@group(0) @binding(0)
-var<storage, read> tilogits: MatrixBatch;
-@group(0) @binding(1)
-var<storage, read> dldsq: MatrixBatch;
-@group(0) @binding(2)
-var<storage, read> dldsk: MatrixBatch;
-@group(0) @binding(3)
-var<storage, read> totdldowa: MatrixBatch;
-@group(0) @binding(4)
-var<storage, read> tqry: MatrixBatch;
-@group(0) @binding(5)
-var<storage, read> tkey: MatrixBatch;
-@group(0) @binding(6)
-var<storage, read> tval: MatrixBatch;
-@group(0) @binding(7)
-var<storage, read_write> ginp: MatrixBatch;
-
-@compute @workgroup_size(1)
-fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
-  ginp.size.x = tval.size.x;
-  ginp.size.y = ilogits.size.x;
-  ginp.size.z = ilogits.size.z;
-
-  var ts = 0f;
-  var li = global_id.z * ilogits.size.y * ilogits.size.x + global_id.y;
-  for (var i = 0u; i < ilogits.size.y; i++) {
-    var s = 0f;
-    var lj = global_id.z * totdldowa.size.x * totdldowa.size.y + i * totdldowa.size.x;
-    var rj = global_id.z * tval.size.x * tval.size.y + global_id.x;
-    for (var j = 0u; j < totdldowa.size.x; j++) {
-      s += totdldowa.data[lj] * tval.data[rj] + dldsq.data[lj] * tkey.data[rj] + dldsk.data[lj] * tqry.data[rj];
-      lj++;
-      rj += tval.size.x;
-    }
-    ts += ilogits.data[li] * s;
-    li += ilogits.size.x;
-  }
-
-  ginp.data[global_id.z * ginp.size.y * ginp.size.x + global_id.y * ginp.size.x + global_id.x] = ts;
-}

shader/gradtattw.wgsl

@@ -1,28 +0,0 @@
-struct MatrixBatch {
-  size: vec4<u32>,
-  data: array<f32>,
-}
-
-@group(0) @binding(0)
-var<storage, read> e: MatrixBatch;
-@group(0) @binding(1)
-var<storage, read> dlds: MatrixBatch;
-@group(0) @binding(2)
-var<storage, read_write> gtattw: MatrixBatch;
-
-@compute @workgroup_size(1)
-fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
-  gtattw.size.z = e.size.z;
-  gtattw.size.y = e.size.x;
-  gtattw.size.x = e.size.x;
-
-  var s = 0f;
-  var li = global_id.z * e.size.x * e.size.y + global_id.x;
-  var ri = global_id.z * dlds.size.x * dlds.size.y + global_id.y;
-  for (var i = 0u; i < e.size.y; i++) {
-    s += e.data[li] * dlds.data[ri];
-    li += e.size.x;
-    ri += dlds.size.x;
-  }
-  gtattw.data[global_id.z * gtattw.size.y * gtattw.size.x + global_id.y * gtattw.size.x + global_id.x] = s;
-}

shader/gradtout.wgsl

@@ -1,28 +0,0 @@
-struct MatrixBatch {
-  size: vec4<u32>,
-  data: array<f32>,
-}
-
-@group(0) @binding(0)
-var<storage, read> a: MatrixBatch;
-@group(0) @binding(1)
-var<storage, read> dldo: MatrixBatch;
-@group(0) @binding(2)
-var<storage, read_write> gtout: MatrixBatch;
-
-@compute @workgroup_size(1)
-fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
-  gtout.size = a.size;
-  gtout.size.y = dldo.size.x;
-
-  var s = 0f;
-  var ai = global_id.z * a.size.y * a.size.x + global_id.x;
-  var di = global_id.z * dldo.size.x * dldo.size.y + global_id.y;
-  for (var i = 0u; i < a.size.y; i++) {
-    s += a.data[ai] * dldo.data[di];
-    ai += a.size.x;
-    di += dldo.size.x;
-  }
-
-  gtout.data[global_id.z * gtout.size.y * gtout.size.x + global_id.y * gtout.size.x + global_id.x] = s;
-}

shader/layernormgrad.wgsl

@@ -0,0 +1,51 @@
+struct MatrixBatch {
+  size: vec4<u32>,
+  data: array<vec4<f32>>,
+}
+
+@group(0) @binding(0) var<storage, read> m: MatrixBatch;
+@group(0) @binding(1) var<storage, read> dldm: MatrixBatch;
+@group(0) @binding(2) var<storage, read_write> grad: MatrixBatch;
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+  grad.size = m.size;
+
+  let vc = m.size.x / 4;
+  let ro = (global_id.z * m.size.y + global_id.y) * vc;
+
+  var sum = vec4(0f);
+  var sumsq = vec4(0f);
+  for (var i = 0u; i < vc; i++) {
+    let v = m.data[ro + i];
+    sum += v;
+    sumsq += v * v;
+  }
+
+  let mean = (sum.x + sum.y + sum.z + sum.w) / f32(m.size.x);
+  let variance = (sumsq.x + sumsq.y + sumsq.z + sumsq.w) / f32(m.size.x) - mean * mean;
+  let inv_std = inverseSqrt((sumsq.x + sumsq.y + sumsq.z + sumsq.w) / f32(m.size.x) - mean * mean + 1e-5);
+
+  // Recompute x̂
+  var sum_grad = vec4(0.0);
+  var sum_grad_xhat = vec4(0.0);
+
+  for (var i = 0u; i < vc; i++) {
+    let v = m.data[ro + i];
+    let g = dldm.data[ro + i];
+    let xhat = (v - vec4(mean)) * vec4(inv_std);
+    sum_grad += g;
+    sum_grad_xhat += g * xhat;
+  }
+
+  let mg = (sum_grad.x + sum_grad.y + sum_grad.z + sum_grad.w) / f32(m.size.x);
+  let mgxhat = (sum_grad_xhat.x + sum_grad_xhat.y + sum_grad_xhat.z + sum_grad_xhat.w) / f32(m.size.x);
+
+  for (var i = 0u; i < vc; i++) {
+    let v = m.data[ro + i];
+    let g = dldm.data[ro + i];
+    let xhat = (v - vec4(mean)) * vec4(inv_std);
+    let dx = vec4(inv_std) * (g - vec4(mg) - xhat * vec4(mgxhat));
+    grad.data[ro + i] = dx;
+  }
+}

shader/multrans.wgsl

@@ -0,0 +1,32 @@
+alias number = f32;
+
+struct MatrixBatch {
+  size: vec4<u32>,
+  data: array<number>,
+}
+
+@group(0) @binding(0)
+var<storage, read> l: MatrixBatch;
+@group(0) @binding(1)
+var<storage, read> r: MatrixBatch;
+@group(0) @binding(2)
+var<storage, read_write> result: MatrixBatch;
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+  result.size.x = l.size.y;
+  result.size.y = r.size.x;
+  result.size.z = l.size.z;
+
+  var s = 0f;
+  var li = global_id.z * l.size.x * l.size.y + global_id.x * l.size.x;
+  var ri = global_id.z * r.size.x * r.size.y + global_id.y;
+  for (var i = 0u; i < l.size.x; i++) {
+    s += l.data[li] * r.data[ri];
+    li++;
+    ri += r.size.x;
+  }
+
+  let l = global_id.z * result.size.y * result.size.x;
+  result.data[l + global_id.y * result.size.x + global_id.x] = s;
+}

shader/normalize.wgsl

@@ -0,0 +1,24 @@
+struct MatrixBatch {
+  size: vec4<u32>,
+  data: array<vec4<f32>>,
+}
+
+@group(0) @binding(0)
+var<storage, read> c: MatrixBatch;
+@group(0) @binding(1)
+var<storage, read_write> result: MatrixBatch;
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+  result.size = c.size;
+
+  let vc = c.size.x / 4;
+  let ro = (global_id.z * c.size.y + global_id.y) * vc;
+
+  var s = vec4(0f);
+  for (var x = 0u; x < vc; x++) {
+    s += c.data[ro + x] * c.data[ro + x];
+  }
+
+  result.data[ro + global_id.x] = c.data[ro + global_id.x] * vec4(inverseSqrt((s.x + s.y + s.z + s.w) / f32(c.size.x) + 1e-5f));
+}

shader/normh.wgsl

@@ -0,0 +1,30 @@
+alias number = f32;
+
+struct MatrixBatch {
+  size: vec4<u32>,
+  data: array<number>,
+}
+
+struct VectorBatch {
+  size: vec4<u32>,
+  data: array<number>,
+}
+
+@group(0) @binding(0)
+var<storage, read> m: MatrixBatch;
+@group(0) @binding(1)
+var<storage, read_write> result: VectorBatch;
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+  result.size.x = m.size.y;
+  result.size.z = m.size.z;
+
+  var s = 0f;
+  let lo = global_id.z * m.size.x * m.size.y + global_id.x * m.size.x;
+  for (var i = 0u; i < m.size.x; i++) {
+    s += m.data[lo + i] * m.data[lo + i];
+  }
+
+  result.data[global_id.z * result.size.x + global_id.x] = s;
+}

shader/normv.wgsl

@@ -0,0 +1,29 @@
+alias number = f32;
+
+struct VectorBatch {
+  size: vec4<u32>,
+  data: array<number>,
+}
+
+struct ScalarBatch {
+  size: vec4<u32>,
+  data: array<number>,
+}
+
+@group(0) @binding(0)
+var<storage, read> v: VectorBatch;
+@group(0) @binding(1)
+var<storage, read_write> result: ScalarBatch;
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+  result.size.z = v.size.z;
+
+  var s = 0f;
+  let lo = global_id.z * v.size.x;
+  for (var i = 0u; i < v.size.x; i++) {
+    s += v.data[lo + i];
+  }
+
+  result.data[global_id.z] = sqrt(s);
+}

shader/otdldo.wgsl

@@ -1,29 +0,0 @@
-struct MatrixBatch {
-  size: vec4<u32>,
-  data: array<f32>,
-}
-
-@group(0) @binding(0)
-var<storage, read> dldo: MatrixBatch;
-@group(0) @binding(1)
-var<storage, read> tout: MatrixBatch;
-@group(0) @binding(2)
-var<storage, read_write> otdldo: MatrixBatch;
-
-@compute @workgroup_size(1)
-fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
-  otdldo.size.x = dldo.size.y;
-  otdldo.size.y = tout.size.x;
-  otdldo.size.z = dldo.size.z;
-
-  var s = 0f;
-  var li = (global_id.z * dldo.size.y + global_id.x) * dldo.size.x;
-  var ri = global_id.z * tout.size.x * tout.size.y + global_id.y;
-  for (var i = 0u; i < dldo.size.x; i++) {
-    s += dldo.data[li] * tout.data[ri];
-    li++;
-    ri += tout.size.x;
-  }
-
-  otdldo.data[(global_id.z * otdldo.size.y + global_id.y) * otdldo.size.x + global_id.x] = s;
-}

shader/softmaxgrad.wgsl

@@ -0,0 +1,26 @@
+struct MatrixBatch {
+  size: vec4<u32>,
+  data: array<vec4<f32>>,
+}
+
+@group(0) @binding(0)
+var<storage, read> m: MatrixBatch;
+@group(0) @binding(1)
+var<storage, read> dldm: MatrixBatch;
+@group(0) @binding(2)
+var<storage, read_write> grad: MatrixBatch;
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+  grad.size = m.size;
+
+  let vc = m.size.x / 4;
+  let ro = (global_id.z * m.size.y + global_id.y) * vc;
+
+  var s = 0f;
+  for (var x = 0u; x < vc; x++) {
+    s += dot(m.data[ro + x], dldm.data[ro + x]);
+  }
+
+  grad.data[ro + global_id.x] = m.data[ro + global_id.x] * (dldm.data[ro + global_id.x] - vec4(s));
+}

shader/totdldowa.wgsl

@@ -1,29 +0,0 @@
-struct MatrixBatch {
-  size: vec4<u32>,
-  data: array<f32>,
-}
-
-@group(0) @binding(0)
-var<storage, read> wa: MatrixBatch;
-@group(0) @binding(1)
-var<storage, read> otdldo: MatrixBatch;
-@group(0) @binding(2)
-var<storage, read_write> totdldowa: MatrixBatch;
-
-@compute @workgroup_size(1)
-fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
-  totdldowa.size.x = otdldo.size.y;
-  totdldowa.size.y = wa.size.x;
-  totdldowa.size.z = otdldo.size.z;
-
-  var s = 0f;
-  var li = global_id.z * otdldo.size.x * otdldo.size.y + global_id.x * otdldo.size.x;
-  var ri = global_id.z * wa.size.x * wa.size.y + global_id.y;
-  for (var i = 0u; i < otdldo.size.x; i++) {
-    s += otdldo.data[li] * wa.data[ri];
-    li++;
-    ri += wa.size.x;
-  }
-
-  totdldowa.data[(global_id.z * totdldowa.size.y + global_id.y) * totdldowa.size.x + global_id.x] = s;
-}

shader/translmul.wgsl

@@ -0,0 +1,31 @@
+alias number = f32;
+
+struct MatrixBatch {
+  size: vec4<u32>,
+  data: array<number>,
+}
+
+@group(0) @binding(0)
+var<storage, read> l: MatrixBatch;
+@group(0) @binding(1)
+var<storage, read> r: MatrixBatch;
+@group(0) @binding(2)
+var<storage, read_write> result: MatrixBatch;
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+  result.size.x = r.size.x;
+  result.size.y = l.size.x;
+  result.size.z = l.size.z;
+
+  var s = 0f;
+  var li = global_id.z * l.size.y * l.size.x + global_id.y;
+  var ri = global_id.z * r.size.y * r.size.x + global_id.x;
+  for (var i = 0u; i < l.size.y; i++) {
+    s += l.data[li] * r.data[ri];
+    li += l.size.x;
+    ri += r.size.x;
+  }
+
+  result.data[global_id.z * result.size.y * result.size.x + global_id.y * result.size.x + global_id.x] = s;
+}

shader/transrmul.wgsl

@@ -0,0 +1,34 @@
+alias number = f32;
+
+struct PackedMatrixBatch {
+  size: vec4<u32>,
+  data: array<vec4<number>>,
+}
+
+struct MatrixBatch {
+  size: vec4<u32>,
+  data: array<number>,
+}
+
+@group(0) @binding(0)
+var<storage, read> l: PackedMatrixBatch;
+@group(0) @binding(1)
+var<storage, read> tr: PackedMatrixBatch;
+@group(0) @binding(2)
+var<storage, read_write> result: MatrixBatch;
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+  result.size = l.size;
+  result.size.x = tr.size.y;
+
+  var s = 0f;
+  let vc = l.size.x / 4;
+  let lo = (global_id.z * l.size.y + global_id.y) * vc;
+  let ro = (global_id.z * tr.size.y + global_id.x) * vc;
+  for (var i = 0u; i < vc; i++) {
+    s += dot(l.data[lo + i], tr.data[ro + i]);
+  }
+
+  result.data[global_id.z * result.size.y * result.size.x + global_id.y * result.size.x + global_id.x] = s;
+}

src/data/gpu/gpu-struct.ts

@@ -8,14 +8,24 @@ interface ComputePipeline {
 
 const PIPELINES: Map<GPUDevice, Map<string, ComputePipeline>> = new Map();
 
-const COMPWISE_SHADER = {
+const COMPWISE1_SHADER = {
+  sqr: fs.readFileSync('shader/csqr.wgsl').toString(),
+  sqrt: fs.readFileSync('shader/csqrt.wgsl').toString(),
+};
+
+const COMPWISE2_SHADER = {
   sum: fs.readFileSync('shader/csum.wgsl').toString(),
   sub: fs.readFileSync('shader/csub.wgsl').toString(),
   mul: fs.readFileSync('shader/cmul.wgsl').toString(),
   div: fs.readFileSync('shader/cdiv.wgsl').toString(),
+  pow: fs.readFileSync('shader/cpow.wgsl').toString(),
   scale: fs.readFileSync('shader/cscale.wgsl').toString(),
 };
 
+const COMPWISE3_SHADER = {
+  sum: fs.readFileSync('shader/csum3.wgsl').toString(),
+};
+
 export function createComputePipeline(
   device: GPUDevice,
   name: string,
@@ -145,12 +155,14 @@ export abstract class GPUStruct<A extends Uint32Array | Int32Array | Float32Arra
     await this._device.queue.onSubmittedWorkDone();
 
     await this._outbuffer.mapAsync(GPUMapMode.READ);
-    this.data.set(new this._array(this._outbuffer.getMappedRange()).slice(this.size.length));
+    this.data.set(new this._array(this._outbuffer.getMappedRange()).slice(this.size.length, this.size.length + this.data.length));
     this._outbuffer.unmap();
     this._dirty = null;
     this._op = null;
   }
 
+  public abstract get(): Promise<unknown>;
+
   public buffer(): GPUBuffer {
     if (this._dirty === 'front') {
       this.write();
@@ -188,7 +200,9 @@ export abstract class GPUStruct<A extends Uint32Array | Int32Array | Float32Arra
 
     pass.setPipeline(pipeline);
     pass.setBindGroup(0, this._device.createBindGroup({
-      layout, entries: [
+      label: `${name}${ip ? '_inplace_' : '_'}${this.type}`,
+      layout,
+      entries: [
         ...args.map((s, i) => ({ binding: i, resource: { buffer: s.buffer() } })),
         { binding: args.length + +ip, resource: { buffer: this.buffer() } },
       ],
@@ -247,44 +261,94 @@ export abstract class GPUStruct<A extends Uint32Array | Int32Array | Float32Arra
     }
   }
 
-  public static compwise<
+  public static compwise1<
+    A extends Uint32Array | Int32Array | Float32Array,
+    S extends GPUStruct<A>,
+  >(op: keyof typeof COMPWISE1_SHADER, v: S, result: S): S {
+    this.assertColocated(v, result);
+    this.assertPackable(4, v, result);
+
+    return result.calculate(`${op}1`, COMPWISE1_SHADER[op], result === v ? [] : [v], result === v, 4);
+  }
+
+  public static csqr<
+    A extends Uint32Array | Int32Array | Float32Array,
+    S extends GPUStruct<A>,
+  >(v: S, result: S): S {
+    return GPUStruct.compwise1('sqr', v, result);
+  }
+
+  public static csqrt<
+    A extends Uint32Array | Int32Array | Float32Array,
+    S extends GPUStruct<A>,
+  >(v: S, result: S): S {
+    return GPUStruct.compwise1('sqrt', v, result);
+  }
+
+  public static compwise2<
     A extends Uint32Array | Int32Array | Float32Array,
     S extends GPUStruct<A>,
-  >(op: keyof typeof COMPWISE_SHADER, l: S, r: S, result: S): S {
+  >(op: keyof typeof COMPWISE2_SHADER, l: S, r: S, result: S): S {
     this.assertUnique(l, r);
     this.assertUnique(l, result);
     this.assertColocated(l, r, result);
     this.assertPackable(4, l, r, result);
 
-    return result.calculate(op, COMPWISE_SHADER[op], result === r ? [l] : [l, r], result === r, 4);
+    return result.calculate(`${op}2`, COMPWISE2_SHADER[op], result === r ? [l] : [l, r], result === r, 4);
+  }
+
+  public static compwise3<
+    A extends Uint32Array | Int32Array | Float32Array,
+    S extends GPUStruct<A>,
+  >(op: keyof typeof COMPWISE3_SHADER, l: S, m: S, r: S, result: S): S {
+    this.assertUnique(l, m, r);
+    this.assertUnique(l, m, result);
+    this.assertColocated(l, m, r, result);
+    this.assertPackable(4, l, m, r, result);
+
+    return result.calculate(`${op}3`, COMPWISE3_SHADER[op], result === r ? [l, m] : [l, m, r], result === r, 4);
   }
 
   public static csum<
     A extends Uint32Array | Int32Array | Float32Array,
     S extends GPUStruct<A>,
   >(l: S, r: S, result: S): S {
-    return GPUStruct.compwise('sum', l, r, result);
+    return GPUStruct.compwise2('sum', l, r, result);
+  }
+
+  public static csum3<
+    A extends Uint32Array | Int32Array | Float32Array,
+    S extends GPUStruct<A>,
+  >(l: S, m: S, r: S, result: S): S {
+    return GPUStruct.compwise3('sum', l, m, r, result);
   }
 
   public static csub<
     A extends Uint32Array | Int32Array | Float32Array,
     S extends GPUStruct<A>,
   >(l: S, r: S, result: S): S {
-    return GPUStruct.compwise('sub', l, r, result);
+    return GPUStruct.compwise2('sub', l, r, result);
   }
 
   public static cmul<
     A extends Uint32Array | Int32Array | Float32Array,
     S extends GPUStruct<A>,
   >(l: S, r: S, result: S): S {
-    return GPUStruct.compwise('mul', l, r, result);
+    return GPUStruct.compwise2('mul', l, r, result);
   }
 
   public static cdiv<
     A extends Uint32Array | Int32Array | Float32Array,
     S extends GPUStruct<A>,
   >(l: S, r: S, result: S): S {
-    return GPUStruct.compwise('div', l, r, result);
+    return GPUStruct.compwise2('div', l, r, result);
+  }
+
+  public static cpow<
+    A extends Uint32Array | Int32Array | Float32Array,
+    S extends GPUStruct<A>,
+  >(l: S, r: S, result: S): S {
+    return GPUStruct.compwise2('pow', l, r, result);
   }
 
   public static cscale<
@@ -296,6 +360,6 @@ export abstract class GPUStruct<A extends Uint32Array | Int32Array | Float32Arra
     this.assertColocated(l, r, result);
     this.assertPackable(4, r, result);
 
-    return result.calculate('scale', COMPWISE_SHADER.scale, result === r ? [l] : [l, r], result === r, 4);
+    return result.calculate('scale', COMPWISE2_SHADER.scale, result === r ? [l] : [l, r], result === r, 4);
   }
 }

src/data/gpu/matrix-batch.ts

@@ -9,6 +9,11 @@ const SHADER_SCALE = fs.readFileSync('shader/bscale.wgsl').toString();
 const SHADER_ISUM = fs.readFileSync('shader/isum.wgsl').toString();
 const SHADER_MUL = fs.readFileSync('shader/mul.wgsl').toString();
 const SHADER_TRANS = fs.readFileSync('shader/transpose.wgsl').toString();
+const SHADER_TRANSLMUL = fs.readFileSync('shader/translmul.wgsl').toString();
+const SHADER_TRANSRMUL = fs.readFileSync('shader/transrmul.wgsl').toString();
+const SHADER_MULTRANS = fs.readFileSync('shader/multrans.wgsl').toString();
+const SHADER_NORMH = fs.readFileSync('shader/normh.wgsl').toString();
+const SHADER_NORMV = fs.readFileSync('shader/normv.wgsl').toString();
 const SHADER_AVG = fs.readFileSync('shader/avg.wgsl').toString();
 const SHADER_TOT = fs.readFileSync('shader/tot.wgsl').toString();
 
@@ -173,6 +178,16 @@ export class MatrixBatch<
     return GPUStruct.csum(l, r, result);
   }
 
+  public static sum3<
+    L extends number,
+    H extends number,
+    W extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>,
+  >(l: MatrixBatch<L, H, W, A, T>, m: MatrixBatch<L, H, W, A, T>, r: MatrixBatch<L, H, W, A, T>, result: MatrixBatch<L, H, W, A, T> = new MatrixBatch(l._device, l._array, l.layers, l.height, l.width)): MatrixBatch<L, H, W, A, T> {
+    return GPUStruct.csum3(l, m, r, result);
+  }
+
   public static sub<
     L extends number,
     H extends number,
@@ -253,6 +268,84 @@ export class MatrixBatch<
     return result.calculate('transpose', SHADER_TRANS, [m]);
   }
 
+  public static translmul<
+    L extends number,
+    H extends number,
+    S extends number,
+    W extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>>,
+  >(
+    l: MatrixBatch<L, S, H, A, T>,
+    r: MatrixBatch<L, S, W, A, T>,
+    result: MatrixBatch<L, H, W, A, T> = new MatrixBatch(l._device, l._array, l.layers, l.width, r.width),
+  ): MatrixBatch<L, H, W, A, T> {
+    this.assertUnique(l, r, result);
+    this.assertColocated(l, r, result);
+    this.assertPackable(1, l, r, result);
+
+    return result.calculate('translmul', SHADER_TRANSLMUL, [l, r]);
+  }
+
+  public static transrmul<
+    L extends number,
+    H extends number,
+    S extends number,
+    W extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>>,
+  >(
+    l: MatrixBatch<L, H, S, A, T>,
+    tr: MatrixBatch<L, W, S, A, T>,
+    result: MatrixBatch<L, H, W, A, T> = new MatrixBatch(l._device, l._array, l.layers, l.height, tr.height),
+  ): MatrixBatch<L, H, W, A, T> {
+    this.assertUnique(l, tr, result);
+    this.assertColocated(l, tr, result);
+    this.assertPackable(1, l, tr, result);
+
+    return result.calculate('transrmul', SHADER_TRANSRMUL, [l, tr]);
+  }
+
+  public static multrans<
+    L extends number,
+    H extends number,
+    S extends number,
+    W extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>>,
+  >(
+    l: MatrixBatch<L, W, S, A, T>,
+    r: MatrixBatch<L, S, H, A, T>,
+    result: MatrixBatch<L, H, W, A, T> = new MatrixBatch(l._device, l._array, l.layers, r.width, l.height),
+  ): MatrixBatch<L, H, W, A, T> {
+    this.assertUnique(l, r, result);
+    this.assertColocated(l, r, result);
+    this.assertPackable(1, l, r, result);
+
+    return result.calculate('multrans', SHADER_MULTRANS, [l, r]);
+  }
+
+  public static norm<
+    L extends number,
+    H extends number,
+    W extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>>,
+  >(
+    m: MatrixBatch<L, H, W, A, T>,
+    t: VectorBatch<L, H, A, T> = new VectorBatch(m._device, m._array, m.layers, m.height),
+    result: ScalarBatch<L, A, T> = new ScalarBatch(m._device, m._array, m.layers),
+  ): ScalarBatch<L, A, T> {
+    this.assertUnique(m, result);
+    this.assertColocated(m, result);
+    this.assertPackable(4, m);
+
+    t.calculate('normh', SHADER_NORMH, [m]);
+    result.calculate('normv', SHADER_NORMV, [t]);
+
+    return result;
+  }
+
   public static avg<
     L extends number,
     H extends number,

src/data/gpu/matrix.ts

@@ -8,6 +8,11 @@ import { Vector } from './vector';
 const SHADER_ISUM = fs.readFileSync('shader/isum.wgsl').toString();
 const SHADER_MUL = fs.readFileSync('shader/mul.wgsl').toString();
 const SHADER_TRANS = fs.readFileSync('shader/transpose.wgsl').toString();
+const SHADER_TRANSLMUL = fs.readFileSync('shader/translmul.wgsl').toString();
+const SHADER_TRANSRMUL = fs.readFileSync('shader/transrmul.wgsl').toString();
+const SHADER_MULTRANS = fs.readFileSync('shader/multrans.wgsl').toString();
+const SHADER_NORMH = fs.readFileSync('shader/normh.wgsl').toString();
+const SHADER_NORMV = fs.readFileSync('shader/normv.wgsl').toString();
 const SHADER_UNAVG = fs.readFileSync('shader/unavg.wgsl').toString();
 
 export class Matrix<
@@ -133,6 +138,22 @@ export class Matrix<
     }
   }
 
+  static esqr<
+    H extends number,
+    W extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+  >(v: Matrix<H, W, A>, result: Matrix<H, W, A> = new Matrix<H, W, A>(v._device, v._array, v.height, v.width)): Matrix<H, W, A> {
+    return GPUStruct.csqr(v, result);
+  }
+
+  static esqrt<
+    H extends number,
+    W extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+  >(v: Matrix<H, W, A>, result: Matrix<H, W, A> = new Matrix<H, W, A>(v._device, v._array, v.height, v.width)): Matrix<H, W, A> {
+    return GPUStruct.csqrt(v, result);
+  }
+
   public static sum<
     H extends number,
     W extends number,
@@ -142,6 +163,15 @@ export class Matrix<
     return GPUStruct.csum(l, r, result);
   }
 
+  public static sum3<
+    H extends number,
+    W extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>,
+  >(l: Matrix<H, W, A, T>, m: Matrix<H, W, A, T>, r: Matrix<H, W, A, T>, result: Matrix<H, W, A, T> = new Matrix(l._device, l._array, l.height, l.width)): Matrix<H, W, A, T> {
+    return GPUStruct.csum3(l, m, r, result);
+  }
+
   public static sub<
     H extends number,
     W extends number,
@@ -170,6 +200,15 @@ export class Matrix<
     return result.calculate('isum', SHADER_ISUM, [lis, m], false, 4);
   }
 
+  public static ediv<
+    H extends number,
+    W extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>,
+  >(l: Matrix<H, W, A, T>, r: Matrix<H, W, A, T>, result: Matrix<H, W, A, T> = new Matrix(l._device, l._array, l.height, l.width)): Matrix<H, W, A, T> {
+    return GPUStruct.cdiv(l, r, result);
+  }
+
   public static scale<
     H extends number,
     W extends number,
@@ -213,6 +252,80 @@ export class Matrix<
     return result.calculate('transpose', SHADER_TRANS, [m]);
   }
 
+  public static translmul<
+    H extends number,
+    S extends number,
+    W extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>>,
+  >(
+    l: Matrix<S, H, A, T>,
+    r: Matrix<S, W, A, T>,
+    result: Matrix<H, W, A, T> = new Matrix(l._device, l._array, l.width, r.width),
+  ): Matrix<H, W, A, T> {
+    this.assertUnique(l, r, result);
+    this.assertColocated(l, r, result);
+    this.assertPackable(1, l, r, result);
+
+    return result.calculate('translmul', SHADER_TRANSLMUL, [l, r]);
+  }
+
+  public static transrmul<
+    H extends number,
+    S extends number,
+    W extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>>,
+  >(
+    l: Matrix<H, S, A, T>,
+    tr: Matrix<W, S, A, T>,
+    result: Matrix<H, W, A, T> = new Matrix(l._device, l._array, l.height, tr.height),
+  ): Matrix<H, W, A, T> {
+    this.assertUnique(l, tr, result);
+    this.assertColocated(l, tr, result);
+    this.assertPackable(1, l, tr, result);
+
+    return result.calculate('transrmul', SHADER_TRANSRMUL, [l, tr]);
+  }
+
+  public static multrans<
+    H extends number,
+    S extends number,
+    W extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>>,
+  >(
+    l: Matrix<W, S, A, T>,
+    r: Matrix<S, H, A, T>,
+    result: Matrix<H, W, A, T> = new Matrix(l._device, l._array, r.width, l.height),
+  ): Matrix<H, W, A, T> {
+    this.assertUnique(l, r, result);
+    this.assertColocated(l, r, result);
+    this.assertPackable(1, l, r, result);
+
+    return result.calculate('multrans', SHADER_MULTRANS, [l, r]);
+  }
+
+  public static norm<
+    H extends number,
+    W extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>>,
+  >(
+    m: Matrix<H, W, A, T>,
+    t: Vector<H, A, T> = new Vector(m._device, m._array, m.height),
+    result: Scalar<A, T> = new Scalar(m._device, m._array),
+  ): Scalar<A, T> {
+    this.assertUnique(m, t, result);
+    this.assertColocated(m, t, result);
+    this.assertPackable(1, m);
+
+    t.calculate('normh', SHADER_NORMH, [m]);
+    result.calculate('normv', SHADER_NORMV, [t]);
+
+    return result;
+  }
+
   public static unavg<
     L extends number,
     H extends number,

src/data/gpu/scalar-batch.ts

@@ -104,6 +104,22 @@ export class ScalarBatch<
     }
   }
 
+  public static sqr<
+    L extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>,
+  >(v: ScalarBatch<L, A, T>, result: ScalarBatch<L, A, T> = new ScalarBatch(v._device, v._array, v.layers)): ScalarBatch<L, A, T> {
+    return GPUStruct.csqr(v, result);
+  }
+
+  public static sqrt<
+    L extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>,
+  >(v: ScalarBatch<L, A, T>, result: ScalarBatch<L, A, T> = new ScalarBatch(v._device, v._array, v.layers)): ScalarBatch<L, A, T> {
+    return GPUStruct.csqrt(v, result);
+  }
+
   public static sum<
     L extends number,
     A extends Uint32Array | Int32Array | Float32Array,
@@ -112,6 +128,14 @@ export class ScalarBatch<
     return GPUStruct.csum(l, r, result);
   }
 
+  public static sum3<
+    L extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>,
+  >(l: ScalarBatch<L, A, T>, m: ScalarBatch<L, A, T>, r: ScalarBatch<L, A, T>, result: ScalarBatch<L, A, T> = new ScalarBatch(l._device, l._array, l.layers)): ScalarBatch<L, A, T> {
+    return GPUStruct.csum3(l, m, r, result);
+  }
+
   public static sub<
     L extends number,
     A extends Uint32Array | Int32Array | Float32Array,
@@ -136,6 +160,14 @@ export class ScalarBatch<
     return GPUStruct.cdiv(l, r, result);
   }
 
+  public static pow<
+    L extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>,
+  >(l: ScalarBatch<L, A, T>, r: ScalarBatch<L, A, T>, result: ScalarBatch<L, A, T> = new ScalarBatch(l._device, l._array, l.layers)): ScalarBatch<L, A, T> {
+    return GPUStruct.cpow(l, r, result);
+  }
+
   public static scale<
     L extends number,
     A extends Uint32Array | Int32Array | Float32Array,

src/data/gpu/scalar.ts

@@ -44,6 +44,20 @@ export class Scalar<
     return this.data[0] as T;
   }
 
+  public static sqr<
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>,
+  >(v: Scalar<A, T>, result: Scalar<A, T> = new Scalar(v._device, v._array)): Scalar<A, T> {
+    return GPUStruct.csqr(v, result);
+  }
+
+  public static sqrt<
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>,
+  >(v: Scalar<A, T>, result: Scalar<A, T> = new Scalar(v._device, v._array)): Scalar<A, T> {
+    return GPUStruct.csqrt(v, result);
+  }
+
   public static sum<
     A extends Uint32Array | Int32Array | Float32Array,
     T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>,
@@ -51,6 +65,13 @@ export class Scalar<
     return GPUStruct.csum(l, r, result);
   }
 
+  public static sum3<
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>,
+  >(l: Scalar<A, T>, m: Scalar<A, T>, r: Scalar<A, T>, result: Scalar<A, T> = new Scalar(l._device, l._array)): Scalar<A, T> {
+    return GPUStruct.csum3(l, m, r, result);
+  }
+
   public static sub<
     A extends Uint32Array | Int32Array | Float32Array,
     T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>,
@@ -72,6 +93,13 @@ export class Scalar<
     return GPUStruct.cdiv(l, r, result);
   }
 
+  public static pow<
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>,
+  >(l: Scalar<A, T>, r: Scalar<A, T>, result: Scalar<A, T> = new Scalar(l._device, l._array)): Scalar<A, T> {
+    return GPUStruct.cpow(l, r, result);
+  }
+
   public static unavg<
     L extends number,
     A extends Uint32Array | Int32Array | Float32Array,

src/data/gpu/vector-batch.ts

@@ -148,6 +148,15 @@ export class VectorBatch<
     return GPUStruct.csum(l, r, result);
   }
 
+  public static sum3<
+    L extends number,
+    D extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>,
+  >(l: VectorBatch<L, D, A, T>, m: VectorBatch<L, D, A, T>, r: VectorBatch<L, D, A, T>, result: VectorBatch<L, D, A, T> = new VectorBatch(l._device, l._array, l.layers, l.dims)): VectorBatch<L, D, A, T> {
+    return GPUStruct.csum3(l, m, r, result);
+  }
+
   public static sub<
     L extends number,
     D extends number,

src/data/gpu/vector.ts

@@ -114,6 +114,14 @@ export class Vector<
     return GPUStruct.csum(l, r, result);
   }
 
+  public static sum3<
+    D extends number,
+    A extends Uint32Array | Int32Array | Float32Array,
+    T extends NonNullable<ReturnType<A['at']>> = NonNullable<ReturnType<A['at']>>,
+  >(l: Vector<D, A, T>, m: Vector<D, A, T>, r: Vector<D, A, T>, result: Vector<D, A, T> = new Vector(l._device, l._array, l.dims)): Vector<D, A, T> {
+    return GPUStruct.csum3(l, m, r, result);
+  }
+
   public static sub<
     D extends number,
     A extends Uint32Array | Int32Array | Float32Array,

src/model-gpu.ts

@@ -1,26 +1,28 @@
 import fs from 'node:fs';
 import { Writable } from 'node:stream';
+import { GPUStruct } from './data/gpu/gpu-struct';
 import { Matrix } from './data/gpu/matrix';
 import { Scalar } from './data/gpu/scalar';
 import { Vector } from './data/gpu/vector';
 import { Matrix as CPUMatrix, Vector as CPUVector, drm, mat, randseq, slope, stddev, vec } from './data/math';
 import { TID, Tokenizer } from './tokenizer';
 
-const SHADER_HOTONES = fs.readFileSync('shader/hotones.wgsl').toString();
+const SHADER_ONE_HOTS = fs.readFileSync('shader/onehots.wgsl').toString();
 const SHADER_EMBED = fs.readFileSync('shader/embed.wgsl').toString();
-const SHADER_ATTW = fs.readFileSync('shader/attw.wgsl').toString();
+const SHADER_CENTER = fs.readFileSync('shader/center.wgsl').toString();
+const SHADER_NORMALIZE = fs.readFileSync('shader/normalize.wgsl').toString();
+const SHADER_TWEIGH = fs.readFileSync('shader/tweigh.wgsl').toString();
 const SHADER_SOFTMAX_EXP = fs.readFileSync('shader/softmaxexp.wgsl').toString();
 const SHADER_SOFTMAX_NORM = fs.readFileSync('shader/softmaxnorm.wgsl').toString();
 const SHADER_DLDO = fs.readFileSync('shader/dldo.wgsl').toString();
-const SHADER_OTDLDO = fs.readFileSync('shader/otdldo.wgsl').toString();
-const SHADER_TOTDLDOWA = fs.readFileSync('shader/totdldowa.wgsl').toString();
-const SHADER_DLDS = fs.readFileSync('shader/dlds.wgsl').toString();
-const SHADER_GRAD_TOUT = fs.readFileSync('shader/gradtout.wgsl').toString();
-const SHADER_GRAD_TATTW = fs.readFileSync('shader/gradtattw.wgsl').toString();
-const SHADER_GRAD_INP = fs.readFileSync('shader/gradinp.wgsl').toString();
+const SHADER_SOFTMAX_GRAD = fs.readFileSync('shader/softmaxgrad.wgsl').toString();
+const SHADER_LAYERNORM_GRAD = fs.readFileSync('shader/layernormgrad.wgsl').toString();
+const SHADER_CLIP = fs.readFileSync('shader/clip.wgsl').toString();
 const SHADER_ARGMAX = fs.readFileSync('shader/argmax.wgsl').toString();
 const SHADER_SAMPLES = fs.readFileSync('shader/samples.wgsl').toString();
 
+const LOGGING = false;
+
 export interface ModelInitData<
   V extends number,
   W extends number,
@@ -39,7 +41,7 @@ export interface ModelData<
 > extends ModelInitData<V, W, D> {
   stage: 'early' | 'mid' | 'end';
   tokenizer: string;
-  samples: number;
+  step: number;
   weights: {
     inp: CPUMatrix<V, D>;
     pos: CPUMatrix<W, D>;
@@ -48,6 +50,22 @@ export interface ModelData<
     tval: CPUMatrix<D, D>;
     out: CPUMatrix<V, D>;
   };
+  momentums: {
+    inp: CPUMatrix<V, D>;
+    pos: CPUMatrix<W, D>;
+    tqry: CPUMatrix<D, D>;
+    tkey: CPUMatrix<D, D>;
+    tval: CPUMatrix<D, D>;
+    out: CPUMatrix<V, D>;
+  };
+  velocities: {
+    inp: CPUMatrix<V, D>;
+    pos: CPUMatrix<W, D>;
+    tqry: CPUMatrix<D, D>;
+    tkey: CPUMatrix<D, D>;
+    tval: CPUMatrix<D, D>;
+    out: CPUMatrix<V, D>;
+  };
 }
 
 interface TrainingPoint {
@@ -108,7 +126,7 @@ export class Model<
   public readonly ws: W;
   public readonly ds: D;
 
-  protected samples: number = 0;
+  protected step: number = 0;
   protected stage: 'early' | 'mid' | 'end' = 'early';
 
   protected readonly weights: {
@@ -120,20 +138,49 @@ export class Model<
     readonly tout: Matrix<V, D, Float32Array>;
   };
 
+  protected readonly momentums: {
+    readonly inp: Matrix<V, D, Float32Array>;
+    readonly pos: Matrix<W, D, Float32Array>;
+    readonly tqry: Matrix<D, D, Float32Array>;
+    readonly tkey: Matrix<D, D, Float32Array>;
+    readonly tval: Matrix<D, D, Float32Array>;
+    readonly tout: Matrix<V, D, Float32Array>;
+  };
+
+  protected readonly velocities: {
+    readonly inp: Matrix<V, D, Float32Array>;
+    readonly pos: Matrix<W, D, Float32Array>;
+    readonly tqry: Matrix<D, D, Float32Array>;
+    readonly tkey: Matrix<D, D, Float32Array>;
+    readonly tval: Matrix<D, D, Float32Array>;
+    readonly tout: Matrix<V, D, Float32Array>;
+  };
+
   protected readonly processing: {
+    readonly one: Scalar<Float32Array>;
     readonly ds: Scalar<Float32Array>;
+    readonly vs: Scalar<Float32Array>;
     readonly lr: Scalar<Float32Array>;
+    readonly mbeta: Scalar<Float32Array>;
+    readonly vbeta: Scalar<Float32Array>;
+    readonly step: Scalar<Float32Array>;
+    readonly thresholdL: Scalar<Float32Array>;
+    readonly thresholdM: Scalar<Float32Array>;
+    readonly thresholdS: Scalar<Float32Array>;
     readonly itids: Vector<W, Int32Array, TID>;
     readonly ilogits: Matrix<W, V, Float32Array>;
     readonly e: Matrix<W, D, Float32Array>;
+    readonly ce: Matrix<W, D, Float32Array>;
+    readonly ne: Matrix<W, D, Float32Array>;
+    readonly tne: Matrix<D, W, Float32Array>;
     readonly q: Matrix<W, D, Float32Array>;
     readonly k: Matrix<W, D, Float32Array>;
-    readonly tk: Matrix<D, W, Float32Array>;
     readonly qtk: Matrix<W, W, Float32Array>;
     readonly v: Matrix<W, D, Float32Array>;
     readonly s: Matrix<W, W, Float32Array>;
     readonly se: Matrix<W, W, Float32Array>;
     readonly wa: Matrix<W, W, Float32Array>;
+    readonly ar: Matrix<W, D, Float32Array>;
     readonly a: Matrix<W, D, Float32Array>;
     readonly o: Matrix<W, V, Float32Array>;
     readonly temp: Scalar<Float32Array>;
@@ -141,21 +188,39 @@ export class Model<
     readonly on: Matrix<W, V, Float32Array>;
     readonly ologits: Matrix<W, V, Float32Array>;
     readonly dldo: Matrix<W, V, Float32Array>;
+    readonly dlda: Matrix<W, D, Float32Array>;
     readonly otdldo: Matrix<D, W, Float32Array>;
-    readonly totdldowa: Matrix<W, D, Float32Array>;
+    readonly dldv: Matrix<W, D, Float32Array>;
     readonly dlds: Matrix<W, W, Float32Array>;
     readonly dldsk: Matrix<W, D, Float32Array>;
     readonly dldsq: Matrix<W, D, Float32Array>;
     readonly ginp: Matrix<V, D, Float32Array>;
+    readonly ginpnormh: Vector<V, Float32Array>;
+    readonly ginpnorm: Scalar<Float32Array>;
     readonly gtqry: Matrix<D, D, Float32Array>;
+    readonly gtqrynormh: Vector<D, Float32Array>;
+    readonly gtqrynorm: Scalar<Float32Array>;
     readonly gtkey: Matrix<D, D, Float32Array>;
+    readonly gtkeynormh: Vector<D, Float32Array>;
+    readonly gtkeynorm: Scalar<Float32Array>;
     readonly gtval: Matrix<D, D, Float32Array>;
+    readonly gtvalnormh: Vector<D, Float32Array>;
+    readonly gtvalnorm: Scalar<Float32Array>;
     readonly gtout: Matrix<V, D, Float32Array>;
+    readonly gtoutnormh: Vector<V, Float32Array>;
+    readonly gtoutnorm: Scalar<Float32Array>;
     readonly rnd: Vector<W, Float32Array>;
     readonly otids: Vector<W, Int32Array, TID>;
     readonly ttids: Vector<W, Int32Array, TID>;
   };
 
+  protected async log<A extends Uint32Array | Int32Array | Float32Array>(name: string, s: GPUStruct<A>, sample: number = 4): Promise<void> {
+    if (LOGGING) {
+      const v = await s.get() as number[];
+      console.log(name, Array.isArray(v) ? v.slice(0, sample) : v);
+    }
+  }
+
   protected softmax<H extends number, W extends number>(
     m: Matrix<H, W, Float32Array>,
     e: Matrix<H, W, Float32Array> = new Matrix(this.device, Float32Array, m.height, m.width),
@@ -166,56 +231,111 @@ export class Model<
     return r;
   }
 
+  protected softmax_grad<H extends number, W extends number>(
+    m: Matrix<H, W, Float32Array>,
+    dldm: Matrix<H, W, Float32Array>,
+    result: Matrix<H, W, Float32Array> = new Matrix(this.device, Float32Array, m.height, m.width),
+  ): Matrix<H, W, Float32Array> {
+    result.calculate('softmaxgrad', SHADER_SOFTMAX_GRAD, [m, dldm]);
+    return result;
+  }
+
+  protected async layernorm<H extends number, W extends number>(
+    m: Matrix<H, W, Float32Array>,
+    c: Matrix<H, W, Float32Array> = new Matrix(this.device, Float32Array, m.height, m.width),
+    r: Matrix<H, W, Float32Array> = new Matrix(this.device, Float32Array, m.height, m.width),
+  ): Promise<Matrix<H, W, Float32Array>> {
+    c.calculate('center', SHADER_CENTER, [m], false, 4);
+    r.calculate('normalize', SHADER_NORMALIZE, [m], false, 4);
+    await this.log('layernorm c', c);
+    await this.log('layernorm r', r);
+    return r;
+  }
+
+  protected layernorm_grad<H extends number, W extends number>(
+    m: Matrix<H, W, Float32Array>,
+    dldm: Matrix<H, W, Float32Array>,
+    result: Matrix<H, W, Float32Array> = new Matrix(this.device, Float32Array, m.height, m.width),
+  ): Matrix<H, W, Float32Array> {
+    result.calculate('layernormgrad', SHADER_LAYERNORM_GRAD, [m, dldm]);
+    return result;
+  }
+
   protected async ilogits(itids: Vector<W, Int32Array, TID> = this.processing.itids): Promise<Matrix<W, V, Float32Array>> {
-    this.processing.ilogits.calculate('hotones', SHADER_HOTONES, [itids]);
+    this.processing.ilogits.calculate('onehots', SHADER_ONE_HOTS, [itids]);
+    await this.log('itids', itids);
+    await this.log('ilogits', this.processing.ilogits);
     return this.processing.ilogits;
   }
 
   protected async e(ilogits: Matrix<W, V, Float32Array> = this.processing.ilogits): Promise<Matrix<W, D, Float32Array>> {
     this.processing.e.calculate('embed', SHADER_EMBED, [ilogits, this.weights.inp, this.weights.pos]);
+    await this.log('inp', this.weights.inp);
+    await this.log('pos', this.weights.pos);
+    await this.log('e', this.processing.e);
     return this.processing.e;
   }
 
-  protected q(e: Matrix<W, D, Float32Array> = this.processing.e): Matrix<W, D, Float32Array> {
-    this.processing.q.calculate('attw', SHADER_ATTW, [e, this.weights.tqry]);
+  protected async ne(e: Matrix<W, D, Float32Array> = this.processing.e): Promise<Matrix<W, D, Float32Array>> {
+    this.layernorm(e, this.processing.ce, this.processing.ne);
+    await this.log('ne', this.processing.ne);
+    return this.processing.ne;
+  }
+
+  protected async q(ne: Matrix<W, D, Float32Array> = this.processing.ne): Promise<Matrix<W, D, Float32Array>> {
+    this.processing.q.calculate('tqry', SHADER_TWEIGH, [ne, this.weights.tqry]);
+    await this.log('tqry', this.weights.tqry);
+    await this.log('q', this.processing.q);
     return this.processing.q;
   }
 
-  protected k(e: Matrix<W, D, Float32Array> = this.processing.e): Matrix<W, D, Float32Array> {
-    this.processing.k.calculate('attw', SHADER_ATTW, [e, this.weights.tkey]);
+  protected async k(ne: Matrix<W, D, Float32Array> = this.processing.ne): Promise<Matrix<W, D, Float32Array>> {
+    this.processing.k.calculate('tkey', SHADER_TWEIGH, [ne, this.weights.tkey]);
+    await this.log('tkey', this.weights.tkey);
+    await this.log('k', this.processing.k);
     return this.processing.k;
   }
 
-  protected v(e: Matrix<W, D, Float32Array> = this.processing.e): Matrix<W, D, Float32Array> {
-    this.processing.v.calculate('attw', SHADER_ATTW, [e, this.weights.tval]);
+  protected async v(ne: Matrix<W, D, Float32Array> = this.processing.ne): Promise<Matrix<W, D, Float32Array>> {
+    this.processing.v.calculate('tval', SHADER_TWEIGH, [ne, this.weights.tval]);
+    await this.log('tval', this.weights.tval);
+    await this.log('v', this.processing.v);
     return this.processing.v;
   }
 
-  protected s(
+  protected async s(
     q: Matrix<W, D, Float32Array> = this.processing.q,
     k: Matrix<W, D, Float32Array> = this.processing.k,
-  ): Matrix<W, W, Float32Array> {
-    Matrix.transpose(k, this.processing.tk);
-    Matrix.mul(q, this.processing.tk, this.processing.qtk);
+  ): Promise<Matrix<W, W, Float32Array>> {
+    Matrix.transrmul(q, k, this.processing.qtk);
     Matrix.scale(this.processing.ds, this.processing.qtk, this.processing.s);
+    await this.log('s', this.processing.s);
     return this.processing.s;
   }
 
   protected async wa(s: Matrix<W, W, Float32Array> = this.processing.s): Promise<Matrix<W, W, Float32Array>> {
     this.softmax(s, this.processing.se, this.processing.wa);
+    await this.log('se', this.processing.se);
+    await this.log('wa', this.processing.wa);
     return this.processing.wa;
   }
 
-  protected a(
+  protected async a(
+    e: Matrix<W, D, Float32Array> = this.processing.e,
     wa: Matrix<W, W, Float32Array> = this.processing.wa,
     v: Matrix<W, D, Float32Array> = this.processing.v,
-  ): Matrix<W, D, Float32Array> {
-    Matrix.mul(wa, v, this.processing.a);
+  ): Promise<Matrix<W, D, Float32Array>> {
+    Matrix.mul(wa, v, this.processing.ar);
+    Matrix.sum(e, this.processing.ar, this.processing.a);
+    await this.log('ar', this.processing.ar);
+    await this.log('a', this.processing.a);
     return this.processing.a;
   }
 
-  protected o(a: Matrix<W, D, Float32Array> = this.processing.a): Matrix<W, V, Float32Array> {
-    this.processing.o.calculate('attw', SHADER_ATTW, [a, this.weights.tout]);
+  protected async o(a: Matrix<W, D, Float32Array> = this.processing.a): Promise<Matrix<W, V, Float32Array>> {
+    this.processing.o.calculate('tout', SHADER_TWEIGH, [a, this.weights.tout]);
+    await this.log('tout', this.weights.tout);
+    await this.log('o', this.processing.o);
     return this.processing.o;
   }
 
@@ -250,9 +370,41 @@ export class Model<
     return probs.reduce((l: number, ps: CPUVector<V>, i: number) => l - Math.log(ps[target[i]] + 1e-10), 0) / probs.length;
   }
 
+  protected async adm<H extends number, W extends number>(
+    step: number,
+    m: Matrix<H, W, Float32Array>,
+    v: Matrix<H, W, Float32Array>,
+    g: Matrix<H, W, Float32Array>,
+  ): Promise<Matrix<H, W, Float32Array>> {
+    await this.processing.step.set(step);
+
+    Matrix.sum(
+      Matrix.scale(this.processing.mbeta, m),
+      Matrix.scale(Scalar.sub(this.processing.one, this.processing.mbeta), g),
+      m,
+    );
+
+    Matrix.sum(
+      Matrix.scale(this.processing.vbeta, v),
+      Matrix.scale(Scalar.sub(this.processing.one, this.processing.vbeta), Matrix.esqr(g)),
+      v,
+    );
+
+    const mhat = Matrix.scale(Scalar.div(this.processing.one, Scalar.sub(this.processing.one, Scalar.pow(this.processing.mbeta, this.processing.step))), m);
+    const vhat = Matrix.scale(Scalar.div(this.processing.one, Scalar.sub(this.processing.one, Scalar.pow(this.processing.vbeta, this.processing.step))), v);
+
+    Matrix.ediv(
+      mhat,
+      Matrix.sum(Matrix.esqrt(vhat), new Matrix(this.device, Float32Array, vhat.height, vhat.width, mat(vhat.height, vhat.width, () => 1e-8))),
+      g,
+    );
+
+    return g;
+  }
+
   protected async upd(
     ilogits: Matrix<W, V, Float32Array> = this.processing.ilogits,
-    e: Matrix<W, D, Float32Array> = this.processing.e,
+    ne: Matrix<W, D, Float32Array> = this.processing.ne,
     q: Matrix<W, D, Float32Array> = this.processing.q,
     k: Matrix<W, D, Float32Array> = this.processing.k,
     v: Matrix<W, D, Float32Array> = this.processing.v,
@@ -267,42 +419,60 @@ export class Model<
     this.processing.dldo.calculate('dldo', SHADER_DLDO, [
       ologits,
       ttids,
-      this.processing.lr,
+      this.processing.one,
     ]);
 
-    this.processing.otdldo.calculate('otdldo', SHADER_OTDLDO, [this.processing.dldo, this.weights.tout]);
+    Matrix.mul(this.processing.dldo, this.weights.tout, this.processing.dlda);
 
-    this.processing.dlds.calculate('dlds', SHADER_DLDS, [
-      this.processing.ds,
+    this.softmax_grad(
       wa,
-      v,
-      this.processing.otdldo,
-    ]);
+      Matrix.transrmul(this.processing.dlda, v),
+      this.processing.dlds,
+    );
+    Matrix.scale(this.processing.ds, this.processing.dlds, this.processing.dlds);
 
+    Matrix.translmul(this.processing.dlds, q, this.processing.dldsq);
     Matrix.mul(this.processing.dlds, k, this.processing.dldsk);
-    Matrix.mul(Matrix.transpose(this.processing.dlds), q, this.processing.dldsq);
-
-    this.processing.ginp.calculate('gradinp', SHADER_GRAD_INP, [
-      ilogits,
-      this.processing.dldsq,
-      this.processing.dldsk,
-      this.processing.totdldowa,
-      this.weights.tqry,
-      this.weights.tkey,
-      this.weights.tval,
-    ]);
-    Matrix.sum(this.processing.ginp, this.weights.inp, this.weights.inp);
+    Matrix.translmul(wa, this.processing.dlda, this.processing.dldv);
+    Matrix.mul(
+      Matrix.transpose(ilogits),
+      Matrix.sum(
+        Matrix.scale(new Scalar(this.device, Float32Array, 0.5), this.processing.dlda),
+        this.layernorm_grad(
+          ne,
+          Matrix.sum3(
+            Matrix.mul(this.processing.dldsk, this.weights.tqry),
+            Matrix.mul(this.processing.dldsq, this.weights.tkey),
+            Matrix.mul(this.processing.dldv, this.weights.tval),
+          ),
+        ),
+      ),
+      this.processing.ginp,
+    );
+    await this.adm(this.step, this.momentums.inp, this.velocities.inp, this.processing.ginp);
+    Matrix.scale(this.processing.lr, this.processing.ginp, this.processing.ginp);
+
+    Matrix.translmul(this.processing.dldsk, ne, this.processing.gtqry);
+    await this.adm(this.step, this.momentums.tqry, this.velocities.tqry, this.processing.gtqry);
+    Matrix.scale(this.processing.lr, this.processing.gtqry, this.processing.gtqry);
+
+    Matrix.translmul(this.processing.dldsq, ne, this.processing.gtkey);
+    await this.adm(this.step, this.momentums.tkey, this.velocities.tkey, this.processing.gtkey);
+    Matrix.scale(this.processing.lr, this.processing.gtkey, this.processing.gtkey);
+
+    Matrix.translmul(this.processing.dldv, ne, this.processing.gtval);
+    await this.adm(this.step, this.momentums.tval, this.velocities.tval, this.processing.gtval);
+    Matrix.scale(this.processing.lr, this.processing.gtval, this.processing.gtval);
+
+    Matrix.translmul(this.processing.dldo, a, this.processing.gtout);
+    await this.adm(this.step, this.momentums.tout, this.velocities.tout, this.processing.gtout);
+    Matrix.scale(this.processing.lr, this.processing.gtout, this.processing.gtout);
 
-    this.processing.gtout.calculate('gradtout', SHADER_GRAD_TOUT, [a, this.processing.dldo]);
-    Matrix.sum(this.processing.gtout, this.weights.tout, this.weights.tout);
-    this.processing.gtqry.calculate('gradtqry', SHADER_GRAD_TATTW, [e, this.processing.dldsk]);
+    Matrix.sum(this.processing.ginp, this.weights.inp, this.weights.inp);
     Matrix.sum(this.processing.gtqry, this.weights.tqry, this.weights.tqry);
-    this.processing.gtkey.calculate('gradtkey', SHADER_GRAD_TATTW, [e, this.processing.dldsq]);
     Matrix.sum(this.processing.gtkey, this.weights.tkey, this.weights.tkey);
-
-    this.processing.totdldowa.calculate('totdldowa', SHADER_TOTDLDOWA, [wa, this.processing.otdldo]);
-    this.processing.gtval.calculate('gradtval', SHADER_GRAD_TATTW, [e, this.processing.totdldowa]);
     Matrix.sum(this.processing.gtval, this.weights.tval, this.weights.tval);
+    Matrix.sum(this.processing.gtout, this.weights.tout, this.weights.tout);
   }
 
   public constructor(
@@ -325,7 +495,23 @@ export class Model<
         tqry: new Matrix(this.device, Float32Array, mat(this.ds, this.ds, () => (Math.random() - 0.5) * Math.sqrt(1 / this.ds))),
         tkey: new Matrix(this.device, Float32Array, mat(this.ds, this.ds, () => (Math.random() - 0.5) * Math.sqrt(1 / this.ds))),
         tval: new Matrix(this.device, Float32Array, mat(this.ds, this.ds, () => (Math.random() - 0.5) * Math.sqrt(1 / this.ds))),
-        tout: new Matrix(this.device, Float32Array, mat(this.vs, this.ds, () => (Math.random() - 0.5) * Math.sqrt(1 / this.ds))),
+        tout: new Matrix(this.device, Float32Array, mat(this.vs, this.ds, () => (Math.random() - 0.5) * Math.sqrt(1 / this.ds / this.vs))),
+      };
+      this.momentums = {
+        inp: new Matrix(this.device, Float32Array, mat(this.vs, this.ds, () => 0)),
+        pos: new Matrix(this.device, Float32Array, mat(this.ws, this.ds, () => 0)),
+        tqry: new Matrix(this.device, Float32Array, mat(this.ds, this.ds, () => 0)),
+        tkey: new Matrix(this.device, Float32Array, mat(this.ds, this.ds, () => 0)),
+        tval: new Matrix(this.device, Float32Array, mat(this.ds, this.ds, () => 0)),
+        tout: new Matrix(this.device, Float32Array, mat(this.vs, this.ds, () => 0)),
+      };
+      this.velocities = {
+        inp: new Matrix(this.device, Float32Array, mat(this.vs, this.ds, () => 0)),
+        pos: new Matrix(this.device, Float32Array, mat(this.ws, this.ds, () => 0)),
+        tqry: new Matrix(this.device, Float32Array, mat(this.ds, this.ds, () => 0)),
+        tkey: new Matrix(this.device, Float32Array, mat(this.ds, this.ds, () => 0)),
+        tval: new Matrix(this.device, Float32Array, mat(this.ds, this.ds, () => 0)),
+        tout: new Matrix(this.device, Float32Array, mat(this.vs, this.ds, () => 0)),
       };
     } else {
       this.weights = {
@@ -336,44 +522,83 @@ export class Model<
         tval: new Matrix(this.device, Float32Array, data.weights.tval),
         tout: new Matrix(this.device, Float32Array, data.weights.out),
       };
+      this.momentums = {
+        inp: new Matrix(this.device, Float32Array, data.momentums.inp),
+        pos: new Matrix(this.device, Float32Array, data.momentums.pos),
+        tqry: new Matrix(this.device, Float32Array, data.momentums.tqry),
+        tkey: new Matrix(this.device, Float32Array, data.momentums.tkey),
+        tval: new Matrix(this.device, Float32Array, data.momentums.tval),
+        tout: new Matrix(this.device, Float32Array, data.momentums.out),
+      };
+      this.velocities = {
+        inp: new Matrix(this.device, Float32Array, data.velocities.inp),
+        pos: new Matrix(this.device, Float32Array, data.velocities.pos),
+        tqry: new Matrix(this.device, Float32Array, data.velocities.tqry),
+        tkey: new Matrix(this.device, Float32Array, data.velocities.tkey),
+        tval: new Matrix(this.device, Float32Array, data.velocities.tval),
+        tout: new Matrix(this.device, Float32Array, data.velocities.out),
+      };
       this.stage = data.stage;
-      this.samples = data.samples;
+      this.step = data.step;
       if (data.tokenizer !== `${tokenizer.tokens.length}${tokenizer.tokens[0]}`) {
         throw new Error(`The model was trained against a different tokenizer ("${data.tokenizer}").`);
       }
     }
 
     this.processing = {
+      one: new Scalar<Float32Array>(this.device, Float32Array, 1),
       ds: new Scalar<Float32Array>(this.device, Float32Array, 1 / Math.sqrt(this.ds)),
+      vs: new Scalar<Float32Array>(this.device, Float32Array, 1 / Math.sqrt(this.vs)),
       lr: new Scalar<Float32Array>(this.device, Float32Array, 0),
+      mbeta: new Scalar<Float32Array>(this.device, Float32Array, 0.9),
+      vbeta: new Scalar<Float32Array>(this.device, Float32Array, 0.999),
+      step: new Scalar<Float32Array>(this.device, Float32Array, 0),
+
+      thresholdL: new Scalar<Float32Array>(this.device, Float32Array, 1e10),
+      thresholdM: new Scalar<Float32Array>(this.device, Float32Array, 1e10),
+      thresholdS: new Scalar<Float32Array>(this.device, Float32Array, 1e10),
       itids: new Vector<W, Int32Array, TID>(this.device, Int32Array, this.ws),
       ilogits: new Matrix<W, V, Float32Array>(this.device, Float32Array, this.ws, this.vs),
       e: new Matrix<W, D, Float32Array>(this.device, Float32Array, this.ws, this.ds),
+      ce: new Matrix<W, D, Float32Array>(this.device, Float32Array, this.ws, this.ds),
+      ne: new Matrix<W, D, Float32Array>(this.device, Float32Array, this.ws, this.ds),
+      tne: new Matrix<D, W, Float32Array>(this.device, Float32Array, this.ds, this.ws),
       q: new Matrix<W, D, Float32Array>(this.device, Float32Array, this.ws, this.ds),
       k: new Matrix<W, D, Float32Array>(this.device, Float32Array, this.ws, this.ds),
-      tk: new Matrix<D, W, Float32Array>(this.device, Float32Array, this.ds, this.ws),
       qtk: new Matrix<W, W, Float32Array>(this.device, Float32Array, this.ws, this.ws),
       v: new Matrix<W, D, Float32Array>(this.device, Float32Array, this.ws, this.ds),
       s: new Matrix<W, W, Float32Array>(this.device, Float32Array, this.ws, this.ws),
       se: new Matrix<W, W, Float32Array>(this.device, Float32Array, this.ws, this.ws),
       wa: new Matrix<W, W, Float32Array>(this.device, Float32Array, this.ws, this.ws),
       a: new Matrix<W, D, Float32Array>(this.device, Float32Array, this.ws, this.ds),
+      ar: new Matrix<W, D, Float32Array>(this.device, Float32Array, this.ws, this.ds),
       o: new Matrix<W, V, Float32Array>(this.device, Float32Array, this.ws, this.vs),
       temp: new Scalar<Float32Array>(this.device, Float32Array, 1),
       tempo: new Matrix<W, V, Float32Array>(this.device, Float32Array, this.ws, this.vs),
       on: new Matrix<W, V, Float32Array>(this.device, Float32Array, this.ws, this.vs),
       ologits: new Matrix<W, V, Float32Array>(this.device, Float32Array, this.ws, this.vs),
       dldo: new Matrix<W, V, Float32Array>(this.device, Float32Array, this.ws, this.vs),
+      dlda: new Matrix<W, D, Float32Array>(this.device, Float32Array, this.ws, this.ds),
       otdldo: new Matrix<D, W, Float32Array>(this.device, Float32Array, this.ds, this.ws),
-      totdldowa: new Matrix<W, D, Float32Array>(this.device, Float32Array, this.ws, this.ds),
+      dldv: new Matrix<W, D, Float32Array>(this.device, Float32Array, this.ws, this.ds),
       dlds: new Matrix<W, W, Float32Array>(this.device, Float32Array, this.ws, this.ws),
       dldsk: new Matrix<W, D, Float32Array>(this.device, Float32Array, this.ws, this.ds),
       dldsq: new Matrix<W, D, Float32Array>(this.device, Float32Array, this.ws, this.ds),
       ginp: new Matrix(this.device, Float32Array, this.vs, this.ds),
+      ginpnormh: new Vector(this.device, Float32Array, this.vs),
+      ginpnorm: new Scalar(this.device, Float32Array),
       gtqry: new Matrix(this.device, Float32Array, this.ds, this.ds),
+      gtqrynormh: new Vector(this.device, Float32Array, this.ds),
+      gtqrynorm: new Scalar(this.device, Float32Array),
       gtkey: new Matrix(this.device, Float32Array, this.ds, this.ds),
+      gtkeynormh: new Vector(this.device, Float32Array, this.ds),
+      gtkeynorm: new Scalar(this.device, Float32Array),
       gtval: new Matrix(this.device, Float32Array, this.ds, this.ds),
+      gtvalnormh: new Vector(this.device, Float32Array, this.ds),
+      gtvalnorm: new Scalar(this.device, Float32Array),
       gtout: new Matrix(this.device, Float32Array, this.vs, this.ds),
+      gtoutnormh: new Vector(this.device, Float32Array, this.vs),
+      gtoutnorm: new Scalar(this.device, Float32Array),
       rnd: new Vector<W, Float32Array>(this.device, Float32Array, this.ws),
       otids: new Vector<W, Int32Array, TID>(this.device, Int32Array, this.ws),
       ttids: new Vector<W, Int32Array, TID>(this.device, Int32Array, this.ws),
@@ -458,7 +683,9 @@ export class Model<
     const tids = this.tokenizer.tokenize(data);
     const batches = tids.length - this.ws - 1;
 
-    let lr = 0.003;
+    const seed = tids.slice(tids.length / 2, tids.length / 2 + 100);
+
+    const lr = 0.003;
 
     for (let epoch = 0; epoch < epochs; epoch++) {
       const idxs = randseq(tids.length);
@@ -471,15 +698,16 @@ export class Model<
 
         const ilogits = await this.ilogits(this.processing.itids);
         const e = await this.e(ilogits);
-        const q = this.q(e);
-        const k = this.k(e);
-        const v = this.v(e);
-        const s = this.s(q, k);
+        const ne = await this.ne(e);
+        const q = await this.q(ne);
+        const k = await this.k(ne);
+        const v = await this.v(ne);
+        const s = await this.s(q, k);
         const wa = await this.wa(s);
-        const a = this.a(wa, v);
-        const o = this.o(a);
+        const a = await this.a(e, wa, v);
+        const o = await this.o(a);
         const ologits = await this.ologits(o);
-        await this.upd(ilogits, e, q, k, v, wa, a, ologits, this.processing.ttids, lr);
+        await this.upd(ilogits, ne, q, k, v, wa, a, ologits, this.processing.ttids, lr);
 
         const queued: number = Date.now();
 
@@ -487,7 +715,7 @@ export class Model<
 
         const end: number = Date.now();
 
-        this.samples += 1;
+        this.step += 1;
         times.push(end - begin);
 
         if (
@@ -496,18 +724,18 @@ export class Model<
           || end - logged > statsInterval
         ) {
           const point: TrainingPoint = {
-            index: this.samples,
+            index: this.step,
             timestamp: Date.now(),
             epoch,
             batch,
             lr,
             loss: 0,
             accuracy: 0,
-            ginp: (await this.processing.ginp.get()).reduce((s, r) => r.reduce((s, v) => s + v * v, s), 0) / lr / lr,
-            gqry: (await this.processing.gtqry.get()).reduce((s, r) => r.reduce((s, v) => s + v * v, s), 0) / lr / lr,
-            gkey: (await this.processing.gtkey.get()).reduce((s, r) => r.reduce((s, v) => s + v * v, s), 0) / lr / lr,
-            gval: (await this.processing.gtval.get()).reduce((s, r) => r.reduce((s, v) => s + v * v, s), 0) / lr / lr,
-            gout: (await this.processing.gtout.get()).reduce((s, r) => r.reduce((s, v) => s + v * v, s), 0) / lr / lr,
+            ginp: await this.processing.ginpnorm.get(),
+            gqry: await this.processing.gtqrynorm.get(),
+            gkey: await this.processing.gtkeynorm.get(),
+            gval: await this.processing.gtvalnorm.get(),
+            gout: await this.processing.gtoutnorm.get(),
             queueing: 0,
             computing: 0,
             inferring: 0,
@@ -528,27 +756,27 @@ export class Model<
 
           losses.push(point.loss);
 
-          if (this.stage === 'early') {
-            lr = Math.max(0.001, Math.min(0.003, point.loss * 0.001));
-
-            if (losses.length > aw && (slope(losses.slice(-aw)) > -0.005 || stddev(losses.slice(-aw)) < 0.2)) {
-              logs?.write('Switching to mid stage.\n');
-              this.stage = 'mid';
-            }
-          } else if (this.stage === 'mid') {
-            lr = Math.max(0.0003, lr * 0.98);
-
-            if (losses.length > aw && Math.abs(slope(losses.slice(-aw))) < 0.002 && stddev(losses.slice(-aw)) < 0.05) {
-              logs?.write('Switching to end stage.\n');
-              this.stage = 'end';
-            }
-          } else {
-            lr = Math.max(1e-4, lr * 0.98);
-          }
+          // if (this.stage === 'early') {
+          //   lr = Math.max(0.001, Math.min(0.003, point.loss * 0.001));
+
+          //   if (losses.length > aw && (slope(losses.slice(-aw)) > -0.005 || stddev(losses.slice(-aw)) < 0.2)) {
+          //     logs?.write('Switching to mid stage.\n');
+          //     this.stage = 'mid';
+          //   }
+          // } else if (this.stage === 'mid') {
+          //   lr = Math.max(0.0003, lr * 0.98);
+
+          //   if (losses.length > aw && Math.abs(slope(losses.slice(-aw))) < 0.002 && stddev(losses.slice(-aw)) < 0.05) {
+          //     logs?.write('Switching to end stage.\n');
+          //     this.stage = 'end';
+          //   }
+          // } else {
+          //   lr = Math.max(1e-4, lr * 0.98);
+          // }
 
           logs?.write(`Done ${epoch * batches + batch + 1} / ${epochs * batches}, ETA ${Math.round(drm(times) * (epochs * batches - times.length) / 1000)}s\n`);
           stats?.write(`${Object.values(point).join(',')}\n`);
-          await this.debug(this.samples, this.processing.itids, this.processing.ttids, o, ologits, logits);
+          await this.debug(this.step, this.processing.itids, this.processing.ttids, o, ologits, logits);
 
           logged = end;
         }
@@ -567,19 +795,19 @@ export class Model<
 
           samples?.write(`Epoch ${epoch + 1}:\n`);
           samples?.write('Sample (temp=0.3):\n');
-          samples?.write(this.tokenizer.detokenize(await this.run(this.tokenizer.tokenize('Hello?'), { temp: 0.3 })).join(''));
+          samples?.write(this.tokenizer.detokenize(await this.run(seed, { temp: 0.3 })).join(''));
           samples?.write('\nEnd of sample.\n');
           samples?.write('Sample (temp=0.7):\n');
-          samples?.write(this.tokenizer.detokenize(await this.run(this.tokenizer.tokenize('Hello?'), { temp: 0.7 })).join(''));
+          samples?.write(this.tokenizer.detokenize(await this.run(seed, { temp: 0.7 })).join(''));
           samples?.write('\nEnd of sample.\n');
           samples?.write('Sample (temp=1):\n');
-          samples?.write(this.tokenizer.detokenize(await this.run(this.tokenizer.tokenize('Hello?'), { temp: 1 })).join(''));
+          samples?.write(this.tokenizer.detokenize(await this.run(seed, { temp: 1 })).join(''));
           samples?.write('\nEnd of sample.\n');
           samples?.write('Sample (temp=1.3):\n');
-          samples?.write(this.tokenizer.detokenize(await this.run(this.tokenizer.tokenize('Hello?'), { temp: 1.3 })).join(''));
+          samples?.write(this.tokenizer.detokenize(await this.run(seed, { temp: 1.3 })).join(''));
           samples?.write('\nEnd of sample.\n\n');
           samples?.write('Sample (argmax):\n');
-          samples?.write(this.tokenizer.detokenize(await this.run(this.tokenizer.tokenize('Hello?'), { temp: 'max' })).join(''));
+          samples?.write(this.tokenizer.detokenize(await this.run(seed, { temp: 'max' })).join(''));
           samples?.write('\nEnd of sample.\n\n');
 
           snapshotted = end;
@@ -606,13 +834,14 @@ export class Model<
     await this.processing.ttids.set([...itids.slice(1), 0]);
     const ilogits = await this.ilogits(this.processing.itids);
     const e = await this.e(ilogits);
-    const q = this.q(e);
-    const k = this.k(e);
-    const v = this.v(e);
-    const s = this.s(q, k);
+    const ne = await this.ne(e);
+    const q = await this.q(ne);
+    const k = await this.k(ne);
+    const v = await this.v(ne);
+    const s = await this.s(q, k);
     const wa = await this.wa(s);
-    const a = this.a(wa, v);
-    const o = this.o(a);
+    const a = await this.a(e, wa, v);
+    const o = await this.o(a);
     const ologits = await this.ologits(o, temp === 'max' ? 1 : temp);
     const out = temp === 'max' ? this.otids(ologits) : this.otids(await this.rnd(), ologits);
     const otids = await out.get();
@@ -648,7 +877,7 @@ export class Model<
       ws: this.ws,
       ds: this.ds,
       stage: this.stage,
-      samples: this.samples,
+      step: this.step,
       weights: {
         inp: await this.weights.inp.get(),
         pos: await this.weights.pos.get(),
@@ -657,6 +886,22 @@ export class Model<
         tval: await this.weights.tval.get(),
         out: await this.weights.tout.get(),
       },
+      momentums: {
+        inp: await this.momentums.inp.get(),
+        pos: await this.momentums.pos.get(),
+        tqry: await this.momentums.tqry.get(),
+        tkey: await this.momentums.tkey.get(),
+        tval: await this.momentums.tval.get(),
+        out: await this.momentums.tout.get(),
+      },
+      velocities: {
+        inp: await this.velocities.inp.get(),
+        pos: await this.velocities.pos.get(),
+        tqry: await this.velocities.tqry.get(),
+        tkey: await this.velocities.tkey.get(),
+        tval: await this.velocities.tval.get(),
+        out: await this.velocities.tout.get(),
+      },
     };
   }
 }