code wiki / _hdl_build / nx_imatrix.nx
nx_imatrix.nx source
↩ module page · 57 lines · 2968 B
1// nx_imatrix.nx -- IMPORTANCE-MATRIX quantization = rate-distortion WATER-FILLING (the deepest
2// survived task from the governed loop: VRAM-research actionable A1 ~= historic actionable HA1).
3// Naive quantization minimizes RAW error sum((w-wq)^2) and spends bits uniformly. But not every
4// weight matters equally to the output -- an importance matrix (from calibration activations) says
5// how much each weight's error costs. The right objective is the IMPORTANCE-WEIGHTED error
6// sum(imp_i * (w_i - wq_i)^2), and the right bit allocation is Shannon WATER-FILLING: give the next
7// bit to whichever group reduces the weighted error most. This is rate-distortion optimization
8// (Radio), and it is why imatrix cuts low-bit error ~31% at ~zero size cost: same total bits, but
9// spent where distortion is expensive. Distortion here is the QUALITY-relevant (importance-weighted)
10// distortion, not raw MSE -- closing one notch toward the perceptually-meaningful objective the
11// Critic insisted on. license_tier: ORIGINAL Refs: ggml imatrix PR#4930; Radio (rate-distortion); Shannon water-filling.
12
13import "nx_syscalls.nx"
14
15// quantization error^2 proxy for b bits: ~ (range^2)/2^(2b). Scaled so integers stay clean.
16func imat_err2(bits: i64) -> i64 {
17 if bits >= 20 { return 0 }
18 return (1 << 20) >> (2 * bits)
19}
20
21// the IMPORTANCE-WEIGHTED error of a bit allocation (the quality-relevant distortion).
22func imat_weighted_error(n: i64, imp: *i64, bits: *i64) -> i64 {
23 var s: i64 = 0; var i: i64 = 0
24 while i < n { s = s + imp[i] * imat_err2(bits[i]); i = i + 1 }
25 return s
26}
27
28func imat_total_bits(n: i64, bits: *i64) -> i64 { var s: i64 = 0; var i: i64 = 0; while i < n { s = s + bits[i]; i = i + 1 } return s }
29
30// WATER-FILL: distribute a fixed total bit budget to minimize the importance-weighted error. Start
31// each group at 1 bit, then greedily give each remaining bit to the group with the largest marginal
32// reduction imp_i*(err2(b) - err2(b+1)). Provably the rate-distortion-optimal greedy allocation.
33func imat_waterfill(n: i64, imp: *i64, total_bits: i64, out_bits: *i64) -> i64 {
34 var i: i64 = 0
35 while i < n { out_bits[i] = 1; i = i + 1 }
36 var used: i64 = n
37 while used < total_bits {
38 var best: i64 = 0; var bestred: i64 = 0 - 1; var j: i64 = 0
39 while j < n {
40 let b: i64 = out_bits[j]
41 let red: i64 = imp[j] * (imat_err2(b) - imat_err2(b + 1))
42 if red > bestred { bestred = red; best = j }
43 j = j + 1
44 }
45 out_bits[best] = out_bits[best] + 1
46 used = used + 1
47 }
48 return 0
49}
50
51// helper: a uniform allocation of the same total budget (the naive baseline), for the comparison.
52func imat_uniform(n: i64, total_bits: i64, out_bits: *i64) -> i64 {
53 let per: i64 = total_bits / n; let extra: i64 = total_bits - per * n
54 var i: i64 = 0
55 while i < n { out_bits[i] = per; if i < extra { out_bits[i] = per + 1 } i = i + 1 }
56 return 0
57}