Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 51 additions & 0 deletions crates/bgz17/src/base17.rs
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,25 @@ impl Base17 {
d
}

/// PCDVQ-informed L1: weight sign dimension 20x over mantissa.
///
/// From arxiv 2506.05432: direction (sign) is 20x more sensitive to
/// quantization than magnitude. BF16 decomposition maps to polar:
/// dim 0 = sign (direction), dims 1-6 = exponent (magnitude scale),
/// dims 7-16 = mantissa (fine detail).
///
/// Returns weighted L1 distance. Higher values for direction mismatches.
#[inline]
pub fn l1_weighted(&self, other: &Base17) -> u32 {
let mut d = 0u32;
for i in 0..BASE_DIM {
let diff = (self.dims[i] as i32 - other.dims[i] as i32).unsigned_abs();
let weight = if i == 0 { 20 } else if i < 7 { 3 } else { 1 };
d += diff * weight;
}
d
}

/// Sign-bit agreement (out of 17).
#[inline]
pub fn sign_agreement(&self, other: &Base17) -> u32 {
Expand Down Expand Up @@ -325,4 +344,36 @@ mod tests {
let b = Base17::from_bytes(&bytes);
assert_eq!(a, b);
}

#[test]
fn test_l1_weighted_sign_dim_dominates() {
// Difference only in dim 0 (sign) vs only in dim 10 (mantissa)
let a = Base17 { dims: [0; 17] };
let mut b_sign = Base17 { dims: [0; 17] };
b_sign.dims[0] = 100; // sign dimension diff = 100
let mut b_mant = Base17 { dims: [0; 17] };
b_mant.dims[10] = 100; // mantissa dimension diff = 100

let d_sign = a.l1_weighted(&b_sign);
let d_mant = a.l1_weighted(&b_mant);

// Sign diff should be 20× the mantissa diff
assert_eq!(d_sign, 100 * 20);
assert_eq!(d_mant, 100 * 1);
assert!(d_sign > d_mant * 10, "sign should dominate: {} vs {}", d_sign, d_mant);
}

#[test]
fn test_l1_weighted_self_zero() {
let a = Base17 { dims: [100, -50, 30, 0, 10, -20, 40, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10] };
assert_eq!(a.l1_weighted(&a), 0);
}

#[test]
fn test_l1_weighted_geq_l1() {
// Weighted L1 should be >= plain L1 (all weights >= 1)
let a = Base17 { dims: [100, -50, 30, 0, 10, -20, 40, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10] };
let b = Base17 { dims: [-50, 30, 0, 10, -20, 40, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 100] };
assert!(a.l1_weighted(&b) >= a.l1(&b));
}
}
23 changes: 23 additions & 0 deletions crates/bgz17/src/distance_matrix.rs
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,29 @@ impl DistanceMatrix {
self.data[a as usize * self.k + b as usize]
}

/// Build with PCDVQ direction-weighted L1 distances.
///
/// From arxiv 2506.05432: direction is 20x more sensitive to quantization.
/// Uses `Base17::l1_weighted()` instead of `Base17::l1()`.
pub fn build_pcdvq_weighted(palette: &Palette) -> Self {
let k = palette.len();
let mut data = vec![0u16; k * k];

for i in 0..k {
for j in (i + 1)..k {
let d = palette.entries[i].l1_weighted(&palette.entries[j]);
// Scale to u16. Max weighted L1 = 20 + 6*3 + 10*1 = 48 per dim,
// but using i16 range: max ≈ 20*65535 + ... ≈ 2.7M
let max_weighted = 20u64 * 65535 + 6 * 3 * 65535 + 10 * 65535;
let scaled = ((d as u64 * 65535) / max_weighted).min(65535) as u16;
data[i * k + j] = scaled;
data[j * k + i] = scaled;
}
}

DistanceMatrix { data, k }
}

