Switch to Adam optimizer:

WIP

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/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/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/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;
+}

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,7 +155,7 @@ 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;
@@ -190,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() } },
       ],
@@ -249,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 COMPWISE_SHADER, l: S, r: S, result: S): S {
+  >(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 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<
@@ -298,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,7 +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_TMUL = fs.readFileSync('shader/tmul.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();
 
@@ -174,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,
@@ -254,7 +268,26 @@ export class MatrixBatch<
     return result.calculate('transpose', SHADER_TRANS, [m]);
   }
 
-  public static tmul<
+  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,
@@ -270,7 +303,47 @@ export class MatrixBatch<
     this.assertColocated(l, tr, result);
     this.assertPackable(1, l, tr, result);
 
-    return result.calculate('tmul', SHADER_TMUL, [l, tr]);
+    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<

src/data/gpu/matrix.ts

@@ -8,7 +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_TMUL = fs.readFileSync('shader/tmul.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<
@@ -134,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,
@@ -143,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,
@@ -171,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,
@@ -198,24 +236,6 @@ export class Matrix<
     return result.calculate('mul', SHADER_MUL, [l, r]);
   }
 
-  public static tmul<
-    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('tmul', SHADER_TMUL, [l, tr]);
-  }
-
   public static transpose<
     H extends number,
     W extends number,
@@ -232,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

@@ -17,6 +17,7 @@ const SHADER_SOFTMAX_NORM = fs.readFileSync('shader/softmaxnorm.wgsl').toString(
 const SHADER_DLDO = fs.readFileSync('shader/dldo.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();
 
@@ -40,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>;
@@ -49,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 {
@@ -109,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: {
@@ -121,14 +138,41 @@ 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 qtk: Matrix<W, W, Float32Array>;
@@ -144,16 +188,27 @@ 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>;
@@ -161,7 +216,8 @@ export class Model<
 
   protected async log<A extends Uint32Array | Int32Array | Float32Array>(name: string, s: GPUStruct<A>, sample: number = 4): Promise<void> {
     if (LOGGING) {
-      console.log(name, (await s.get() as number[]).slice(0, sample));
+      const v = await s.get() as number[];
+      console.log(name, Array.isArray(v) ? v.slice(0, sample) : v);
     }
   }
 
@@ -247,32 +303,39 @@ export class Model<
     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> {
+  ): Promise<Matrix<W, W, Float32Array>> {
-    Matrix.tmul(q, k, this.processing.qtk);
+    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> {
+  ): 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> {
+  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;
   }
 
@@ -307,6 +370,38 @@ 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,
     ne: Matrix<W, D, Float32Array> = this.processing.ne,
@@ -324,77 +419,54 @@ export class Model<
     this.processing.dldo.calculate('dldo', SHADER_DLDO, [
       ologits,
       ttids,
-      this.processing.lr,
+      this.processing.one,
     ]);
 
+    Matrix.mul(this.processing.dldo, this.weights.tout, this.processing.dlda);
+
     this.softmax_grad(
       wa,
-      Matrix.mul(this.processing.dldo, Matrix.mul(this.weights.tout, Matrix.transpose(v))),
+      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.translmul(wa, this.processing.dlda, this.processing.dldv);
     Matrix.mul(
       Matrix.transpose(ilogits),
       Matrix.sum(
-        Matrix.mul(this.processing.dldo, this.weights.tout),
+        Matrix.scale(new Scalar(this.device, Float32Array, 0.5), this.processing.dlda),
         this.layernorm_grad(
           ne,
-          Matrix.sum(
+          Matrix.sum3(
-            Matrix.mul(
+            Matrix.mul(this.processing.dldsk, this.weights.tqry),
-              Matrix.mul(
+            Matrix.mul(this.processing.dldsq, this.weights.tkey),
-                this.processing.dlds,
+            Matrix.mul(this.processing.dldv, this.weights.tval),
-                Matrix.scale(this.processing.ds, k),
-              ),
-              this.weights.tqry,
-            )
-            ,
-            Matrix.sum(
-              Matrix.mul(
-                Matrix.mul(
-                  Matrix.transpose(this.processing.dlds),
-                  Matrix.scale(this.processing.ds, q),
-                ),
-                this.weights.tkey,
-              ),
-              Matrix.mul(
-                Matrix.mul(Matrix.transpose(wa), this.processing.dldo),
-                Matrix.mul(this.weights.tout, 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.transpose(
+    Matrix.translmul(this.processing.dldsk, ne, this.processing.gtqry);
-      Matrix.scale(this.processing.ds, Matrix.mul(
+    await this.adm(this.step, this.momentums.tqry, this.velocities.tqry, this.processing.gtqry);
-        Matrix.transpose(ne),
+    Matrix.scale(this.processing.lr, this.processing.gtqry, this.processing.gtqry);
-        Matrix.mul(this.processing.dlds, k),
-      )),
-      this.processing.gtqry,
-    );
 
-    Matrix.transpose(
+    Matrix.translmul(this.processing.dldsq, ne, this.processing.gtkey);
-      Matrix.scale(this.processing.ds, Matrix.mul(
+    await this.adm(this.step, this.momentums.tkey, this.velocities.tkey, this.processing.gtkey);
-        Matrix.transpose(ne),
+    Matrix.scale(this.processing.lr, this.processing.gtkey, this.processing.gtkey);
-        Matrix.mul(Matrix.transpose(this.processing.dlds), q),
-      )),
-      this.processing.gtkey,
-    );
 
-    Matrix.transpose(
+    Matrix.translmul(this.processing.dldv, ne, this.processing.gtval);
-      Matrix.mul(
+    await this.adm(this.step, this.momentums.tval, this.velocities.tval, this.processing.gtval);
-        Matrix.mul(Matrix.transpose(ne), Matrix.transpose(wa)),
+    Matrix.scale(this.processing.lr, this.processing.gtval, this.processing.gtval);
-        Matrix.mul(this.processing.dldo, this.weights.tout),
-      ),
-      this.processing.gtval,
-    );
 
-    Matrix.transpose(
+    Matrix.translmul(this.processing.dldo, a, this.processing.gtout);
-      Matrix.mul(Matrix.transpose(a), this.processing.dldo),
+    await this.adm(this.step, this.momentums.tout, this.velocities.tout, this.processing.gtout);
-      this.processing.gtout,
+    Matrix.scale(this.processing.lr, this.processing.gtout, this.processing.gtout);
-    );
 
     Matrix.sum(this.processing.ginp, this.weights.inp, this.weights.inp);
     Matrix.sum(this.processing.gtqry, this.weights.tqry, this.weights.tqry);
@@ -423,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 = {
@@ -434,21 +522,47 @@ 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),
       qtk: new Matrix<W, W, Float32Array>(this.device, Float32Array, this.ws, this.ws),
@@ -464,16 +578,27 @@ export class Model<
       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),
@@ -560,7 +685,7 @@ export class Model<
 
     const seed = tids.slice(tids.length / 2, tids.length / 2 + 100);
 
-    let lr = 0.003;
+    const lr = 0.003;
 
     for (let epoch = 0; epoch < epochs; epoch++) {
       const idxs = randseq(tids.length);
@@ -577,10 +702,10 @@ export class Model<
         const q = await this.q(ne);
         const k = await this.k(ne);
         const v = await this.v(ne);
-        const s = this.s(q, k);
+        const s = await this.s(q, k);
         const wa = await this.wa(s);
-        const a = this.a(e, wa, v);
+        const a = await this.a(e, wa, v);
-        const o = this.o(a);
+        const o = await this.o(a);
         const ologits = await this.ologits(o);
         await this.upd(ilogits, ne, q, k, v, wa, a, ologits, this.processing.ttids, lr);
 
@@ -590,7 +715,7 @@ export class Model<
 
         const end: number = Date.now();
 
-        this.samples += 1;
+        this.step += 1;
         times.push(end - begin);
 
         if (
@@ -599,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,
+            ginp: await this.processing.ginpnorm.get(),
-            gqry: (await this.processing.gtqry.get()).reduce((s, r) => r.reduce((s, v) => s + v * v, s), 0) / lr / lr,
+            gqry: await this.processing.gtqrynorm.get(),
-            gkey: (await this.processing.gtkey.get()).reduce((s, r) => r.reduce((s, v) => s + v * v, s), 0) / lr / lr,
+            gkey: await this.processing.gtkeynorm.get(),
-            gval: (await this.processing.gtval.get()).reduce((s, r) => r.reduce((s, v) => s + v * v, s), 0) / lr / lr,
+            gval: await this.processing.gtvalnorm.get(),
-            gout: (await this.processing.gtout.get()).reduce((s, r) => r.reduce((s, v) => s + v * v, s), 0) / lr / lr,
+            gout: await this.processing.gtoutnorm.get(),
             queueing: 0,
             computing: 0,
             inferring: 0,
@@ -631,27 +756,27 @@ export class Model<
 
           losses.push(point.loss);
 
-          if (this.stage === 'early') {
+          // if (this.stage === 'early') {
-            lr = Math.max(0.001, Math.min(0.003, point.loss * 0.001));
+          //   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)) {
+          //   if (losses.length > aw && (slope(losses.slice(-aw)) > -0.005 || stddev(losses.slice(-aw)) < 0.2)) {
-              logs?.write('Switching to mid stage.\n');
+          //     logs?.write('Switching to mid stage.\n');
-              this.stage = 'mid';
+          //     this.stage = 'mid';
-            }
+          //   }
-          } else if (this.stage === 'mid') {
+          // } else if (this.stage === 'mid') {
-            lr = Math.max(0.0003, lr * 0.98);
+          //   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) {
+          //   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');
+          //     logs?.write('Switching to end stage.\n');
-              this.stage = 'end';
+          //     this.stage = 'end';
-            }
+          //   }
-          } else {
+          // } else {
-            lr = Math.max(1e-4, lr * 0.98);
+          //   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;
         }
@@ -713,10 +838,10 @@ export class Model<
     const q = await this.q(ne);
     const k = await this.k(ne);
     const v = await this.v(ne);
-    const s = this.s(q, k);
+    const s = await this.s(q, k);
     const wa = await this.wa(s);
-    const a = this.a(e, wa, v);
+    const a = await this.a(e, wa, v);
-    const o = this.o(a);
+    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();
@@ -752,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(),
@@ -761,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(),
+      },
     };
   }
 }