environment/training-source/source-snapshot/src/pca.rs

//! Turning 384-dimensional patch features into something a person can see.
//!
//! The first display mode is the classic DINO visualization: project patch
//! features onto their top three principal components and read those off as
//! RGB. It needs no trained decoder, costs one `[hidden, 3]` matmul, and is
//! semantically meaningful — patches of the same object land on the same
//! colour, which is precisely what the features encode.
//!
//! The basis is fitted **on device** from a captured frame rather than
//! shipped as an asset. Feature statistics depend on what the camera is
//! actually looking at, and a basis fitted on ImageNet-ish photos would
//! waste most of its dynamic range on a living room wall.
//!
//! Fitting uses power iteration with deflation. For three components out of
//! a few hundred tokens that is a handful of milliseconds on the CPU, and
//! it avoids pulling in a linear-algebra dependency for one job.

use meganeura::{Graph, NodeId};

/// Number of principal components, one per colour channel.
pub const COMPONENTS: usize = 3;

/// A fitted projection from feature space to RGB.
#[derive(Debug, Clone)]
pub struct Basis {
    /// Feature-space mean, subtracted before projection.
    pub mean: Vec<f32>,
    /// Row-major `[COMPONENTS, dim]` — the top principal directions.
    pub components: Vec<f32>,
    /// Per-component scale mapping projections into roughly `[0, 1]`.
    pub scale: [f32; COMPONENTS],
    /// Per-component offset, applied after scaling.
    pub offset: [f32; COMPONENTS],
    pub dim: usize,
}

impl Basis {
    /// An identity-ish basis to display before the first fit completes:
    /// three arbitrary orthogonal axes. Produces a picture, just not a
    /// well-conditioned one.
    pub fn placeholder(dim: usize) -> Self {
        let mut components = vec![0.0; COMPONENTS * dim];
        for c in 0..COMPONENTS {
            components[c * dim + c] = 1.0;
        }
        Self {
            mean: vec![0.0; dim],
            components,
            scale: [1.0; COMPONENTS],
            offset: [0.5; COMPONENTS],
            dim,
        }
    }

    /// Fit from a `[tokens, dim]` feature matrix.
    ///
    /// `skip` drops the leading prefix tokens: CLS and register tokens are
    /// not patches, they carry global rather than spatial information, and
    /// including them skews the components away from what is being
    /// displayed.
    pub fn fit(features: &[f32], tokens: usize, dim: usize, skip: usize) -> Self {
        assert_eq!(features.len(), tokens * dim, "feature matrix shape mismatch");
        assert!(skip < tokens, "nothing left after skipping {skip} tokens");
        let rows = tokens - skip;
        let data = &features[skip * dim..];

        // Centre.
        let mut mean = vec![0.0f32; dim];
        for r in 0..rows {
            for d in 0..dim {
                mean[d] += data[r * dim + d];
            }
        }
        for m in &mut mean {
            *m /= rows as f32;
        }
        let mut centred: Vec<f32> = (0..rows * dim)
            .map(|i| data[i] - mean[i % dim])
            .collect();

        let mut components = vec![0.0f32; COMPONENTS * dim];
        let mut projected = vec![0.0f32; COMPONENTS * rows];

        for c in 0..COMPONENTS {
            // Deterministic, non-degenerate start vector. A constant vector
            // would be orthogonal to components that sum to zero, so vary it.
            let mut v: Vec<f32> = (0..dim)
                .map(|d| ((d * 2654435761usize) % 1024) as f32 / 1024.0 - 0.5)
                .collect();
            normalize(&mut v);

            let mut scores = vec![0.0f32; rows];
            for _ in 0..48 {
                // scores = X v ; v' = Xᵀ scores  — power iteration on XᵀX
                // without ever forming the dim×dim covariance matrix.
                for r in 0..rows {
                    let row = &centred[r * dim..(r + 1) * dim];
                    scores[r] = row.iter().zip(&v).map(|(a, b)| a * b).sum();
                }
                let mut next = vec![0.0f32; dim];
                for r in 0..rows {
                    let s = scores[r];
                    let row = &centred[r * dim..(r + 1) * dim];
                    for d in 0..dim {
                        next[d] += s * row[d];
                    }
                }
                if normalize(&mut next) < 1e-12 {
                    break;
                }
                v = next;
            }

            // Final scores, then deflate so the next iteration finds the
            // next-strongest direction.
            for r in 0..rows {
                let row = &centred[r * dim..(r + 1) * dim];
                let s: f32 = row.iter().zip(&v).map(|(a, b)| a * b).sum();
                scores[r] = s;
                projected[c * rows + r] = s;
            }
            for r in 0..rows {
                let s = scores[r];
                for d in 0..dim {
                    centred[r * dim + d] -= s * v[d];
                }
            }
            components[c * dim..(c + 1) * dim].copy_from_slice(&v);
        }

        // Map each component to [0, 1] using a robust range rather than
        // min/max: a single outlier patch would otherwise flatten the whole
        // image into a narrow band of colour.
        let mut lo = [0.0f32; COMPONENTS];
        let mut span = [0.0f32; COMPONENTS];
        for c in 0..COMPONENTS {
            let mut col: Vec<f32> = projected[c * rows..(c + 1) * rows].to_vec();
            col.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
            lo[c] = col[(rows as f32 * 0.02) as usize];
            let hi = col[((rows as f32 * 0.98) as usize).min(rows - 1)];
            span[c] = hi - lo[c];
        }

        // A component carrying almost no variance relative to the leading
        // one carries no visual information, and normalizing it to full
        // range would amplify numerical noise into a psychedelic channel.
        // This is not a contrived case: it is what a blank wall looks like.
        // Map such a channel to flat mid-grey instead.
        let dominant = span.iter().copied().fold(0.0f32, f32::max);
        let floor = dominant * 1e-3;
        let mut scale = [0.0f32; COMPONENTS];
        let mut offset = [0.5f32; COMPONENTS];
        for c in 0..COMPONENTS {
            if span[c] > floor && span[c] > f32::MIN_POSITIVE {
                scale[c] = 1.0 / span[c];
                offset[c] = -lo[c] / span[c];
            }
        }

        Self {
            mean,
            components,
            scale,
            offset,
            dim,
        }
    }

