diff --git a/include/rabitqlib/quantization/rabitq_impl.hpp b/include/rabitqlib/quantization/rabitq_impl.hpp index a59a698..bd05ee3 100644 --- a/include/rabitqlib/quantization/rabitq_impl.hpp +++ b/include/rabitqlib/quantization/rabitq_impl.hpp @@ -434,10 +434,10 @@ inline T ex_bits_code( ipnorm_inv = quantize_ex(abs_res.data(), ex_code, dim, ex_bits); } - // revert codes for negative dims + // Revert codes for negative or zero dims int32_t mask = (1 << ex_bits) - 1; for (size_t j = 0; j < dim; ++j) { - if (res_arr.data()[j] < 0) { + if (res_arr.data()[j] <= 0) { TP tmp = ex_code[j]; ex_code[j] = (~tmp) & mask; } diff --git a/tests/unit/rabitqlib/quantization/rabitq_test.cpp b/tests/unit/rabitqlib/quantization/rabitq_test.cpp index 22231be..f8103a4 100644 --- a/tests/unit/rabitqlib/quantization/rabitq_test.cpp +++ b/tests/unit/rabitqlib/quantization/rabitq_test.cpp @@ -6,6 +6,7 @@ #include #include #include +#include namespace rabitqlib::quant { namespace { @@ -92,5 +93,80 @@ TEST(RabitqOneBitTest, FullQuantizationInitializesCodesAndFactors) { EXPECT_TRUE(std::isfinite(f_error)); } +// An exactly-zero residual coordinate must be encoded consistently by both halves of +// the split code. one_bit_code() uses (residual > 0) and so calls zero negative; if +// ex_bits_code() calls it positive, the assembled code lands 2^ex_bits away from where +// the factors assume it is. The == ||r||^2 invariant cannot see this, +// because a zero coordinate contributes nothing to that inner product. +TEST(RabitqSignConventionTest, ZeroResidualCoordinateGetsSmallestMagnitudeCode) { + constexpr size_t kDim = 64; + constexpr size_t kTotalBits = 8; + constexpr size_t kExBits = kTotalBits - 1; + + std::array data{}; + for (size_t i = 0; i < kDim; ++i) { + data[i] = static_cast(i % 7) - 3.0F; // zero at i % 7 == 3 + } + + std::array code{}; + float f_add = 0; + float f_rescale = 0; + float f_error = 0; + quantize_full_single( + data.data(), kDim, kTotalBits, code.data(), f_add, f_rescale, f_error + ); + + // centered code is code + cb, cb = -(2^total_bits - 1)/2 + const float cb = -((1 << kExBits) - 0.5F); + size_t n_zero = 0; + for (size_t i = 0; i < kDim; ++i) { + if (data[i] != 0.0F) { + continue; + } + ++n_zero; + EXPECT_FLOAT_EQ(static_cast(code[i]) + cb, -0.5F) + << "zero residual at dim " << i << " encoded as " << +code[i]; + } + ASSERT_GT(n_zero, 0U) << "test vector must contain exact zeros"; +} + +TEST(RabitqSignConventionTest, ZeroResidualDoesNotInflateReconstructionNorm) { + constexpr size_t kDim = 64; + constexpr size_t kTotalBits = 8; + constexpr size_t kExBits = kTotalBits - 1; + + std::array data{}; + for (size_t i = 0; i < kDim; ++i) { + data[i] = static_cast((i * 37) % 23) - 11.0F; // zeros at rem 11 + } + + std::array code{}; + float f_add = 0; + float f_rescale = 0; + float f_error = 0; + quantize_full_single( + data.data(), kDim, kTotalBits, code.data(), f_add, f_rescale, f_error + ); + + const double cb = -((1 << kExBits) - 0.5); + const double c = f_rescale / -2.0; + double nsq = 0; + double qq = 0; + size_t n_zero = 0; + for (size_t i = 0; i < kDim; ++i) { + const double q = static_cast(code[i]) + cb; + nsq += static_cast(data[i]) * data[i]; + qq += q * q; + n_zero += (data[i] == 0.0F); + } + ASSERT_GT(n_zero, 0U) << "test vector must contain exact zeros"; + + // ||c*xu_cb||^2 / ||r||^2 - 1 == tan^2 of the angle between the residual and its + // reconstruction. A mismatched sign convention inflates this by orders of magnitude. + const double tan_sq = ((c * c * qq) / nsq) - 1.0; + EXPECT_LT(tan_sq, 1e-3) << "reconstruction norm inflated, tan^2 = " + << std::to_string(tan_sq); +} + } // namespace } // namespace rabitqlib::quant