yo_vector/

rotate.rs

//! The random rotation every vector goes through before it is quantised
//! (`10` section 3).
//!
//! RaBitQ quantises a coordinate to its sign, and a sign carries information
//! only when the coordinates are all about the same size. Real embeddings are
//! not like that: a handful of dimensions hold most of the energy, and the sign
//! of the rest is close to a coin toss. A random rotation fixes it, because
//! rotating a vector by a random orthogonal matrix spreads its length evenly
//! over the coordinates while leaving every distance and every angle exactly
//! where it was. That is the whole reason the estimator's error bound holds.
//!
//! ```
//! use yo_vector::Rotation;
//!
//! let r = Rotation::new(8, 42);
//! let mut v = [1.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0];
//! r.apply(&mut v);
//!
//! // The length did not move, and neither did anything else about the vector.
//! let len: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
//! assert!((len - 1.0).abs() < 1e-5);
//! // The one coordinate that held everything now holds a share of it.
//! assert!(v.iter().all(|x| x.abs() < 0.95));
//! ```
//!
//! # Why this shape and not a matrix
//!
//! The obvious rotation is a dense `D` by `D` orthogonal matrix from a QR
//! decomposition, and it costs `D^2` multiplies per vector. At 768 dimensions
//! that is 590 thousand multiplies to insert one vector, and the ingest target
//! is fifty thousand vectors a second on one core, which asks for 29 GFLOP/s of
//! nothing but rotation. It does not fit, so the rotation has to be structured.
//!
//! The usual structured answer is a random sign flip followed by a Hadamard
//! transform, which is `D log D` and lovely, and needs `D` to be a power of two.
//! Padding 768 up to 1024 would make a one bit code 128 bytes instead of 96,
//! which is a third more index for every vector in the collection, so that is
//! not free either.
//!
//! What is here is the same idea without the power of two. A round flips the
//! sign of a random half of the coordinates, pairs them all up at random, and
//! replaces each pair with its sum and its difference over the square root of
//! two. Every step of that is orthogonal by construction rather than
//! approximately orthogonal after rounding, and a quarter turn on a pair splits
//! whatever it was holding evenly between the two, which is the part that
//! actually spreads a spike out. Pairing at random rather than by a fixed stride
//! is what lets any coordinate reach any other, so `log2(D)` rounds and a couple
//! more is enough.
//!
//! The angle is not random, and that is deliberate. A round with a random angle
//! keeps roughly nine tenths of what the larger side was holding, so a spike
//! only decays like `0.9^rounds` and it takes something like sixty rounds at 256
//! dimensions to flatten. A quarter turn halves it every time and takes eight.
//! The randomness the estimator needs comes from the pairing and the signs, and
//! there is plenty of it.
//!
//! # It is not written down anywhere
//!
//! A rotation is `dim` and a seed, and both live in the collection's catalogue
//! entry. Rebuilding it is deterministic on every machine and every target,
//! because the generator underneath is, so the file never holds the tables and
//! two processes that open the same collection compute the same rotation.

use yo_common::Rng;

/// A quarter turn on a pair is a sum and a difference, both over this.
const INV_ROOT2: f32 = core::f32::consts::FRAC_1_SQRT_2;

/// One sweep: a sign for every coordinate, a pairing of all of them, and which
/// way round each pair is turned.
///
/// The pairs are stored rather than the permutation they came from, because
/// rotating in place off a pair list is one pass over the vector and no scratch
/// buffer, where a permutation would want somewhere to write the shuffled copy.
struct Round {
    /// One bit per coordinate, set when its sign is flipped before the pairing.
    flip: Vec<u64>,
    /// `(i, j)` for each pair, which together cover every index once, or every
    /// index but one when the dimension is odd.
    pairs: Vec<(u32, u32)>,
    /// One bit per pair, set when the difference goes first rather than second.
    turn: Vec<u64>,
}

/// A random orthogonal transform, rebuilt from its seed rather than stored.
pub struct Rotation {
    dim: usize,
    seed: u64,
    rounds: Vec<Round>,
}

impl Rotation {
    /// The rotation a collection of `dim` dimensional vectors uses, from the
    /// seed in its catalogue entry.
    ///
    /// # Panics
    ///
    /// If `dim` is zero. A collection of vectors with no coordinates is a
    /// mistake made somewhere further up rather than a case to handle.
    #[must_use]
    pub fn new(dim: usize, seed: u64) -> Rotation {
        assert!(dim > 0, "a vector has at least one dimension");
        let mut rng = Rng::new(seed);
        let sweeps = sweeps(dim);
        let mut rounds = Vec::with_capacity(sweeps);
        // The identity, shuffled fresh for every round, which is what makes the
        // pairing different each time and so lets a coordinate reach across the
        // whole vector rather than staying in the half it started in.
        let mut order: Vec<u32> = (0..dim as u32).collect();
        for _ in 0..sweeps {
            shuffle(&mut order, &mut rng);
            let half = dim / 2;
            let pairs = (0..half)
                .map(|k| (order[2 * k], order[2 * k + 1]))
                .collect();
            rounds.push(Round {
                flip: bits(dim, &mut rng),
                pairs,
                turn: bits(half, &mut rng),
            });
        }
        Rotation { dim, seed, rounds }
    }

    /// How many coordinates a vector this rotates has.
    #[must_use]
    pub fn dim(&self) -> usize {
        self.dim
    }

    /// The seed this was built from, which is what a catalogue stores.
    #[must_use]
    pub fn seed(&self) -> u64 {
        self.seed
    }

