environment/training-source/source-snapshot/tests/semantics.rs

//! Do the features actually mean anything?
//!
//! The point of DINO features is that patches of the same thing land near
//! each other in feature space regardless of position, and patches of
//! different things do not. That property is what the PCA colouring
//! displays and what any downstream use would rely on, and it is
//! surprisingly sensitive to exactly the mistakes that are easy to make
//! here: a transposed projection, a mis-paired RoPE half, or a patch
//! flattening in the wrong channel order all leave the magnitudes looking
//! healthy while destroying the structure.
//!
//! So rather than compare against a reference dump — which would need
//! torch on the machine — this checks the property directly, using a scene
//! with known regions.
//!
//! Requires the checkpoint. Set `DINOVISION_WEIGHTS` to a
//! `model.safetensors`; without it the test reports itself skipped, since
//! failing for a missing asset would be noise.

use dinovision::dinov3::Config;

/// Cosine similarity, the metric DINO features are usually compared under.
fn cosine(a: &[f32], b: &[f32]) -> f32 {
    let dot: f32 = a.iter().zip(b).map(|(x, y)| x * y).sum();
    let na: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
    let nb: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();
    dot / (na * nb).max(f32::MIN_POSITIVE)
}

/// A scene with four unambiguous regions, drawn as flat colour blocks so
/// the expected grouping is not a matter of opinion. Quadrants, at 224²:
/// top-left red, top-right green, bottom-left blue, bottom-right a fine
/// checkerboard (same mean colour as the background, different texture).
fn quadrant_scene(size: usize) -> Vec<u8> {
    let mut rgb = vec![0u8; size * size * 3];
    let half = size / 2;
    for y in 0..size {
        for x in 0..size {
            let px = match (x < half, y < half) {
                (true, true) => [200, 40, 40],
                (false, true) => [40, 180, 60],
                (true, false) => [50, 70, 210],
                (false, false) => {
                    let on = ((x / 4) + (y / 4)) % 2 == 0;
                    if on { [230, 230, 230] } else { [30, 30, 30] }
                }
            };
            let i = (y * size + x) * 3;
            rgb[i..i + 3].copy_from_slice(&px);
        }
    }
    rgb
}

/// Index of the patch at grid position (gx, gy), skipping prefix tokens.
fn patch(features: &[f32], config: &Config, gx: usize, gy: usize) -> Vec<f32> {
    let t = config.num_prefix_tokens() + gy * config.grid() + gx;
    let h = config.hidden_size;
    features[t * h..(t + 1) * h].to_vec()
}

#[test]
fn features_group_by_content_not_position() {
    let Ok(weights) = std::env::var("DINOVISION_WEIGHTS") else {
        eprintln!("skipped: set DINOVISION_WEIGHTS to a model.safetensors to run this");
        return;
    };
    let weights = std::path::PathBuf::from(weights);
    if !weights.exists() {
        eprintln!("skipped: {} does not exist", weights.display());
        return;
    }

    let config = Config::vits16().at_resolution(224);
    let gpu = dinovision::init_context(None).expect("GPU context");
    let (mut session, _) = dinovision::bench::build_encoder_session(gpu, &config, None);
    let model = meganeura::data::safetensors::SafeTensorsModel::load(weights).expect("read weights");
    dinovision::weights::load_encoder(&mut session, &model, &config).expect("bind weights");

    let rgb = quadrant_scene(config.image_size);
    let patches = dinovision::preprocess::patches_from_rgb8(&rgb, &config);
    session.set_input("patches", &patches);
    session.step();
    session.wait();
    let features = session.read_output(config.num_tokens() * config.hidden_size);
    assert!(
        features.iter().all(|v| v.is_finite()),
        "features contain non-finite values"
    );

    // Sample two well-separated patches inside each quadrant, staying away
    // from the boundaries where receptive fields mix regions.
    let g = config.grid();
    let q = g / 4;
    let regions = [
        ("red", [(q, q), (q + 1, q + 1)]),
        ("green", [(g - q - 1, q), (g - q - 2, q + 1)]),
        ("blue", [(q, g - q - 1), (q + 1, g - q - 2)]),
        ("checker", [(g - q - 1, g - q - 1), (g - q - 2, g - q - 2)]),
    ];

    // Within a region, two patches of the same material should be close.
    let mut worst_within = 1.0f32;
    for (name, pts) in &regions {
        let a = patch(&features, &config, pts[0].0, pts[0].1);
        let b = patch(&features, &config, pts[1].0, pts[1].1);
        let c = cosine(&a, &b);
        eprintln!("within {name:>8}: {c:.3}");
        worst_within = worst_within.min(c);
    }

    // Across regions, patches of different material should be further apart
    // than any same-material pair.
    let mut best_across = -1.0f32;
    for i in 0..regions.len() {
        for j in (i + 1)..regions.len() {
            let a = patch(&features, &config, regions[i].1[0].0, regions[i].1[0].1);
            let b = patch(&features, &config, regions[j].1[0].0, regions[j].1[0].1);
            let c = cosine(&a, &b);
            eprintln!("across {:>8}/{:<8}: {c:.3}", regions[i].0, regions[j].0);
            best_across = best_across.max(c);
        }
    }

    eprintln!("worst within-region {worst_within:.3}, best across-region {best_across:.3}");
    assert!(
        worst_within > best_across,
        "features do not separate content: the least similar same-region pair \
         ({worst_within:.3}) scored below the most similar different-region pair \
         ({best_across:.3}). A scrambled patch order or transposed projection \
         looks exactly like this."
    );
}