geographdb-core 0.5.4

//! Divergence measurement between attention modes for the approximate
//! transformer math / tiebreak-correction experiment.
//!
//! The fast CPU-native path restricts attention to the current token plus its
//! geometric neighbours.  This module quantifies how much that approximation
//! drifts away from the full hybrid attention graph, both in logits and in the
//! final hidden state used by the MLP head.

use super::attention::GraphAttentionClassifier;

/// Per-mode divergence against the hybrid baseline.
#[derive(Debug, Clone, Default, PartialEq)]
pub struct ModeDivergence {
    /// Mean squared error between approximate and baseline logits.
    pub logit_mse: f32,
    /// KL divergence from the baseline distribution to the approximate one.
    pub logit_kl: f32,
    /// Euclidean distance between the final hidden states.
    pub hidden_l2: f32,
    /// Whether the approximate and baseline predictions agree.
    pub top1_agreement: bool,
}

/// Divergence metrics comparing geometric-only and mixed attention against the
/// full hybrid baseline.
#[derive(Debug, Clone, Default, PartialEq)]
pub struct DivergenceMetrics {
    pub geometric_vs_hybrid: ModeDivergence,
    pub mixed_vs_hybrid: ModeDivergence,
}

/// Per-position hidden-state divergence for a single approximate mode.
#[derive(Debug, Clone, Default, PartialEq)]
pub struct PerPositionModeDivergence {
    pub position: usize,
    /// Euclidean distance between the approximate and baseline hidden states.
    pub hidden_l2: f32,
    /// Mean squared error per hidden dimension.
    pub hidden_mse: f32,
    /// Mean absolute error per hidden dimension.
    pub hidden_mae: f32,
    /// Cosine similarity between the approximate and baseline hidden states.
    pub cosine_similarity: f32,
}

/// Per-position hidden-state divergence comparing geometric-only and mixed
/// attention against the full hybrid baseline.
#[derive(Debug, Clone, Default, PartialEq)]
pub struct PerPositionDivergenceMetrics {
    pub geometric_vs_hybrid: Vec<PerPositionModeDivergence>,
    pub mixed_vs_hybrid: Vec<PerPositionModeDivergence>,
}

/// Compare hybrid, geometric-only, and mixed attention for the same input.
///
/// `per_position_geometric[j]` selects geometric-only attention for position
/// `j`; the remaining positions use full self-attention.  The returned metrics
/// record both the pure geometric approximation and the mixed approximation.
pub fn measure_divergence(
    model: &GraphAttentionClassifier,
    token_ids: &[u32],
    neighbors: &[&[usize]],
    per_position_geometric: &[bool],
) -> DivergenceMetrics {
    let baseline_logits = model.forward_hybrid(token_ids, neighbors);
    let geometric_logits = model.forward_geometric_only(token_ids, neighbors);
    let mixed_logits = model.forward_mixed(token_ids, neighbors, per_position_geometric);

    let baseline_hidden = model.attend(token_ids, neighbors, false);
    let geometric_hidden = model.attend(token_ids, neighbors, true);
    let mixed_hidden = model.attend_mixed(token_ids, neighbors, per_position_geometric);

    let last_baseline = baseline_hidden.last().expect("empty context");
    let last_geometric = geometric_hidden.last().expect("empty context");
    let last_mixed = mixed_hidden.last().expect("empty context");

    DivergenceMetrics {
        geometric_vs_hybrid: ModeDivergence {
            logit_mse: mse(&geometric_logits, &baseline_logits),
            logit_kl: kl_divergence(&baseline_logits, &geometric_logits),
            hidden_l2: l2(last_geometric, last_baseline),
            top1_agreement: argmax(&geometric_logits) == argmax(&baseline_logits),
        },
        mixed_vs_hybrid: ModeDivergence {
            logit_mse: mse(&mixed_logits, &baseline_logits),
            logit_kl: kl_divergence(&baseline_logits, &mixed_logits),
            hidden_l2: l2(last_mixed, last_baseline),
            top1_agreement: argmax(&mixed_logits) == argmax(&baseline_logits),
        },
    }
}

