nx_vq_train.nx source
↩ module page · 43 lines · 1953 B
1// nx_vq_train.nx -- Lloyd's algorithm (k-means) VQ codebook training: moonshot rung #2. Given a training set of LPC
2// reflection vectors, it iteratively (1) assigns each vector to its nearest codebook entry and (2) moves each centroid
3// to the mean of its assigned vectors -- provably non-increasing distortion. This is the optimal CLASSICAL codebook
4// that neural RVQ generalises: it lowers the VQ distortion at the SAME bitrate (better coverage), the bridge from raw
5// VQ to the learned codebook. Composes nx_vq. Caller owns all scratch (no hidden alloc). license_tier: ORIGINAL
6import "nx_vq.nx"
7
8// pointer to the t-th D-vector inside a flat train[T*D] buffer
9func vqt_vec(train: *i64, t: i64, D: i64) -> *i64 { return (train as i64 + t*D*8) as *i64 }
10
11// one Lloyd iteration over T training vectors. assign[T], sum[K*D], cnt[K] are caller scratch.
12func vqt_iterate(train: *i64, T: i64, cb: *i64, K: i64, D: i64, assign: *i64, sum: *i64, cnt: *i64) -> i64 {
13 var t: i64 = 0
14 while t < T { assign[t] = vq_encode(vqt_vec(train, t, D), cb, K, D); t = t + 1 }
15 var i: i64 = 0
16 while i < K*D { sum[i] = 0; i = i + 1 }
17 i = 0; while i < K { cnt[i] = 0; i = i + 1 }
18 t = 0
19 while t < T {
20 let a: i64 = assign[t]
21 cnt[a] = cnt[a] + 1
22 var j: i64 = 0
23 while j < D { sum[a*D + j] = sum[a*D + j] + train[t*D + j]; j = j + 1 }
24 t = t + 1
25 }
26 var c: i64 = 0
27 while c < K {
28 if cnt[c] > 0 { var j: i64 = 0; while j < D { cb[c*D + j] = sum[c*D + j] / cnt[c]; j = j + 1 } }
29 c = c + 1
30 }
31 return 0
32}
33// total quantisation distortion of cb over the training set (sum of nearest-entry squared distances)
34func vqt_distortion(train: *i64, T: i64, cb: *i64, K: i64, D: i64) -> i64 {
35 var s: i64 = 0
36 var t: i64 = 0
37 while t < T {
38 let v: *i64 = vqt_vec(train, t, D)
39 s = s + vq_dist2(v, cb, vq_encode(v, cb, K, D), D)
40 t = t + 1
41 }
42 return s
43}