    /// The `[dim, COMPONENTS]` matrix the graph's projection matmul wants,
    /// with the per-component scale folded in.
    pub fn weight_matrix(&self) -> Vec<f32> {
        let mut w = vec![0.0f32; self.dim * COMPONENTS];
        for c in 0..COMPONENTS {
            for d in 0..self.dim {
                w[d * COMPONENTS + c] = self.components[c * self.dim + d] * self.scale[c];
            }
        }
        w
    }

    /// The matching bias.
    ///
    /// `(x - mean) @ W * scale + offset` folds into `x @ (W * scale) + b`
    /// with `b = offset - (mean @ W) * scale`, saving a subtraction pass
    /// over every token.
    pub fn bias_vector(&self) -> Vec<f32> {
        let mut b = [0.0f32; COMPONENTS];
        for c in 0..COMPONENTS {
            let dot: f32 = (0..self.dim)
                .map(|d| self.mean[d] * self.components[c * self.dim + d])
                .sum();
            b[c] = self.offset[c] - dot * self.scale[c];
        }
        b.to_vec()
    }

    /// CPU-side projection, for tests and for previewing a fit without a
    /// GPU roundtrip.
    pub fn project(&self, features: &[f32], tokens: usize) -> Vec<f32> {
        let w = self.weight_matrix();
        let b = self.bias_vector();
        let mut out = vec![0.0f32; tokens * COMPONENTS];
        for t in 0..tokens {
            for c in 0..COMPONENTS {
                let mut acc = b[c];
                for d in 0..self.dim {
                    acc += features[t * self.dim + d] * w[d * COMPONENTS + c];
                }
                out[t * COMPONENTS + c] = acc;
            }
        }
        out
    }
}

fn normalize(v: &mut [f32]) -> f32 {
    let norm = v.iter().map(|x| x * x).sum::<f32>().sqrt();
    if norm > 1e-12 {
        for x in v.iter_mut() {
            *x /= norm;
        }
    }
    norm
}