/// Byte size of the matrix.
pub fn byte_size(&self) -> usize {
self.k * self.k * 2
Expand Down
1 change: 1 addition & 0 deletions crates/bgz17/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,7 @@ pub mod palette_semiring;
pub mod palette_matrix;
pub mod palette_csr;
pub mod simd;
pub mod rabitq_compat;

/// Maximum palette size per plane.
pub const MAX_PALETTE_SIZE: usize = 256;
Expand Down
99 changes: 99 additions & 0 deletions crates/bgz17/src/palette.rs
Original file line number Diff line number Diff line change
Expand Up @@ -160,6 +160,58 @@ impl Palette {
Palette { entries: centroids }
}

/// Sigma-band palette: codebook from empirical distribution.
///
/// Each band boundary = sorted-percentile of input patterns.
/// k bands = k entries. Guaranteed no empty clusters (each band covers
/// 100/k percent of data by construction).
///
/// Distribution-free: works for Gaussian, bimodal, skewed, heavy-tailed.
/// Inspired by GQ (arxiv 2512.06609): Target Divergence Constraint
/// → codebook without training, guaranteed uniform utilization.
pub fn from_sigma_bands(patterns: &[Base17], k: usize) -> Self {
let k = k.min(MAX_PALETTE_SIZE).min(patterns.len());
if k == 0 {
return Palette { entries: Vec::new() };
}

// Sort patterns by L1 distance from the centroid (global mean)
let n = patterns.len();
let mut mean = [0i64; 17];
for p in patterns {
for d in 0..17 {
mean[d] += p.dims[d] as i64;
}
}
let centroid = Base17 {
dims: {
let mut dims = [0i16; 17];
for d in 0..17 {
dims[d] = (mean[d] / n as i64) as i16;
}
dims
},
};

// Compute distances from centroid and sort indices by distance
let mut indexed: Vec<(usize, u32)> = patterns
.iter()
.enumerate()
.map(|(i, p)| (i, p.l1(&centroid)))
.collect();
indexed.sort_unstable_by_key(|&(_, d)| d);

// Pick one representative per equal-percentile band
let mut entries = Vec::with_capacity(k);
for band in 0..k {
let center_idx = (band * n / k) + (n / (2 * k));
let idx = center_idx.min(n - 1);
entries.push(patterns[indexed[idx].0].clone());
}

Palette { entries }
}

/// Build three palettes (one per S/P/O plane) from a set of SpoBase17 edges.
pub fn build_spo(edges: &[SpoBase17], k: usize, max_iter: usize) -> (Self, Self, Self) {
let s_patterns: Vec<Base17> = edges.iter().map(|e| e.subject.clone()).collect();
Expand Down Expand Up @@ -303,6 +355,53 @@ mod tests {
assert_eq!(PaletteResolution::Full256.matrix_bytes(), 131072);
}

#[test]
fn test_sigma_band_palette_size() {
let patterns = make_patterns(200);
let palette = Palette::from_sigma_bands(&patterns, 32);
assert_eq!(palette.len(), 32);
}

#[test]
fn test_sigma_band_no_empty() {
let patterns = make_patterns(100);
let palette = Palette::from_sigma_bands(&patterns, 16);
assert_eq!(palette.len(), 16);
// All entries should be distinct (from different percentile bands)
for i in 0..palette.len() {
for j in (i + 1)..palette.len() {
// Not necessarily distinct, but they come from different positions
// At minimum, palette shouldn't be empty
assert!(!palette.entries[i].dims.iter().all(|&d| d == 0) || i == 0);
}
}
}

#[test]
fn test_sigma_band_comparable_to_kmeans() {
let patterns = make_patterns(200);
let sigma = Palette::from_sigma_bands(&patterns, 32);
let kmeans = Palette::build(&patterns, 32, 10);

// Both should produce reasonable assignments
let total_dist_sigma: u64 = patterns.iter().map(|p| {
let idx = sigma.nearest(p);
p.l1(&sigma.entries[idx as usize]) as u64
}).sum();

let total_dist_kmeans: u64 = patterns.iter().map(|p| {
let idx = kmeans.nearest(p);
p.l1(&kmeans.entries[idx as usize]) as u64
}).sum();

// Sigma-band should be within 5× of k-means (it's training-free)
assert!(
total_dist_sigma < total_dist_kmeans * 5,
"sigma {} should be within 5× of kmeans {}",
total_dist_sigma, total_dist_kmeans
);
}

#[test]
fn test_convergence() {
// K-means should converge quickly
Expand Down
Loading