/// Per-position hidden-state divergence between geometric-only / mixed attention
/// and the full hybrid baseline.
///
/// Unlike [`measure_divergence`], this function returns a per-position breakdown
/// of hidden-state drift.  This is useful for identifying which sequence
/// positions are safe to approximate with geometric-only attention and which
/// positions still need full hybrid attention.
pub fn measure_divergence_per_position(
    model: &GraphAttentionClassifier,
    token_ids: &[u32],
    neighbors: &[&[usize]],
    per_position_geometric: &[bool],
) -> PerPositionDivergenceMetrics {
    let baseline_hidden = model.attend(token_ids, neighbors, false);
    let geometric_hidden = model.attend(token_ids, neighbors, true);
    let mixed_hidden = model.attend_mixed(token_ids, neighbors, per_position_geometric);

    PerPositionDivergenceMetrics {
        geometric_vs_hybrid: baseline_hidden
            .iter()
            .zip(geometric_hidden.iter())
            .enumerate()
            .map(|(j, (b, g))| per_position_mode_divergence(j, g, b))
            .collect(),
        mixed_vs_hybrid: baseline_hidden
            .iter()
            .zip(mixed_hidden.iter())
            .enumerate()
            .map(|(j, (b, m))| per_position_mode_divergence(j, m, b))
            .collect(),
    }
}

fn per_position_mode_divergence(
    position: usize,
    approximate: &[f32],
    baseline: &[f32],
) -> PerPositionModeDivergence {
    assert_eq!(approximate.len(), baseline.len());
    let n = approximate.len();
    let mut l2_sq = 0.0f32;
    let mut mae = 0.0f32;
    let mut dot = 0.0f32;
    let mut approx_norm_sq = 0.0f32;
    let mut baseline_norm_sq = 0.0f32;
    for (a, b) in approximate.iter().zip(baseline.iter()) {
        let diff = a - b;
        l2_sq += diff * diff;
        mae += diff.abs();
        dot += a * b;
        approx_norm_sq += a * a;
        baseline_norm_sq += b * b;
    }
    let hidden_l2 = l2_sq.sqrt();
    let hidden_mse = if n > 0 { l2_sq / n as f32 } else { 0.0 };
    let hidden_mae = if n > 0 { mae / n as f32 } else { 0.0 };
    let denom = (approx_norm_sq.sqrt() * baseline_norm_sq.sqrt()).max(1e-12);
    let cosine_similarity = (dot / denom).clamp(-1.0, 1.0);

    PerPositionModeDivergence {
        position,
        hidden_l2,
        hidden_mse,
        hidden_mae,
        cosine_similarity,
    }
}

/// Divergence of the tiebreak-corrected forward pass against the hybrid
/// baseline.
pub fn measure_tiebreak_divergence(
    model: &GraphAttentionClassifier,
    token_ids: &[u32],
    neighbors: &[&[usize]],
    config: &super::attention::TiebreakConfig,
) -> ModeDivergence {
    let baseline_logits = model.forward_hybrid(token_ids, neighbors);
    let corrected_logits = model.forward_tiebreak(token_ids, neighbors, config);

    let baseline_hidden = model.attend(token_ids, neighbors, false);
    let mask = model.tiebreak_mask(token_ids, neighbors, config);
    let corrected_hidden = model.attend_mixed(token_ids, neighbors, &mask);

    let last_baseline = baseline_hidden.last().expect("empty context");
    let last_corrected = corrected_hidden.last().expect("empty context");

    ModeDivergence {
        logit_mse: mse(&corrected_logits, &baseline_logits),
        logit_kl: kl_divergence(&baseline_logits, &corrected_logits),
        hidden_l2: l2(last_corrected, last_baseline),
        top1_agreement: argmax(&corrected_logits) == argmax(&baseline_logits),
    }
}

fn argmax(xs: &[f32]) -> usize {
    xs.iter()
        .enumerate()
        .max_by(|a, b| a.1.partial_cmp(b.1).unwrap_or(std::cmp::Ordering::Equal))
        .map(|(i, _)| i)
        .unwrap_or(0)
}

fn mse(a: &[f32], b: &[f32]) -> f32 {
    assert_eq!(a.len(), b.len());
    let n = a.len() as f32;
    a.iter()
        .zip(b.iter())
        .map(|(x, y)| (x - y) * (x - y))
        .sum::<f32>()
        / n
}

fn l2(a: &[f32], b: &[f32]) -> f32 {
    assert_eq!(a.len(), b.len());
    a.iter()
        .zip(b.iter())
        .map(|(x, y)| (x - y) * (x - y))
        .sum::<f32>()
        .sqrt()
}

fn log_softmax(xs: &[f32]) -> Vec<f32> {
    let max = xs.iter().copied().fold(f32::NEG_INFINITY, f32::max);
    let exps: Vec<f32> = xs.iter().map(|x| (x - max).exp()).collect();
    let log_sum = exps.iter().sum::<f32>().ln();
    xs.iter().map(|x| x - max - log_sum).collect()
}

fn kl_divergence(p_logits: &[f32], q_logits: &[f32]) -> f32 {
    assert_eq!(p_logits.len(), q_logits.len());
    let log_p = log_softmax(p_logits);
    let log_q = log_softmax(q_logits);
    log_p
        .iter()
        .map(|lp| lp.exp())
        .zip(log_p.iter().zip(log_q.iter()))
        .map(|(p, (lp, lq))| p * (lp - lq))
        .sum::<f32>()
}

