MSE distortion per vector <= sqrt(3) * pi/2 * (1/4^b) for unit vectors.
(bits)
| 55 | |
| 56 | @pytest.mark.parametrize("bits", [1, 2, 3, 4]) |
| 57 | def test_distortion_within_paper_bound(bits): |
| 58 | """MSE distortion per vector <= sqrt(3) * pi/2 * (1/4^b) for unit vectors.""" |
| 59 | d = 128 |
| 60 | cb = LloydMaxCodebook(d, bits) |
| 61 | sigma = 1.0 / math.sqrt(d) |
| 62 | |
| 63 | n_samples = 5000 |
| 64 | x = torch.randn(n_samples, d) |
| 65 | x = x / torch.norm(x, dim=-1, keepdim=True) |
| 66 | |
| 67 | rotated = x # with identity rotation, still approximately N(0, 1/d) |
| 68 | indices = cb.quantize(rotated) |
| 69 | reconstructed = cb.dequantize(indices) |
| 70 | mse = ((rotated - reconstructed) ** 2).sum(dim=-1).mean().item() |
| 71 | |
| 72 | upper_bound = math.sqrt(3) * math.pi / 2 * (1.0 / (4**bits)) |
| 73 | assert mse < upper_bound * 2.0, ( |
| 74 | f"bits={bits}: MSE={mse:.6f} exceeds 2× paper bound {upper_bound:.6f}" |
| 75 | ) |
| 76 | |
| 77 | |
| 78 | def test_roundtrip_identity(): |
nothing calls this directly
no test coverage detected