    /// Rotate a vector where it lies.
    ///
    /// # Panics
    ///
    /// If `v` is not [`Rotation::dim`] long.
    pub fn apply(&self, v: &mut [f32]) {
        assert_eq!(
            v.len(),
            self.dim,
            "this rotation is for {} dimensions and was handed {}",
            self.dim,
            v.len()
        );
        for round in &self.rounds {
            // A sign is the top bit of the float, so flipping one is an xor and
            // there is no branch to mispredict on a coin toss.
            for (i, c) in v.iter_mut().enumerate() {
                let sign = ((round.flip[i / 64] >> (i % 64)) & 1) << 31;
                *c = f32::from_bits(c.to_bits() ^ (sign as u32));
            }
            for (k, &(i, j)) in round.pairs.iter().enumerate() {
                let (i, j) = (i as usize, j as usize);
                let (a, b) = (v[i], v[j]);
                let (sum, difference) = ((a + b) * INV_ROOT2, (a - b) * INV_ROOT2);
                if (round.turn[k / 64] >> (k % 64)) & 1 == 0 {
                    v[i] = sum;
                    v[j] = difference;
                } else {
                    v[i] = difference;
                    v[j] = sum;
                }
            }
        }
    }
}

/// How many sweeps it takes for a spike to be spread over the whole vector.
///
/// A round splits whatever a coordinate is holding between two, so the set one
/// can have reached after `k` of them is `2^k` wide. `log2(dim)` rounded up is
/// where that covers the vector, and the two on top are because the pairings are
/// drawn independently and so overlap near the end rather than tiling neatly.
fn sweeps(dim: usize) -> usize {
    ((usize::BITS - (dim - 1).leading_zeros()) as usize + 2).max(4)
}

/// Fisher and Yates, so that every pairing is equally likely.
fn shuffle(order: &mut [u32], rng: &mut Rng) {
    for i in (1..order.len()).rev() {
        order.swap(i, rng.below(i + 1));
    }
}

/// `n` coin tosses, packed.
fn bits(n: usize, rng: &mut Rng) -> Vec<u64> {
    (0..n.div_ceil(64)).map(|_| rng.next_u64()).collect()
}

#[cfg(test)]
mod tests {
    use super::*;

    fn dot(a: &[f32], b: &[f32]) -> f32 {
        a.iter().zip(b).map(|(x, y)| x * y).sum()
    }

    /// A number in `[0, 1)` from the generator, for the tests that want one.
    fn unit(rng: &mut Rng) -> f32 {
        (rng.next_u64() >> 40) as f32 / (1u32 << 24) as f32
    }

    /// A handful of vectors that are not all alike.
    fn sample(dim: usize, n: usize, seed: u64) -> Vec<Vec<f32>> {
        let mut rng = Rng::new(seed);
        (0..n)
            .map(|_| (0..dim).map(|_| unit(&mut rng) * 2.0 - 1.0).collect())
            .collect()
    }

    #[test]
    fn a_rotation_keeps_every_length_and_every_angle() {
        let r = Rotation::new(64, 7);
        let vs = sample(64, 8, 11);
        for v in &vs {
            let mut spun = v.clone();
            r.apply(&mut spun);
            let before = dot(v, v).sqrt();
            let after = dot(&spun, &spun).sqrt();
            assert!((before - after).abs() < 1e-4, "{before} became {after}");
        }
        // And the angle between any two of them, which is what a distance is
        // made of and so is the property that actually has to survive.
        for a in 0..vs.len() {
            for b in 0..a {
                let mut x = vs[a].clone();
                let mut y = vs[b].clone();
                let before = dot(&x, &y);
                r.apply(&mut x);
                r.apply(&mut y);
                let after = dot(&x, &y);
                assert!((before - after).abs() < 1e-3, "{before} became {after}");
            }
        }
    }

    #[test]
    fn the_same_seed_is_the_same_rotation() {
        let v = sample(32, 1, 3).pop().expect("one vector");
        let mut a = v.clone();
        let mut b = v;
        Rotation::new(32, 99).apply(&mut a);
        Rotation::new(32, 99).apply(&mut b);
        assert_eq!(a, b);

        let mut c = a.clone();
        Rotation::new(32, 100).apply(&mut c);
        assert_ne!(a, c, "two seeds should not be one rotation");
    }

    /// The point of the whole thing: a vector whose length sits in one
    /// coordinate comes out with it spread over all of them, which is what
    /// makes the sign of a coordinate worth a bit.
    #[test]
    fn a_spike_comes_out_flat() {
        for dim in [64usize, 256, 768] {
            let r = Rotation::new(dim, 5);
            let mut v = vec![0.0f32; dim];
            v[0] = 1.0;
            r.apply(&mut v);

            let even = 1.0 / (dim as f32).sqrt();
            let biggest = v.iter().fold(0.0f32, |m, x| m.max(x.abs()));
            assert!(biggest < even * 5.0, "{dim}: {biggest} against {even}");
            // And nothing is left sitting at zero either, which is the failure
            // a fixed pairing would have: half the vector never touched.
            let alive = v.iter().filter(|x| x.abs() > even / 4.0).count();
            assert!(alive > dim * 3 / 4, "{dim}: only {alive} coordinates moved");
        }
    }

    #[test]
    fn an_odd_dimension_still_rotates() {
        let r = Rotation::new(7, 1);
        let mut v = [1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0];
        let before = dot(&v, &v).sqrt();
        r.apply(&mut v);
        let after = dot(&v, &v).sqrt();
        assert!((before - after).abs() < 1e-3);
    }

    #[test]
    fn one_dimension_is_only_a_sign() {
        // There is nothing to pair it with, so all a round can do is flip it,
        // and the length still comes out where it went in.
        let r = Rotation::new(1, 1);
        let mut v = [3.0f32];
        r.apply(&mut v);
        assert_eq!(v[0].abs(), 3.0);
    }
}