#[cfg(test)]
mod tests {
    use super::super::attention::{GraphAttentionClassifier, TiebreakConfig};
    use super::*;

    fn dummy_model() -> GraphAttentionClassifier {
        GraphAttentionClassifier::new(8, 4, 6, 3, 2, 42, None, false)
    }

    fn neighbors_for(tokens: &[u32]) -> Vec<Vec<usize>> {
        tokens
            .iter()
            .map(|&i| vec![((i as usize + 1) % 8), ((i as usize + 2) % 8)])
            .collect()
    }

    #[test]
    fn mixed_with_all_full_matches_hybrid() {
        let model = dummy_model();
        let tokens = vec![0, 1, 2, 3];
        let neighbors = neighbors_for(&tokens);
        let nbr_refs: Vec<&[usize]> = neighbors.iter().map(|v| v.as_slice()).collect();
        let mask = vec![false; tokens.len()];

        let m = measure_divergence(&model, &tokens, &nbr_refs, &mask);

        assert!(m.mixed_vs_hybrid.logit_mse < 1e-6);
        assert!(m.mixed_vs_hybrid.hidden_l2 < 1e-6);
        assert!(m.mixed_vs_hybrid.logit_kl.abs() < 1e-6);
        assert!(m.mixed_vs_hybrid.top1_agreement);
    }

    #[test]
    fn geometric_only_divergence_is_non_negative() {
        let model = dummy_model();
        let tokens = vec![0, 1, 2, 3];
        let neighbors = neighbors_for(&tokens);
        let nbr_refs: Vec<&[usize]> = neighbors.iter().map(|v| v.as_slice()).collect();
        let mask = vec![true; tokens.len()];

        let m = measure_divergence(&model, &tokens, &nbr_refs, &mask);

        assert!(m.geometric_vs_hybrid.logit_mse >= 0.0);
        assert!(m.geometric_vs_hybrid.logit_kl >= -1e-6);
        assert!(m.geometric_vs_hybrid.hidden_l2 >= 0.0);
    }

    #[test]
    fn tiebreak_divergence_is_measurable() {
        let model = dummy_model();
        let tokens = vec![0, 1, 2, 3];
        let neighbors = neighbors_for(&tokens);
        let nbr_refs: Vec<&[usize]> = neighbors.iter().map(|v| v.as_slice()).collect();
        let config = TiebreakConfig::default();

        let d = measure_tiebreak_divergence(&model, &tokens, &nbr_refs, &config);

        assert!(d.logit_mse >= 0.0);
        assert!(d.logit_kl >= -1e-6);
        assert!(d.hidden_l2 >= 0.0);
    }

    #[test]
    fn per_position_full_hybrid_is_zero() {
        let model = dummy_model();
        let tokens = vec![0, 1, 2, 3];
        let neighbors = neighbors_for(&tokens);
        let nbr_refs: Vec<&[usize]> = neighbors.iter().map(|v| v.as_slice()).collect();
        let mask = vec![false; tokens.len()];

        let per_pos = measure_divergence_per_position(&model, &tokens, &nbr_refs, &mask);

        assert_eq!(per_pos.mixed_vs_hybrid.len(), tokens.len());
        for d in &per_pos.mixed_vs_hybrid {
            assert!(
                d.hidden_l2 < 1e-6,
                "position {}: hidden_l2 = {}",
                d.position,
                d.hidden_l2
            );
            assert!(d.hidden_mse < 1e-12);
            assert!(d.hidden_mae < 1e-6);
            assert!((d.cosine_similarity - 1.0).abs() < 1e-6);
        }
    }

    #[test]
    fn per_position_geometric_is_non_negative() {
        let model = dummy_model();
        let tokens = vec![0, 1, 2, 3];
        let neighbors = neighbors_for(&tokens);
        let nbr_refs: Vec<&[usize]> = neighbors.iter().map(|v| v.as_slice()).collect();
        let mask = vec![true; tokens.len()];

        let per_pos = measure_divergence_per_position(&model, &tokens, &nbr_refs, &mask);

        assert_eq!(per_pos.geometric_vs_hybrid.len(), tokens.len());
        for d in &per_pos.geometric_vs_hybrid {
            assert!(d.hidden_l2 >= 0.0);
            assert!(d.hidden_mse >= 0.0);
            assert!(d.hidden_mae >= 0.0);
            assert!(d.cosine_similarity >= -1.0 && d.cosine_similarity <= 1.0);
        }
    }
}