-
Notifications
You must be signed in to change notification settings - Fork 44
Expand file tree
/
Copy pathtest_codebook_theory.cpp
More file actions
163 lines (132 loc) · 5.38 KB
/
Copy pathtest_codebook_theory.cpp
File metadata and controls
163 lines (132 loc) · 5.38 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
/**
* test_codebook_theory.cpp -- Lloyd-Max codebook verification
*
* Verifies that the hardcoded codebook centroids match published
* optimal Lloyd-Max values for N(0,1), checks symmetry, and
* measures actual MSE against theoretical optimum.
*/
#include <gtest/gtest.h>
#include <cmath>
#include <vector>
#include <random>
#include <cstring>
#include <algorithm>
extern "C" {
#include "turboquant/turboquant.h"
/* Codebook API */
const float* tq_codebook_centroids(int bits);
int tq_codebook_levels(int bits);
void tq_codebook_quantize(const float* src, uint8_t* dst_indices,
int n, int bits, float inv_std);
void tq_codebook_dequantize(const uint8_t* indices, float* dst,
int n, int bits, float inv_std);
}
/* ============================================================
* Test 1: 2-bit centroids match N(0,1) Lloyd-Max optimal
* Literature values: [-1.5104, -0.4528, 0.4528, 1.5104]
* ============================================================ */
TEST(CodebookTheory, TwoBitCentroidsMatchLiterature) {
const float expected[4] = {-1.5104f, -0.4528f, 0.4528f, 1.5104f};
const float* actual = tq_codebook_centroids(2);
ASSERT_NE(actual, nullptr);
EXPECT_EQ(tq_codebook_levels(2), 4);
for (int i = 0; i < 4; i++) {
EXPECT_NEAR(actual[i], expected[i], 0.001f)
<< "2-bit centroid[" << i << "] mismatch";
}
}
/* ============================================================
* Test 2: 3-bit centroids match N(0,1) Lloyd-Max optimal
* ============================================================ */
TEST(CodebookTheory, ThreeBitCentroidsMatchLiterature) {
const float expected[8] = {
-2.1520f, -1.3440f, -0.7560f, -0.2451f,
0.2451f, 0.7560f, 1.3440f, 2.1520f
};
const float* actual = tq_codebook_centroids(3);
ASSERT_NE(actual, nullptr);
EXPECT_EQ(tq_codebook_levels(3), 8);
for (int i = 0; i < 8; i++) {
EXPECT_NEAR(actual[i], expected[i], 0.001f)
<< "3-bit centroid[" << i << "] mismatch";
}
}
/* ============================================================
* Test 3: Codebook symmetry: centroid[i] = -centroid[n-1-i]
* ============================================================ */
TEST(CodebookTheory, SymmetryProperty) {
for (int bits = 1; bits <= 4; bits++) {
const float* c = tq_codebook_centroids(bits);
int n = tq_codebook_levels(bits);
ASSERT_NE(c, nullptr);
ASSERT_GT(n, 0);
for (int i = 0; i < n / 2; i++) {
EXPECT_NEAR(c[i], -c[n - 1 - i], 1e-5f)
<< "Symmetry violated for " << bits << "-bit codebook at index " << i;
}
}
}
/* ============================================================
* Test 4: Actual MSE of 2-bit codebook on N(0,1) samples
* Theoretical optimal MSE for 2-bit on N(0,1): ~0.1175
* Verify actual MSE is within 1.20x of theory.
* ============================================================ */
TEST(CodebookTheory, TwoBitMSEWithinTheoretical) {
const int N = 100000;
const float theoretical_mse = 0.1175f;
std::mt19937 rng(42);
std::normal_distribution<float> dist(0.0f, 1.0f);
std::vector<float> samples(N);
for (int i = 0; i < N; i++) {
samples[i] = dist(rng);
}
/* Quantize with inv_std = 1.0 (already N(0,1)) */
std::vector<uint8_t> indices(N);
tq_codebook_quantize(samples.data(), indices.data(), N, 2, 1.0f);
/* Dequantize */
std::vector<float> reconstructed(N);
tq_codebook_dequantize(indices.data(), reconstructed.data(), N, 2, 1.0f);
/* Compute MSE */
double mse = 0.0;
for (int i = 0; i < N; i++) {
double d = (double)samples[i] - (double)reconstructed[i];
mse += d * d;
}
mse /= N;
printf(" 2-bit codebook MSE: %.6f (theoretical: %.4f, ratio: %.3fx)\n",
mse, theoretical_mse, mse / theoretical_mse);
/* MSE should be within 1.20x of theoretical optimal */
EXPECT_LT(mse, theoretical_mse * 1.20)
<< "2-bit MSE " << mse << " exceeds 1.20x theoretical " << theoretical_mse;
/* MSE should not be lower than theoretical (sanity check) */
EXPECT_GT(mse, theoretical_mse * 0.90)
<< "2-bit MSE " << mse << " is suspiciously lower than theoretical";
}
/* ============================================================
* Test 5: 3-bit MSE verification
* Theoretical optimal MSE for 3-bit on N(0,1): ~0.0344
* ============================================================ */
TEST(CodebookTheory, ThreeBitMSEWithinTheoretical) {
const int N = 100000;
const float theoretical_mse = 0.0344f;
std::mt19937 rng(123);
std::normal_distribution<float> dist(0.0f, 1.0f);
std::vector<float> samples(N);
for (int i = 0; i < N; i++) {
samples[i] = dist(rng);
}
std::vector<uint8_t> indices(N);
tq_codebook_quantize(samples.data(), indices.data(), N, 3, 1.0f);
std::vector<float> reconstructed(N);
tq_codebook_dequantize(indices.data(), reconstructed.data(), N, 3, 1.0f);
double mse = 0.0;
for (int i = 0; i < N; i++) {
double d = (double)samples[i] - (double)reconstructed[i];
mse += d * d;
}
mse /= N;
printf(" 3-bit codebook MSE: %.6f (theoretical: %.4f, ratio: %.3fx)\n",
mse, theoretical_mse, mse / theoretical_mse);
EXPECT_LT(mse, theoretical_mse * 1.20)
<< "3-bit MSE " << mse << " exceeds 1.20x theoretical " << theoretical_mse;
}