/// Append the colour projection to an encoder output.
///
/// Declares `pca.weight` and `pca.bias` as parameters, so refitting the
/// basis is a `set_parameter` call rather than a graph rebuild.
///
/// Returns the `[tokens, COMPONENTS]` colour node. Folding this into the
/// graph rather than projecting on the CPU keeps the per-frame readback at
/// `tokens * 3` floats instead of `tokens * 384` — about 3 KB rather than
/// 300 KB, which matters on a bus shared with the compositor.
pub fn add_projection(g: &mut Graph, features: NodeId, hidden: usize) -> NodeId {
    let w = g.parameter("pca.weight", &[hidden, COMPONENTS]);
    let b = g.parameter("pca.bias", &[COMPONENTS]);
    let projected = g.matmul(features, w);
    g.bias_add(projected, b)
}

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

    /// Features lying on a known plane must be recovered by the top two
    /// components, with the projection spreading across the output range.
    #[test]
    fn recovers_a_planted_subspace() {
        let dim = 32;
        let tokens = 200;
        let mut features = vec![0.0f32; tokens * dim];
        for t in 0..tokens {
            let a = (t as f32 / tokens as f32) * 2.0 - 1.0;
            let b = ((t * 7 % tokens) as f32 / tokens as f32) * 2.0 - 1.0;
            for d in 0..dim {
                // Two strong directions plus a much weaker third.
                features[t * dim + d] = if d == 3 {
                    5.0 * a
                } else if d == 11 {
                    4.0 * b
                } else {
                    0.01 * ((d as f32) * 0.1 + a)
                };
            }
        }

        let basis = Basis::fit(&features, tokens, dim, 0);
        // The first two components should be dominated by dims 3 and 11.
        let c0 = &basis.components[0..dim];
        let c1 = &basis.components[dim..2 * dim];
        let strongest = |c: &[f32]| {
            c.iter()
                .enumerate()
                .max_by(|a, b| a.1.abs().partial_cmp(&b.1.abs()).unwrap())
                .unwrap()
                .0
        };
        let (s0, s1) = (strongest(c0), strongest(c1));
        assert!(
            (s0 == 3 && s1 == 11) || (s0 == 11 && s1 == 3),
            "expected components along dims 3 and 11, got {s0} and {s1}"
        );
    }

    #[test]
    fn components_are_orthonormal() {
        let dim = 24;
        let tokens = 90;
        let features: Vec<f32> = (0..tokens * dim)
            .map(|i| ((i * 37 % 101) as f32 / 101.0 - 0.5) * (1.0 + (i % 5) as f32))
            .collect();
        let basis = Basis::fit(&features, tokens, dim, 0);

        for a in 0..COMPONENTS {
            let va = &basis.components[a * dim..(a + 1) * dim];
            let norm: f32 = va.iter().map(|x| x * x).sum::<f32>().sqrt();
            assert!((norm - 1.0).abs() < 1e-3, "component {a} norm {norm}");
            for b in (a + 1)..COMPONENTS {
                let vb = &basis.components[b * dim..(b + 1) * dim];
                let dot: f32 = va.iter().zip(vb).map(|(x, y)| x * y).sum();
                assert!(dot.abs() < 1e-2, "components {a},{b} not orthogonal: {dot}");
            }
        }
    }

    /// The folded weight/bias form must agree with an explicit
    /// centre-project-scale-offset, since the graph only sees the folded one.
    #[test]
    fn folded_weights_match_explicit_form() {
        let dim = 16;
        let tokens = 60;
        let features: Vec<f32> = (0..tokens * dim)
            .map(|i| ((i * 13 % 97) as f32 / 97.0 - 0.5) * 3.0)
            .collect();
        let basis = Basis::fit(&features, tokens, dim, 0);
        let folded = basis.project(&features, tokens);

        for t in 0..tokens {
            for c in 0..COMPONENTS {
                let explicit: f32 = (0..dim)
                    .map(|d| {
                        (features[t * dim + d] - basis.mean[d]) * basis.components[c * dim + d]
                    })
                    .sum::<f32>()
                    * basis.scale[c]
                    + basis.offset[c];
                let got = folded[t * COMPONENTS + c];
                assert!(
                    (got - explicit).abs() < 1e-3,
                    "token {t} component {c}: folded {got} vs explicit {explicit}"
                );
            }
        }
    }

    /// The robust range should put the bulk of patches inside [0, 1],
    /// otherwise the display is either washed out or clipped.
    #[test]
    fn projection_spans_the_display_range() {
        let dim = 48;
        let tokens = 300;
        let features: Vec<f32> = (0..tokens * dim)
            .map(|i| {
                let t = (i / dim) as f32;
                let d = (i % dim) as f32;
                (t * 0.017 + d * 0.31).sin() * 2.0
            })
            .collect();
        let basis = Basis::fit(&features, tokens, dim, 0);
        let out = basis.project(&features, tokens);

        for c in 0..COMPONENTS {
            let vals: Vec<f32> = (0..tokens).map(|t| out[t * COMPONENTS + c]).collect();
            let inside = vals.iter().filter(|v| (0.0..=1.0).contains(*v)).count();
            let frac = inside as f32 / tokens as f32;
            assert!(frac > 0.9, "component {c}: only {:.0}% inside [0,1]", frac * 100.0);
        }
    }

    /// A near-featureless scene — a blank wall — must not be amplified into
    /// a noise field. Degenerate components go flat grey instead.
    #[test]
    fn degenerate_components_stay_neutral() {
        let dim = 32;
        let tokens = 120;
        // Rank-1 data: only the first component carries any variance.
        let mut features = vec![0.0f32; tokens * dim];
        for t in 0..tokens {
            let a = t as f32 / tokens as f32;
            for d in 0..dim {
                features[t * dim + d] = a * (d as f32 * 0.05).cos();
            }
        }

        let basis = Basis::fit(&features, tokens, dim, 0);
        let out = basis.project(&features, tokens);
        assert!(out.iter().all(|v| v.is_finite()), "projection produced non-finite values");

        // Components 1 and 2 have no signal; every token should sit at the
        // neutral value rather than spanning the range.
        for c in 1..COMPONENTS {
            let vals: Vec<f32> = (0..tokens).map(|t| out[t * COMPONENTS + c]).collect();
            let lo = vals.iter().copied().fold(f32::INFINITY, f32::min);
            let hi = vals.iter().copied().fold(f32::NEG_INFINITY, f32::max);
            assert!(
                (hi - lo) < 1e-3 && (lo - 0.5).abs() < 1e-3,
                "component {c} should be flat mid-grey, spans [{lo}, {hi}]"
            );
        }
    }

    #[test]
    fn placeholder_is_usable_before_the_first_fit() {
        let basis = Basis::placeholder(384);
        let features = vec![0.25f32; 8 * 384];
        let out = basis.project(&features, 8);
        assert_eq!(out.len(), 8 * COMPONENTS);
        assert!(out.iter().all(|v| v.is_finite()));
    }
}